1// This file exports a single function, which is used to read the
2// pixel data from an image.
3//
4// This includes resizing the image to a smaller size (~400×400) for
5// faster downstream computations.
6//
7// It returns a Vec<Lab>, which can be passed to the k-means process.
8
9use std::fmt::Display;
10use std::fs::File;
11use std::io::BufReader;
12use std::path::PathBuf;
13
14use image::codecs::gif::GifDecoder;
15use image::codecs::webp::WebPDecoder;
16use image::imageops::FilterType;
17use image::{AnimationDecoder, DynamicImage, Frame, ImageFormat};
18use palette::cast::from_component_slice;
19use palette::{IntoColor, Lab, Srgba};
20
21pub fn get_image_colors(path: &PathBuf) -> Result<Vec<Lab>, GetImageColorsErr> {
22    let format = get_format(path)?;
23
24    let f = File::open(path)?;
25    let reader = BufReader::new(f);
26
27    let image_bytes = match format {
28        ImageFormat::Gif => {
29            let decoder = GifDecoder::new(reader)?;
30            get_bytes_for_animated_image(decoder)
31        }
32
33        ImageFormat::WebP if is_animated_webp(path) => {
34            let decoder = WebPDecoder::new(reader)?;
35            get_bytes_for_animated_image(decoder)
36        }
37
38        format => {
39            let decoder = image::load(reader, format)?;
40            get_bytes_for_static_image(decoder)
41        }
42    };
43
44    let lab: Vec<Lab> = from_component_slice::<Srgba<u8>>(&image_bytes)
45        .iter()
46        .map(|x| x.into_format::<_, f32>().into_color())
47        .collect();
48
49    Ok(lab)
50}
51
52#[derive(Debug)]
53pub enum GetImageColorsErr {
54    IoError(std::io::Error),
55    ImageError(image::ImageError),
56    GetFormatError(String),
57}
58
59impl Display for GetImageColorsErr {
60    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
61        match self {
62            GetImageColorsErr::IoError(io_error) => write!(f, "{}", io_error),
63            GetImageColorsErr::ImageError(image_error) => write!(f, "{}", image_error),
64            GetImageColorsErr::GetFormatError(format_error) => write!(f, "{}", format_error),
65        }
66    }
67}
68
69impl From<std::io::Error> for GetImageColorsErr {
70    fn from(e: std::io::Error) -> GetImageColorsErr {
71        return GetImageColorsErr::IoError(e);
72    }
73}
74
75impl From<image::ImageError> for GetImageColorsErr {
76    fn from(e: image::ImageError) -> GetImageColorsErr {
77        return GetImageColorsErr::ImageError(e);
78    }
79}
80
81fn get_format(path: &PathBuf) -> Result<ImageFormat, GetImageColorsErr> {
82    let format = match path.extension() {
83        Some(ext) => Ok(image::ImageFormat::from_extension(ext)),
84        None => Err(GetImageColorsErr::GetFormatError(
85            "Path has no file extension, so could not determine image format".to_string(),
86        )),
87    };
88
89    match format {
90        Ok(Some(format)) => Ok(format),
91        Ok(None) => Err(GetImageColorsErr::GetFormatError(
92            "Unable to determine image format from file extension".to_string(),
93        )),
94        Err(e) => Err(e),
95    }
96}
97
98/// Returns true if a WebP file is animated, and false otherwise.
99///
100/// This function assumes that the file exists and can be opened.
101fn is_animated_webp(path: &PathBuf) -> bool {
102    let f = File::open(path).unwrap();
103    let reader = BufReader::new(f);
104
105    match image_webp::WebPDecoder::new(reader) {
106        // num_frames() returns the number of frames in the animation,
107        // or zero if the image is not animated.
108        // See https://docs.rs/image-webp/0.2.0/image_webp/struct.WebPDecoder.html#method.num_frames
109        Ok(decoder) => decoder.num_frames() > 0,
110        Err(_) => false,
111    }
112}
113
114fn get_bytes_for_static_image(img: DynamicImage) -> Vec<u8> {
115    // Resize the image after we open it.  For this tool I'd rather get a good answer
116    // quickly than a great answer slower.
117    //
118    // The choice of max dimension is arbitrary.  Making it smaller means you get
119    // faster results, but possibly at the loss of quality.
120    //
121    // The nearest neighbour algorithm produces images that don't look as good,
122    // but it's much much faster and the loss of quality is unlikely to be
123    // an issue when looking for dominant colours.
124    //
125    // Note: when trying to work out what's "fast enough", make sure you use release
126    // mode.  The image/k-means operations are significantly faster (=2 orders
127    // of magnitude) than in debug mode.
128    //
129    // See https://docs.rs/image/0.23.14/image/imageops/enum.FilterType.html
130    let resized_img = img.resize(400, 400, FilterType::Nearest);
131
132    resized_img.into_rgba8().into_raw()
133}
134
135fn get_bytes_for_animated_image<'a>(decoder: impl AnimationDecoder<'a>) -> Vec<u8> {
136    let frames: Vec<Frame> = decoder.into_frames().collect_frames().unwrap();
137    assert!(frames.len() > 0);
138
139    // If the image is animated, we want to make sure we look at multiple
140    // frames when choosing the dominant colour.
141    //
142    // We don't want to pass all the frames to the k-means analysis, because
143    // that would be incredibly memory-intensive and is unnecessary -- see
144    // previous comments about wanting a good enough answer quickly.
145    //
146    // For that reason, we select a sample of up to 50 frames and use those
147    // as the basis for analysis.
148    //
149    // How this works: it tells us we should be looking at the nth frame.
150    // Examples:
151    //
152    //      frame count | nth frame | comment
153    //      ------------+-----------+---------
154    //      1           |     1     | in a 1-frame GIF, look at the only frame
155    //      25          |     1     | look at every frame
156    //      50          |     2     | look at every second frame
157    //      78          |     3     | look at every third frame
158    //
159    // I'm sure there's a more idiomatic way to do this, but it was late
160    // when I wrote this and it seems to work.
161    //
162    let nth_frame = if frames.len() <= 50 {
163        1
164    } else {
165        ((frames.len() as f32) / (25 as f32)) as i32
166    };
167
168    let selected_frames = frames
169        .iter()
170        .enumerate()
171        .filter(|(i, _)| (*i as f32 / nth_frame as f32).floor() == (*i as f32 / nth_frame as f32))
172        .map(|(_, frame)| frame);
173
174    // Now we go through the frames and extract all the pixels.  The k-means
175    // process doesn't care about position, so we can concatenate the pixels
176    // for each frame into one big Vec.
177    //
178    // As with non-GIF images, we resize the images down before loading them.
179    // We resize to a smaller frame in GIFs because if there are multiple
180    // frames, we don't care as much about individual frames, and we want
181    // to avoid a large Vec<u8> in memory.
182    let resize = if frames.len() == 1 { 400 } else { 100 };
183
184    selected_frames
185        .map(|frame| {
186            DynamicImage::ImageRgba8(frame.buffer().clone())
187                .resize(resize, resize, FilterType::Nearest)
188                .into_rgba8()
189                .into_raw()
190        })
191        .into_iter()
192        .flatten()
193        .collect()
194}
195
196#[cfg(test)]
197mod test {
198    use std::path::PathBuf;
199
200    use crate::get_image_colors::get_image_colors;
201
202    // This image comes from https://stacks.wellcomecollection.org/peering-through-mri-scans-of-fruit-and-veg-part-1-a2e8b07bde6f
203    //
204    // I don't remember how I got these images, but for some reason they
205    // caused v1.1.2 to fall over.  This is a test that they can still be
206    // processed correctly.
207    #[test]
208    fn it_gets_colors_for_animated_gif() {
209        let colors = get_image_colors(&PathBuf::from("./src/tests/garlic.gif"));
210
211        assert!(colors.is_ok());
212        assert!(colors.unwrap().len() > 0);
213    }
214
215    #[test]
216    fn get_colors_for_static_gif() {
217        let colors = get_image_colors(&PathBuf::from("./src/tests/yellow.gif"));
218
219        assert!(colors.is_ok());
220        assert!(colors.unwrap().len() > 0);
221    }
222
223    #[test]
224    fn get_colors_for_static_webp() {
225        let colors = get_image_colors(&PathBuf::from("./src/tests/purple.webp"));
226
227        assert!(colors.is_ok());
228        assert!(colors.unwrap().len() > 0);
229    }
230
231    #[test]
232    fn get_colors_for_animated_webp() {
233        let colors = get_image_colors(&PathBuf::from("./src/tests/animated_squares.webp"));
234
235        assert!(colors.is_ok());
236        assert!(colors.unwrap().len() > 0);
237    }
238}