allow changing difference algorithm

7 files changed, 192 insertions, 158 deletions

diff --git a/.rustfmt.toml b/.rustfmt.toml
new file mode 100644
index 0000000..218e203
--- /dev/null
+++ b/.rustfmt.toml
@@ -0,0 +1 @@
+hard_tabs = true
diff --git a/Cargo.lock b/Cargo.lock
index 76ad636..f80a769 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3,15 +3,28 @@
 version = 3
 
 [[package]]
-name = "ahash"
-version = "0.7.6"
+name = "adler"
+version = "1.0.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fcb51a0695d8f838b1ee009b3fbf66bda078cd64590202a864a8f3e8c4315c47"
-dependencies = [
- "getrandom",
- "once_cell",
- "version_check",
-]
+checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe"
+
+[[package]]
+name = "bitflags"
+version = "1.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a"
+
+[[package]]
+name = "bytemuck"
+version = "1.14.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "374d28ec25809ee0e23827c2ab573d729e293f281dfe393500e7ad618baa61c6"
+
+[[package]]
+name = "camino"
+version = "1.1.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c59e92b5a388f549b863a7bea62612c09f24c8393560709a54558a9abdfb3b9c"
 
 [[package]]
 name = "cfg-if"
@@ -23,40 +36,79 @@ checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
 name = "colorsquash"
 version = "0.1.0"
 dependencies = [
- "ahash",
+ "rgb",
 ]
 
 [[package]]
-name = "getrandom"
-version = "0.2.8"
+name = "crc32fast"
+version = "1.3.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c05aeb6a22b8f62540c194aac980f2115af067bfe15a0734d7277a768d396b31"
+checksum = "b540bd8bc810d3885c6ea91e2018302f68baba2129ab3e88f32389ee9370880d"
 dependencies = [
  "cfg-if",
- "libc",
- "wasi",
 ]
 
 [[package]]
-name = "libc"
-version = "0.2.138"
+name = "fdeflate"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d329bdeac514ee06249dabc27877490f17f5d371ec693360768b838e19f3ae10"
+dependencies = [
+ "simd-adler32",
+]
+
+[[package]]
+name = "flate2"
+version = "1.0.27"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "db6d7e329c562c5dfab7a46a2afabc8b987ab9a4834c9d1ca04dc54c1546cef8"
+checksum = "c6c98ee8095e9d1dcbf2fcc6d95acccb90d1c81db1e44725c6a984b1dbdfb010"
+dependencies = [
+ "crc32fast",
+ "miniz_oxide",
+]
+
+[[package]]
+name = "miniz_oxide"
+version = "0.7.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e7810e0be55b428ada41041c41f32c9f1a42817901b4ccf45fa3d4b6561e74c7"
+dependencies = [
+ "adler",
+ "simd-adler32",
+]
 
 [[package]]
-name = "once_cell"
-version = "1.16.0"
+name = "png"
+version = "0.17.10"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "86f0b0d4bf799edbc74508c1e8bf170ff5f41238e5f8225603ca7caaae2b7860"
+checksum = "dd75bf2d8dd3702b9707cdbc56a5b9ef42cec752eb8b3bafc01234558442aa64"
+dependencies = [
+ "bitflags",
+ "crc32fast",
+ "fdeflate",
+ "flate2",
+ "miniz_oxide",
+]
 
 [[package]]
-name = "version_check"
-version = "0.9.4"
+name = "rgb"
+version = "0.8.36"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f"
+checksum = "20ec2d3e3fc7a92ced357df9cebd5a10b6fb2aa1ee797bf7e9ce2f17dffc8f59"
+dependencies = [
+ "bytemuck",
+]
 
 [[package]]
-name = "wasi"
-version = "0.11.0+wasi-snapshot-preview1"
+name = "simd-adler32"
+version = "0.3.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423"
+checksum = "d66dc143e6b11c1eddc06d5c423cfc97062865baf299914ab64caa38182078fe"
+
+[[package]]
+name = "squash"
+version = "0.1.0"
+dependencies = [
+ "camino",
+ "png",
+]
diff --git a/Cargo.toml b/Cargo.toml
index e1a36ba..dc8c6d0 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,20 +1,24 @@
 [package]
 name = "colorsquash"
 version = "0.1.0"
-authors = ["Genny Z <gen@nyble.dev>"]
+authors = ["gennyble <gen@nyble.dev>"]
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
 #image = "0.23.14"
-ahash = "0.7.4"
+#ahash = "0.7.4"
 #libc = "0.2.103"
 #rayon = "*"
+rgb = "0.8.36"
 
 #[dependencies.kmeans_colors]
 #version = "0.3"
 #default-features = false
 
 [profile.release]
-debug = true
\ No newline at end of file
+debug = true
+
+[workspace]
+members = ["squash"]
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..0db61d0
--- /dev/null
+++ b/README.md
@@ -0,0 +1,14 @@
+Colorsquash is a colour quantization[^1] crate and algorithm.
+At it's core, it sorts the unique colours that appear in an image
+and selects the most frequent that are sufficiently different.
+
+To put it more clearly:  
+The most frequent colour is always selected and placed into the palette.
+If the second most frequent colour is *different enough*, it will be selected
+as well. If it's not, it is skipped and the third one is tried. This continues
+until it selects the necessary amount of colours.
+
+[^1] (wikipedia: color quantization)[https://en.wikipedia.org/wiki/Color_quantization]
+
+### squash
+A CLI tool to quantize colours. Accepts a path to a PNG; exports an indexed PNG.
\ No newline at end of file
diff --git a/squash/Cargo.toml b/squash/Cargo.toml
new file mode 100644
index 0000000..18f2be3
--- /dev/null
+++ b/squash/Cargo.toml
@@ -0,0 +1,10 @@
+[package]
+name = "squash"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+camino = "1.1.6"
+png = "0.17.10"
diff --git a/squash/src/main.rs b/squash/src/main.rs
new file mode 100644
index 0000000..f328e4d
--- /dev/null
+++ b/squash/src/main.rs
@@ -0,0 +1 @@
+fn main() {}
diff --git a/src/lib.rs b/src/lib.rs
index 3af73f1..a404f38 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,43 +1,12 @@
+use rgb::RGB8;
 use std::collections::HashMap;
-use std::ops::Deref;
 
-use ahash::RandomState;
-
-pub struct ColorCollector {
-    colors: HashMap<Rgb, usize, RandomState>,
-}
-
-impl ColorCollector {
-    pub fn new() -> Self {
-        Self {
-            colors: HashMap::default(),
-        }
-    }
-
-    /// Wants an RGB buffer
-    pub fn add(&mut self, buffer: &[u8]) {
-        for pixel in buffer.chunks(3) {
-            let rgb = Rgb([pixel[0], pixel[1], pixel[2]]);
-
-            match self.colors.get_mut(&rgb) {
-                None => {
-                    self.colors.insert(rgb, 1);
-                }
-                Some(n) => *n += 1,
-            }
-        }
-    }
-
-    pub fn as_squasher<T: Count>(self, max_colors: T) -> Squasher<T> {
-        let sorted = Squasher::<T>::sort(self.colors);
-        Squasher::from_sorted(max_colors, sorted)
-    }
-}
+type DiffFn = dyn Fn(&RGB8, &RGB8) -> f32;
 
 pub struct Squasher<T> {
-    palette: Vec<(Rgb, usize)>,
-    larget_count: usize,
+    palette: Vec<RGB8>,
     map: Vec<T>,
+    difference_fn: Box<DiffFn>,
 }
 
 impl<T: Count> Squasher<T> {
@@ -46,16 +15,28 @@ impl<T: Count> Squasher<T> {
     /// equal to `16MB * std::mem::size_of(T)`
     pub fn new(max_colors: T, buffer: &[u8]) -> Self {
         let sorted = Self::unique_and_sort(buffer);
-        Self::from_sorted(max_colors, sorted)
+        Self::from_sorted(max_colors, sorted, Box::new(rgb_difference))
+    }
+
+    /// Like [Squasher::new] but lets you pass your own difference function
+    /// to compare values while selecting colours. The default difference
+    /// function sums to difference between the RGB channels.
+    pub fn new_with_difference(
+        max_colors: T,
+        buffer: &[u8],
+        difference_fn: &'static DiffFn,
+    ) -> Self {
+        let sorted = Self::unique_and_sort(buffer);
+        Self::from_sorted(max_colors, sorted, Box::new(difference_fn))
     }
 
-    fn from_sorted(max_colors: T, sorted: Vec<(Rgb, usize)>) -> Self {
-        let selected = Self::select_colors(&sorted, max_colors);
+    fn from_sorted(max_colors: T, sorted: Vec<(RGB8, usize)>, difference_fn: Box<DiffFn>) -> Self {
+        let selected = Self::select_colors(&sorted, max_colors, difference_fn.as_ref());
 
         let mut this = Self {
             palette: selected,
-            larget_count: sorted.first().unwrap().1,
             map: vec![T::zero(); 256 * 256 * 256],
+            difference_fn,
         };
 
         this.map_selected(&sorted);
@@ -64,34 +45,44 @@ impl<T: Count> Squasher<T> {
     }
 
     /// Take an RGB image buffer and an output buffer. The function will fill
-    /// the output buffer with indexes into the Palette.
-    pub fn map_image(&mut self, image: &[u8], buffer: &mut [T]) {
+    /// the output buffer with indexes into the Palette. The output buffer should
+    /// be a third of the size of the image buffer.
+    pub fn map(&mut self, image: &[u8], buffer: &mut [T]) {
+        if buffer.len() * 3 < image.len() {
+            panic!("outout buffer too small to fit indexed image");
+        }
+
         // We have to map the colours of this image now because it might contain
         // colours not present in the first image.
         let sorted = Self::unique_and_sort(image);
         self.map_selected(&sorted);
 
         for (idx, color) in image.chunks(3).enumerate() {
-            let index = self.map[color_index(&Rgb([color[0], color[1], color[2]]))];
+            let index = self.map[color_index(&RGB8::new(color[0], color[1], color[2]))];
 
             buffer[idx] = index;
         }
     }
 
-    /// Take an RGB image buffer and an output buffer. The function will fill
-    /// the output buffer with indexes into the Palette.
+    /// Like [Squasher::map] but it doesn't recount the input image. This will
+    /// cause colors the Squasher hasn't seen before to come out as index 0 which
+    /// may be incorrect.
     //TODO: gen- Better name?
     pub fn map_unsafe(&self, image: &[u8], buffer: &mut [T]) {
+        if buffer.len() * 3 < image.len() {
+            panic!("outout buffer too small to fit indexed image");
+        }
+
         for (idx, color) in image.chunks(3).enumerate() {
-            let index = self.map[color_index(&Rgb([color[0], color[1], color[2]]))];
+            let index = self.map[color_index(&RGB8::new(color[0], color[1], color[2]))];
 
             buffer[idx] = index;
         }
     }
 
     /// Retrieve the palette this squasher is working from
-    pub fn palette(&self) -> Vec<Rgb> {
-        self.palette.iter().map(|ahh| ahh.0).collect()
+    pub fn palette(&self) -> &[RGB8] {
+        &self.palette
     }
 
     /// Retrieve the palette as bytes
@@ -99,18 +90,17 @@ impl<T: Count> Squasher<T> {
         self.palette
             .clone()
             .into_iter()
-            .map(|rgb| rgb.0.into_iter())
-            .flatten()
+            .flat_map(|rgb| [rgb.r, rgb.g, rgb.b].into_iter())
             .collect()
     }
 
     /// Takes an image buffer of RGB data and fill the color map
-    fn unique_and_sort(buffer: &[u8]) -> Vec<(Rgb, usize)> {
-        let mut colors: HashMap<Rgb, usize, RandomState> = HashMap::default();
+    fn unique_and_sort(buffer: &[u8]) -> Vec<(RGB8, usize)> {
+        let mut colors: HashMap<RGB8, usize> = HashMap::default();
 
         //count pixels
         for pixel in buffer.chunks(3) {
-            let rgb = Rgb([pixel[0], pixel[1], pixel[2]]);
+            let rgb = RGB8::new(pixel[0], pixel[1], pixel[2]);
 
             match colors.get_mut(&rgb) {
                 None => {
@@ -123,53 +113,48 @@ impl<T: Count> Squasher<T> {
         Self::sort(colors)
     }
 
-    fn sort(map: HashMap<Rgb, usize, RandomState>) -> Vec<(Rgb, usize)> {
-        let mut sorted: Vec<(Rgb, usize)> = map.into_iter().collect();
+    fn sort(map: HashMap<RGB8, usize>) -> Vec<(RGB8, usize)> {
+        let mut sorted: Vec<(RGB8, usize)> = map.into_iter().collect();
         sorted.sort_by(|(colour1, freq1), (colour2, freq2)| {
             freq2
                 .cmp(freq1)
-                .then(colour2[0].cmp(&colour1[0]))
-                .then(colour2[1].cmp(&colour1[1]))
-                .then(colour2[2].cmp(&colour1[2]))
+                .then(colour2.r.cmp(&colour1.r))
+                .then(colour2.g.cmp(&colour1.g))
+                .then(colour2.b.cmp(&colour1.b))
         });
 
         sorted
     }
 
-    fn select_colors(sorted: &[(Rgb, usize)], max_colors: T) -> Vec<(Rgb, usize)> {
+    fn select_colors(sorted: &[(RGB8, usize)], max_colors: T, difference: &DiffFn) -> Vec<RGB8> {
         #[allow(non_snake_case)]
-        let RGB_TOLERANCE: f32 = 0.04 * 256.0;
-        let mut selected_colors: Vec<(Rgb, usize)> = Vec::with_capacity(max_colors.as_usize());
+        let RGB_TOLERANCE: f32 = 0.01 * 768.0;
+        let mut selected_colors: Vec<(RGB8, usize)> = Vec::with_capacity(max_colors.as_usize());
 
         for (key, count) in sorted.iter() {
             if max_colors.le(&selected_colors.len()) {
                 break;
             } else if selected_colors
                 .iter()
-                .all(|color| rgb_difference(key, &color.0) > RGB_TOLERANCE)
+                .all(|color| difference(key, &color.0) > RGB_TOLERANCE)
             {
                 selected_colors.push((*key, *count));
             }
         }
 
         selected_colors
+            .into_iter()
+            .map(|(color, _count)| color)
+            .collect()
     }
 
-    fn map_selected(&mut self, sorted: &[(Rgb, usize)]) {
+    fn map_selected(&mut self, sorted: &[(RGB8, usize)]) {
         for (sorted, _) in sorted {
             let mut min_diff = f32::MAX;
             let mut min_index = usize::MAX;
 
-            for (index, (selected, count)) in self.palette.iter().enumerate() {
-                //let count_weight = *count as f32 / self.larget_count as f32;
-                let diff = rgb_difference(sorted, selected); // - count_weight * 64.0;
-
-                // This is kind of racist genny
-                /*if selected[0] + selected[1] + selected[2] < 72 {
-                    continue;
-                }*/
-
-                //println!("{diff} - {selected:?}");
+            for (index, selected) in self.palette.iter().enumerate() {
+                let diff = (self.difference_fn)(sorted, selected);
 
                 if diff.max(0.0) < min_diff {
                     min_diff = diff;
@@ -182,6 +167,24 @@ impl<T: Count> Squasher<T> {
     }
 }
 
+impl Squasher<u8> {
+    /// Takes an RGB image buffer and writes the indicies to the first third of
+    /// that buffer. The buffer is not resized.
+    ///
+    /// # Returns
+    /// The new size of the image
+    pub fn map_over(&self, image: &mut [u8]) -> usize {
+        for idx in 0..image.len() {
+            let rgb_idx = idx * 3;
+            let color = RGB8::new(image[rgb_idx], image[rgb_idx + 1], image[rgb_idx + 2]);
+            let color_index = self.map[color_index(&color)];
+            image[idx] = color_index;
+        }
+
+        image.len() / 3
+    }
+}
+
 pub trait Count: Copy + Clone {
     fn zero() -> Self;
     fn as_usize(&self) -> usize;
@@ -219,67 +222,16 @@ count_impl!(u32);
 count_impl!(u64);
 count_impl!(usize);
 
-#[derive(Copy, Clone, Debug, Eq, Hash, PartialEq)]
-pub struct Rgb([u8; 3]);
-
-impl Deref for Rgb {
-    type Target = [u8; 3];
-
-    fn deref(&self) -> &Self::Target {
-        &self.0
-    }
-}
-
 #[inline(always)]
-fn color_index(c: &Rgb) -> usize {
-    c.0[0] as usize * (256 * 256) + c.0[1] as usize * 256 + c.0[2] as usize
+fn color_index(c: &RGB8) -> usize {
+    c.r as usize * (256 * 256) + c.g as usize * 256 + c.b as usize
 }
 
+/// The default comparison function. Returns a sum of the channel differences.
+/// I.E. `|a.red - b.red| + |a.green - b.green| + |a.blue - b.blue|`
 #[allow(clippy::many_single_char_names)]
 #[inline(always)]
-fn rgb_difference(a: &Rgb, b: &Rgb) -> f32 {
+pub fn rgb_difference(a: &RGB8, b: &RGB8) -> f32 {
     let absdiff = |a: u8, b: u8| (a as f32 - b as f32).abs();
-
-    /*let hsv1 = pixel_rgb_to_hsv(a);
-    let hsv2 = pixel_rgb_to_hsv(b);*/
-
-    //let diff_max = 3.0;
-
-    absdiff(a.0[0], b.0[0]) + absdiff(a.0[1], b.0[1]) + absdiff(a.0[2], b.0[2])
-    /*(((hsv1.0 / 90.0) - (hsv2.0 / 90.0)).abs()
-    + (hsv1.1 - hsv2.1).abs()
-    + ((hsv1.2 - hsv1.2).abs()))
-    / diff_max*/
-}
-
-fn pixel_rgb_to_hsv(a: &Rgb) -> (f32, f32, f32) {
-    let (r, g, b) = (
-        a.0[0] as f32 / 256.0,
-        a.0[1] as f32 / 256.0,
-        a.0[2] as f32 / 256.0,
-    );
-
-    let value = r.max(g.max(b));
-    let x_min = r.min(g.min(b));
-    let chroma = value - x_min;
-
-    let hue = if chroma == 0.0 {
-        0.0
-    } else if value == r {
-        60.0 * ((g - b) / chroma)
-    } else if value == g {
-        60.0 * (2.0 + (b - r) / chroma)
-    } else if value == b {
-        60.0 * (4.0 + (r - g) / chroma)
-    } else {
-        unreachable!()
-    };
-
-    let value_saturation = if value == 0.0 { 0.0 } else { chroma / value };
-
-    /* Rotate the color wheel counter clockwise to the negative location
-          |       Keep the wheel in place and remove any full rotations
-     _____V____ _____V____
-    |          |          |*/
-    ((hue + 360.0) % 360.0, value_saturation * 2.0, value * 2.0)
+    absdiff(a.r, b.r) + absdiff(a.g, b.g) + absdiff(a.b, b.b)
 }