ID

907dafc

date

2024-09-04 06:52:31+00:00

author

Alex Chan <alex@alexwlchan.net>

parent

e704d87

message

Refactor to add support for animated WebP

changed files

3 files, 52 additions, 36 deletions

• src/get_image_colors.rs
• src/main.rs
• src/tests/animated_squares.webp

Changed files

src/get_image_colors.rs (5623) → src/get_image_colors.rs (5939)

diff --git a/src/get_image_colors.rs b/src/get_image_colors.rs
index a40b013..dfe4cd4 100644
--- a/src/get_image_colors.rs
+++ b/src/get_image_colors.rs
@@ -6,26 +6,56 @@
 //
 // It returns a Vec<Lab>, which can be passed to the k-means process.
 
-use std::ffi::OsStr;
 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 extension = match path.extension().and_then(OsStr::to_str) {
-        Some(ext) => Some(ext.to_lowercase()),
-        None => None,
+    let format = image::ImageFormat::from_extension(path.extension().unwrap());
+
+    let f = match File::open(path) {
+        Ok(im) => im,
+        Err(e) => {
+            eprintln!("{}", e);
+            std::process::exit(1);
+        }
     };
 
-    let image_bytes = match extension {
-        Some(ext) if ext == "gif" => get_bytes_for_gif(&path),
-        _ => get_bytes_for_static_image(&path),
+    let reader = BufReader::new(f);
+
+    let image_bytes = match format {
+        Some(ImageFormat::Gif) => {
+            let decoder = GifDecoder::new(reader).ok().unwrap();
+            get_bytes_for_animated_image(decoder)
+        }
+
+        Some(ImageFormat::WebP) => {
+            let decoder = WebPDecoder::new(reader).ok().unwrap();
+            get_bytes_for_animated_image(decoder)
+        }
+
+        Some(format) => {
+            let decoder = match image::load(reader, format) {
+                Ok(im) => im,
+                Err(e) => {
+                    eprintln!("{}", e);
+                    std::process::exit(1)
+                }
+            };
+            get_bytes_for_static_image(decoder)
+        }
+
+        _ => {
+            eprintln!("The image format could not be determined");
+            std::process::exit(1);
+        }
     };
 
     let lab: Vec<Lab> = from_component_slice::<Srgba<u8>>(&image_bytes)
@@ -36,15 +66,7 @@ pub fn get_image_colors(path: &PathBuf) -> Vec<Lab> {
     lab
 }
 
-fn get_bytes_for_static_image(path: &PathBuf) -> Vec<u8> {
-    let img = match image::open(&path) {
-        Ok(im) => im,
-        Err(e) => {
-            eprintln!("{}", e);
-            std::process::exit(1);
-        }
-    };
-
+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.
     //
@@ -125,22 +147,6 @@ fn get_bytes_for_animated_image<'a>(decoder: impl AnimationDecoder<'a>) -> Vec<u
         .collect()
 }
 
-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();
-
-    get_bytes_for_animated_image(decoder)
-}
-
 #[cfg(test)]
 mod test {
     use std::path::PathBuf;

src/main.rs (8608) → src/main.rs (8922)

diff --git a/src/main.rs b/src/main.rs
index 38593c4..3dc25cd 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -148,11 +148,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
         );
@@ -171,6 +170,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"]);
 

src/tests/animated_squares.webp (0) → src/tests/animated_squares.webp (804)

diff --git a/src/tests/animated_squares.webp b/src/tests/animated_squares.webp
new file mode 100644
index 0000000..97eff04
Binary files /dev/null and b/src/tests/animated_squares.webp differ