Skip to main content

Fix a bug when getting colours for a non-animated WebP

ID
81bb194
date
2024-11-06 21:19:00+00:00
author
Alex Chan <alex@alexwlchan.net>
parent
bca18d1
message
Fix a bug when getting colours for a non-animated WebP

I tried to get dominant colours for a WebP image (a thumbnail of a YouTube
video) and it was failing with an assertion inside kmeans-colors:

    $ RUST_BACKTRACE=1 dominant_colours src/tests/purple.webp
    thread 'main' panicked at kmeans_colors-0.6.0/src/plus_plus.rs:24:5:
    assertion failed: len > 0
    stack backtrace:
       0: _rust_begin_unwind
       1: core::panicking::panic_fmt
       2: core::panicking::panic
       3: kmeans_colors::plus_plus::init_plus_plus
       4: kmeans_colors::kmeans::get_kmeans_hamerly
       5: dominant_colours::main

These are the lines where it's failing [1]:

    let len = buf.len();
    assert!(len > 0);

That `buf` variable is the list of colors we extract from an image in
`get_image_colors`.  I added a couple of internal assertions to this
tool, and discovered we're getting an empty list of colours for
non-animated WebP images.

This wasn't caught by the tests because they were just checking that
`get_image_colors` returns an Ok() result, but not inspecting its value.
I've strengthened the tests to check it returns a non-empty list.

Why are there 0 colours for a non-animated WebP image?  Because we were
assuming we could treat all WebP images as animated, and passing them to
`get_bytes_for_animated_image`.  Then we could extract a list of frames,
and for a non-animated image this would be a list with a single frame:

    let image_bytes = match format {
        ImageFormat::WebP => {
            let decoder = WebPDecoder::new(reader)?;
            get_bytes_for_animated_image(decoder)
        }

        …
    };

    …

    fn get_bytes_for_animated_image<'a>(decoder) -> Vec<u8> {
        let frames = decoder.into_frames().collect_frames().unwrap();

This approach comes from the way the image crate handles GIF images --
for non-animated GIFs, this is a list with a single frame.

But WebP images are different, and we get an empty list here!

Initially I thought this was a bug in the image crate, but this looks like
an intentional design decision -- the underlying `image-webp` crate has
`num_frames = 0` if a WebP is not animated [2].
(By contrast, the Pillow library in Python has `num_frames = 1`.)

This patch:

*   Adds some extra internal assertions to detect if we have an empty
    list of colours, so that can be caught in our code, before we pass
    it to `kmeans-colors`
*   Only uses the animated image decoder for WebP images if they're actually
    animated, and otherwise uses the static image decoder
*   Strengthens the tests for `get_image_colors`

[1]: https://github.com/okaneco/kmeans-colors/blob/03c2f7b4b11ea4e295de8aafe88ad4fc15218b78/src/plus_plus.rs#L23-L24
[2]: https://docs.rs/image-webp/0.2.0/image_webp/struct.WebPDecoder.html#method.num_frames
changed files
6 files, 57 additions, 7 deletions

Changed files

CHANGELOG.md (2094) → CHANGELOG.md (2241)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index bf4b158..19850c3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,9 @@
 # Changelog
 
+## v1.4.1 - 2024-11-06
+
+*   Fix a bug introduced in v1.3.0 where getting colours for non-animated WebP images would fail with an assertion error.
+
 ## v1.4.0 - 2024-10-05
 
 *   `dominant_colours` will now skip printing terminal colours if it detects it's not running in a tty.  This makes it slightly easier to use in automated environments, because you don't need to pass the `--no-palette` flag.

Cargo.lock (17572) → Cargo.lock (17587)

diff --git a/Cargo.lock b/Cargo.lock
index b747cf2..b23e6b7 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -213,11 +213,12 @@ checksum = "fea41bba32d969b513997752735605054bc0dfa92b4c56bf1189f2e174be7a10"
 
 [[package]]
 name = "dominant_colours"
-version = "1.4.0"
+version = "1.4.1"
 dependencies = [
  "assert_cmd",
  "clap",
  "image",
+ "image-webp",
  "kmeans_colors",
  "palette",
  "regex",

Cargo.toml (479) → Cargo.toml (500)

diff --git a/Cargo.toml b/Cargo.toml
index bff0089..88bf30b 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,11 +1,12 @@
 [package]
 name = "dominant_colours"
-version = "1.4.0"
+version = "1.4.1"
 edition = "2018"
 
 [dependencies]
 assert_cmd = "2.0.16"
 clap = { version = "4.5.20", features = ["derive"] }
+image-webp = "0.2.0"
 regex = "1.11.1"
 
 [dependencies.kmeans_colors]

src/find_dominant_colors.rs (2737) → src/find_dominant_colors.rs (2766)

diff --git a/src/find_dominant_colors.rs b/src/find_dominant_colors.rs
index 2e03315..6b08ce8 100644
--- a/src/find_dominant_colors.rs
+++ b/src/find_dominant_colors.rs
@@ -10,6 +10,8 @@ pub fn find_dominant_colors(lab: &Vec<Lab>, max_colors: usize) -> Vec<Lab> {
     let verbose = false;
     let seed: u64 = 0;
 
+    assert!(lab.len() > 0);
+
     let result = get_kmeans_hamerly(max_colors, max_iterations, converge, verbose, lab, seed);
 
     result.centroids

src/get_image_colors.rs (6954) → src/get_image_colors.rs (8241)

diff --git a/src/get_image_colors.rs b/src/get_image_colors.rs
index cf3e96a..ff098ec 100644
--- a/src/get_image_colors.rs
+++ b/src/get_image_colors.rs
@@ -30,7 +30,7 @@ pub fn get_image_colors(path: &PathBuf) -> Result<Vec<Lab>, GetImageColorsErr> {
             get_bytes_for_animated_image(decoder)
         }
 
-        ImageFormat::WebP => {
+        ImageFormat::WebP if is_animated_webp(path) => {
             let decoder = WebPDecoder::new(reader)?;
             get_bytes_for_animated_image(decoder)
         }
@@ -49,6 +49,7 @@ pub fn get_image_colors(path: &PathBuf) -> Result<Vec<Lab>, GetImageColorsErr> {
     Ok(lab)
 }
 
+#[derive(Debug)]
 pub enum GetImageColorsErr {
     IoError(std::io::Error),
     ImageError(image::ImageError),
@@ -94,6 +95,22 @@ fn get_format(path: &PathBuf) -> Result<ImageFormat, GetImageColorsErr> {
     }
 }
 
+/// Returns true if a WebP file is animated, and false otherwise.
+///
+/// This function assumes that the file exists and can be opened.
+fn is_animated_webp(path: &PathBuf) -> bool {
+    let f = File::open(path).unwrap();
+    let reader = BufReader::new(f);
+
+    match image_webp::WebPDecoder::new(reader) {
+        // num_frames() returns the number of frames in the animation,
+        // or zero if the image is not animated.
+        // See https://docs.rs/image-webp/0.2.0/image_webp/struct.WebPDecoder.html#method.num_frames
+        Ok(decoder) => decoder.num_frames() > 0,
+        Err(_) => false,
+    }
+}
+
 fn get_bytes_for_static_image(img: DynamicImage) -> Vec<u8> {
     // Resize the image after we open it.  For this tool I'd rather get a good answer
     // quickly than a great answer slower.
@@ -117,6 +134,7 @@ fn get_bytes_for_static_image(img: DynamicImage) -> Vec<u8> {
 
 fn get_bytes_for_animated_image<'a>(decoder: impl AnimationDecoder<'a>) -> Vec<u8> {
     let frames: Vec<Frame> = decoder.into_frames().collect_frames().unwrap();
+    assert!(frames.len() > 0);
 
     // If the image is animated, we want to make sure we look at multiple
     // frames when choosing the dominant colour.
@@ -187,12 +205,34 @@ mod test {
     // caused v1.1.2 to fall over.  This is a test that they can still be
     // processed correctly.
     #[test]
-    fn it_gets_colors_for_mri_fruit() {
-        assert!(get_image_colors(&PathBuf::from("./src/tests/garlic.gif")).is_ok());
+    fn it_gets_colors_for_animated_gif() {
+        let colors = get_image_colors(&PathBuf::from("./src/tests/garlic.gif"));
+
+        assert!(colors.is_ok());
+        assert!(colors.unwrap().len() > 0);
     }
 
     #[test]
-    fn get_colors_for_webp() {
-        assert!(get_image_colors(&PathBuf::from("./src/tests/purple.webp")).is_ok());
+    fn get_colors_for_static_gif() {
+        let colors = get_image_colors(&PathBuf::from("./src/tests/yellow.gif"));
+
+        assert!(colors.is_ok());
+        assert!(colors.unwrap().len() > 0);
+    }
+
+    #[test]
+    fn get_colors_for_static_webp() {
+        let colors = get_image_colors(&PathBuf::from("./src/tests/purple.webp"));
+
+        assert!(colors.is_ok());
+        assert!(colors.unwrap().len() > 0);
+    }
+
+    #[test]
+    fn get_colors_for_animated_webp() {
+        let colors = get_image_colors(&PathBuf::from("./src/tests/animated_squares.webp"));
+
+        assert!(colors.is_ok());
+        assert!(colors.unwrap().len() > 0);
     }
 }

src/main.rs (10935) → src/main.rs (10964)

diff --git a/src/main.rs b/src/main.rs
index c9efd50..899d235 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -40,6 +40,8 @@ fn main() {
         }
     };
 
+    assert!(lab.len() > 0);
+
     let dominant_colors = find_dominant_colors::find_dominant_colors(&lab, cli.max_colours);
 
     let selected_colors = match cli.background {