Merge pull request #60 from alexwlchan/animated-webp

Add supported for animated WebP images

Author: alexwlchan

CHANGELOG.md

@@ -1,5 +1,10 @@
 # Changelog
 
+## v1.3.0 - 2024-09-04
+
+*   Add support for animated WebP images.
+*   Improve the error messages, especially when dealing with malformed images.
+
 ## v1.2.0 - 2024-05-12
 
 Two new features:

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "dominant_colours"
-version = "1.2.0"
+version = "1.3.0"
 edition = "2018"
 
 [dependencies]

Cargo.toml

@@ -6,40 +6,95 @@
 //
 // It returns a Vec<Lab>, which can be passed to the k-means process.
 
-use std::ffi::OsStr;
+use std::fmt::Display;
 use std::fs::File;
 use std::io::BufReader;
 use std::path::PathBuf;
 
 use image::codecs::gif::GifDecoder;
+use image::codecs::webp::WebPDecoder;
 use image::imageops::FilterType;
-use image::{AnimationDecoder, DynamicImage, Frame};
+use image::{AnimationDecoder, DynamicImage, Frame, ImageFormat};
 use palette::cast::from_component_slice;
 use palette::{IntoColor, Lab, Srgba};
 
-pub fn get_image_colors(path: &PathBuf) -> Vec<Lab> {
-    let image_bytes = match path.extension().and_then(OsStr::to_str) {
-        Some(ext) if ext.to_lowercase() == "gif" => get_bytes_for_gif(&path),
-        _ => get_bytes_for_non_gif(&path),
+pub fn get_image_colors(path: &PathBuf) -> Result<Vec<Lab>, GetImageColorsErr> {
+    let format = get_format(path)?;
+
+    let f = File::open(path)?;
+    let reader = BufReader::new(f);
+
+    let image_bytes = match format {
+        ImageFormat::Gif => {
+            let decoder = GifDecoder::new(reader)?;
+            get_bytes_for_animated_image(decoder)
+        }
+
+        ImageFormat::WebP => {
+            let decoder = WebPDecoder::new(reader)?;
+            get_bytes_for_animated_image(decoder)
+        }
+
+        format => {
+            let decoder = image::load(reader, format)?;
+            get_bytes_for_static_image(decoder)
+        }
     };
 
     let lab: Vec<Lab> = from_component_slice::<Srgba<u8>>(&image_bytes)
         .iter()
         .map(|x| x.into_format::<_, f32>().into_color())
         .collect();
 
-    lab
+    Ok(lab)
 }
 
-fn get_bytes_for_non_gif(path: &PathBuf) -> Vec<u8> {
-    let img = match image::open(&path) {
-        Ok(im) => im,
-        Err(e) => {
-            eprintln!("{}", e);
-            std::process::exit(1);
+pub enum GetImageColorsErr {
+    IoError(std::io::Error),
+    ImageError(image::ImageError),
+    GetFormatError(String),
+}
+
+impl Display for GetImageColorsErr {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            GetImageColorsErr::IoError(io_error) => write!(f, "{}", io_error),
+            GetImageColorsErr::ImageError(image_error) => write!(f, "{}", image_error),
+            GetImageColorsErr::GetFormatError(format_error) => write!(f, "{}", format_error),
         }
+    }
+}
+
+impl From<std::io::Error> for GetImageColorsErr {
+    fn from(e: std::io::Error) -> GetImageColorsErr {
+        return GetImageColorsErr::IoError(e);
+    }
+}
+
+impl From<image::ImageError> for GetImageColorsErr {
+    fn from(e: image::ImageError) -> GetImageColorsErr {
+        return GetImageColorsErr::ImageError(e);
+    }
+}
+
+fn get_format(path: &PathBuf) -> Result<ImageFormat, GetImageColorsErr> {
+    let format = match path.extension() {
+        Some(ext) => Ok(image::ImageFormat::from_extension(ext)),
+        None => Err(GetImageColorsErr::GetFormatError(
+            "Path has no file extension, so could not determine image format".to_string(),
+        )),
     };
 
+    match format {
+        Ok(Some(format)) => Ok(format),
+        Ok(None) => Err(GetImageColorsErr::GetFormatError(
+            "Unable to determine image format from file extension".to_string(),
+        )),
+        Err(e) => Err(e),
+    }
+}
+
+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.
     //
@@ -60,20 +115,10 @@ fn get_bytes_for_non_gif(path: &PathBuf) -> Vec<u8> {
     resized_img.into_rgba8().into_raw()
 }
 
-fn get_bytes_for_gif(path: &PathBuf) -> Vec<u8> {
-    let f = match File::open(path) {
-        Ok(im) => im,
-        Err(e) => {
-            eprintln!("{}", e);
-            std::process::exit(1);
-        }
-    };
-
-    let f = BufReader::new(f);
-
-    let decoder = GifDecoder::new(f).ok().unwrap();
+fn get_bytes_for_animated_image<'a>(decoder: impl AnimationDecoder<'a>) -> Vec<u8> {
+    let frames: Vec<Frame> = decoder.into_frames().collect_frames().unwrap();
 
-    // If the GIF is animated, we want to make sure we look at multiple
+    // If the image is animated, we want to make sure we look at multiple
     // frames when choosing the dominant colour.
     //
     // We don't want to pass all the frames to the k-means analysis, because
@@ -82,8 +127,7 @@ fn get_bytes_for_gif(path: &PathBuf) -> Vec<u8> {
     //
     // For that reason, we select a sample of up to 50 frames and use those
     // as the basis for analysis.
-    let frames: Vec<Frame> = decoder.into_frames().collect_frames().unwrap();
-
+    //
     // How this works: it tells us we should be looking at the nth frame.
     // Examples:
     //
@@ -144,11 +188,11 @@ mod test {
     // processed correctly.
     #[test]
     fn it_gets_colors_for_mri_fruit() {
-        get_image_colors(&PathBuf::from("./src/tests/garlic.gif"));
+        assert!(get_image_colors(&PathBuf::from("./src/tests/garlic.gif")).is_ok());
     }
 
     #[test]
     fn get_colors_for_webp() {
-        get_image_colors(&PathBuf::from("./src/tests/purple.webp"));
+        assert!(get_image_colors(&PathBuf::from("./src/tests/purple.webp")).is_ok());
     }
 }

src/get_image_colors.rs

@@ -31,7 +31,13 @@ struct Cli {
 fn main() {
     let cli = Cli::parse();
 
-    let lab: Vec<Lab> = get_image_colors::get_image_colors(&cli.path);
+    let lab: Vec<Lab> = match get_image_colors::get_image_colors(&cli.path) {
+        Ok(lab) => lab,
+        Err(e) => {
+            eprintln!("{}", e);
+            std::process::exit(1);
+        }
+    };
 
     let dominant_colors = find_dominant_colors::find_dominant_colors(&lab, cli.max_colours);
 
@@ -148,11 +154,10 @@ mod tests {
 
     #[test]
     fn it_looks_at_multiple_frames_in_an_animated_gif() {
-        let output = get_success(&["./src/tests/animated_squares.gif"]);
+        let output = get_success(&["./src/tests/animated_squares.gif", "--no-palette"]);
 
         assert_eq!(
-            output.stdout.matches("\n").count(),
-            2,
+            output.stdout, "#0200ff\n#ff0000\n",
             "stdout = {:?}",
             output.stdout
         );
@@ -170,6 +175,17 @@ mod tests {
         );
     }
 
+    #[test]
+    fn it_looks_at_multiple_frames_in_an_animated_webp() {
+        let output = get_success(&["./src/tests/animated_squares.webp", "--no-palette"]);
+
+        assert_eq!(
+            output.stdout, "#0200ff\n#ff0100\n#ff0002\n",
+            "stdout = {:?}",
+            output.stdout
+        );
+    }
+
     #[test]
     fn it_fails_if_you_pass_an_invalid_max_colours() {
         let output = get_failure(&["./src/tests/red.png", "--max-colours=NaN"]);
@@ -206,7 +222,10 @@ mod tests {
 
         assert_eq!(output.exit_code, 1);
         assert_eq!(output.stdout, "");
-        assert_eq!(output.stderr, "The image format could not be determined\n");
+        assert_eq!(
+            output.stderr,
+            "Unable to determine image format from file extension\n"
+        );
     }
 
     #[test]
@@ -215,7 +234,10 @@ mod tests {
 
         assert_eq!(output.exit_code, 1);
         assert_eq!(output.stdout, "");
-        assert_eq!(output.stderr, "The image format could not be determined\n");
+        assert_eq!(
+            output.stderr,
+            "Unable to determine image format from file extension\n"
+        );
     }
 
     #[test]
@@ -230,6 +252,42 @@ mod tests {
         );
     }
 
+    #[test]
+    fn it_fails_if_you_pass_a_malformed_gif() {
+        let output = get_failure(&["./src/tests/malformed.txt.gif"]);
+
+        assert_eq!(output.exit_code, 1);
+        assert_eq!(output.stdout, "");
+        assert_eq!(
+            output.stderr,
+            "Format error decoding Gif: malformed GIF header\n"
+        );
+    }
+
+    #[test]
+    fn it_fails_if_you_pass_a_malformed_webp() {
+        let output = get_failure(&["./src/tests/malformed.txt.webp"]);
+
+        assert_eq!(output.exit_code, 1);
+        assert_eq!(output.stdout, "");
+        assert_eq!(
+            output.stderr,
+            "Format error decoding WebP: Invalid Chunk header: [82, 73, 70, 70]\n"
+        );
+    }
+
+    #[test]
+    fn it_fails_if_you_pass_a_path_without_a_file_extension() {
+        let output = get_failure(&["./src/tests/noextension"]);
+
+        assert_eq!(output.exit_code, 1);
+        assert_eq!(output.stdout, "");
+        assert_eq!(
+            output.stderr,
+            "Path has no file extension, so could not determine image format\n"
+        );
+    }
+
     #[test]
     fn it_chooses_the_right_color_for_a_dark_background() {
         let output = get_success(&[

src/main.rs

@@ -0,0 +1 @@
+Ceci n'est pas une png