Merge pull request #110 from simon-frankau/exact-quantize

Add support for exact quantisation of images with <= 256 colors

src/common.rs

@@ -2,6 +2,8 @@
 extern crate color_quant;
 
 use std::borrow::Cow;
+use std::collections::HashMap;
+use std::collections::HashSet;
 
 /// Disposal method
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
@@ -215,14 +217,40 @@ impl Frame<'static> {
             }
         }
 
-        let nq = color_quant::NeuQuant::new(speed, 256, pixels);
+        // Attempt to build a palette of all colors. If we go over 256 colors,
+        // switch to the NeuQuant algorithm.
+        let mut colors: HashSet<(u8, u8, u8, u8)> = HashSet::new();
+        for pixel in pixels.chunks_exact(4) {
+            if colors.insert((pixel[0], pixel[1], pixel[2], pixel[3])) && colors.len() > 256 {
+                // > 256 colours, let's use NeuQuant.
+                let nq =  color_quant::NeuQuant::new(speed, 256, pixels);
 
-        Frame {
+                return Frame {
+                    width,
+                    height,
+                    buffer: Cow::Owned(pixels.chunks_exact(4).map(|pix| nq.index_of(pix) as u8).collect()),
+                    palette: Some(nq.color_map_rgb()),
+                    transparent: transparent.map(|t| nq.index_of(&t) as u8),
+                    ..Frame::default()
+                };
+            }
+        }
+
+        // Palette size <= 256 elements, we can build an exact palette.
+        let mut colors_vec: Vec<(u8, u8, u8, u8)> = colors.into_iter().collect();
+        colors_vec.sort();
+        let palette = colors_vec.iter().map(|&(r, g, b, _a)| vec![r, g, b]).flatten().collect();
+        let colors_lookup: HashMap<(u8, u8, u8, u8), u8> =  colors_vec.into_iter().zip(0..).collect();
+
+        let index_of = | pixel: &[u8] |
+            *colors_lookup.get(&(pixel[0], pixel[1], pixel[2], pixel[3])).unwrap();
+
+        return Frame {
             width,
             height,
-            buffer: Cow::Owned(pixels.chunks_exact(4).map(|pix| nq.index_of(pix) as u8).collect()),
-            palette: Some(nq.color_map_rgb()),
-            transparent: transparent.map(|t| nq.index_of(&t) as u8),
+            buffer: Cow::Owned(pixels.chunks_exact(4).map(|pix| index_of(pix)).collect()),
+            palette: Some(palette),
+            transparent: transparent.map(|t| index_of(&t)),
             ..Frame::default()
         }
     }

tests/roundtrip.rs

@@ -1,4 +1,4 @@
-use gif::{Decoder, Encoder, Frame};
+use gif::{ColorOutput, Decoder, Encoder, Frame};
 
 #[test]
 fn encode_roundtrip() {
@@ -53,3 +53,50 @@ fn round_trip_from_image(original: &[u8]) {
         }
     }
 }
+
+#[test]
+fn encode_roundtrip_few_colors() {
+    const WIDTH: u16 = 128;
+    const HEIGHT: u16 = 128;
+
+    // Build an image with a single red pixel, that NeuQuant won't
+    // sample, in order to check that we do appropriatelyq specialise the
+    // few-colors case.
+    let mut pixels: Vec<u8> = vec![255; WIDTH as usize * HEIGHT as usize * 4];
+    // Top-left pixel is always sampled, so use the second pixel.
+    pixels[5] = 0;
+    pixels[6] = 0;
+    // Set speed to 30 to handily avoid sampling that one pixel.
+    //
+    // We clone "pixels", since the parameter is replaced with a
+    // paletted version, and later we want to compare the output with
+    // the original RGBA image.
+    let frames = [Frame::from_rgba_speed(WIDTH, HEIGHT, &mut pixels.clone(), 30)];
+
+    let mut buffer = vec![];
+    {
+        let mut encoder = Encoder::new(&mut buffer, WIDTH, HEIGHT, &[]).unwrap();
+        for frame in &frames {
+            encoder.write_frame(frame).unwrap();
+        }
+    }
+
+    {
+        let mut decoder = {
+            let mut builder = Decoder::<&[u8]>::build();
+            builder.set_color_output(ColorOutput::RGBA);
+            builder.read_info(&buffer[..]).expect("Invalid info encoded")
+        };
+
+        // Only check key fields, assuming "round_trip_from_image"
+        // covers the rest. We are primarily concerned with quantisation.
+        assert_eq!(decoder.width(), WIDTH);
+        assert_eq!(decoder.height(), HEIGHT);
+        let new_frames: Vec<_> = core::iter::from_fn(move || {
+            decoder.read_next_frame().unwrap().cloned()
+        }).collect();
+        assert_eq!(new_frames.len(), 1, "Diverging number of frames");
+        // NB: reference.buffer can't be used as it contains the palette version.
+        assert_eq!(new_frames[0].buffer, pixels);
+    }
+}