Merge pull request #71 from alexwlchan/static-webp-bug

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

Author: alexwlchan

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

@@ -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]

Cargo.toml

@@ -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/find_dominant_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/get_image_colors.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 {