Libify colorquash

2 files changed, 233 insertions, 175 deletions

diff --git a/src/lib.rs b/src/lib.rs
new file mode 100644
index 0000000..8efd6ad
--- /dev/null
+++ b/src/lib.rs
@@ -0,0 +1,233 @@
+use std::collections::HashMap;
+use std::ops::Deref;
+
+use ahash::RandomState;
+
+pub struct Squasher<T> {
+    palette: Vec<(Rgb, usize)>,
+    larget_count: usize,
+    map: Vec<T>,
+}
+
+impl<T: Count> Squasher<T> {
+    /// Creates a new squasher and allocates a new color map. A color map
+    /// contains every 24-bit color and ends up with an amount of memory
+    /// equal to `16MB * std::mem::size_of(T)`
+    pub fn new(max_colors: T, buffer: &[u8]) -> Self {
+        let sorted = Self::unique_and_sort(buffer);
+        let selected = Self::select_colors(&sorted, max_colors);
+
+        let mut this = Self {
+            palette: selected,
+            larget_count: sorted.first().unwrap().1,
+            map: vec![T::zero(); 256 * 256 * 256],
+        };
+
+        this.map_selected(&sorted);
+
+        this
+    }
+
+    /// 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]) {
+        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]]))];
+
+            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()
+    }
+
+    /// Retrieve the palette as bytes
+    pub fn palette_bytes(&self) -> Vec<u8> {
+        self.palette
+            .clone()
+            .into_iter()
+            .map(|rgb| rgb.0.into_iter())
+            .flatten()
+            .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();
+
+        //count pixels
+        for pixel in buffer.chunks(3) {
+            let rgb = Rgb([pixel[0], pixel[1], pixel[2]]);
+
+            match colors.get_mut(&rgb) {
+                None => {
+                    colors.insert(rgb, 1);
+                }
+                Some(n) => *n += 1,
+            }
+        }
+
+        let mut sorted: Vec<(Rgb, usize)> = colors.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]))
+        });
+
+        sorted
+    }
+
+    fn select_colors(sorted: &[(Rgb, usize)], max_colors: T) -> Vec<(Rgb, usize)> {
+        #[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());
+
+        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)
+            {
+                selected_colors.push((*key, *count));
+            }
+        }
+
+        selected_colors
+    }
+
+    fn map_selected(&mut self, sorted: &[(Rgb, 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:?}");
+
+                if diff.max(0.0) < min_diff {
+                    min_diff = diff;
+                    min_index = index;
+                }
+            }
+
+            self.map[color_index(sorted)] = T::from_usize(min_index);
+        }
+    }
+}
+
+pub trait Count: Copy + Clone {
+    fn zero() -> Self;
+    fn as_usize(&self) -> usize;
+    fn from_usize(from: usize) -> Self;
+    fn le(&self, rhs: &usize) -> bool;
+}
+
+macro_rules! count_impl {
+    ($kind:ty) => {
+        impl Count for $kind {
+            fn zero() -> Self {
+                0
+            }
+
+            fn as_usize(&self) -> usize {
+                *self as usize
+            }
+
+            #[inline(always)]
+            fn from_usize(from: usize) -> Self {
+                from as Self
+            }
+
+            #[inline(always)]
+            fn le(&self, rhs: &usize) -> bool {
+                *self as usize <= *rhs
+            }
+        }
+    };
+}
+
+count_impl!(u8);
+count_impl!(u16);
+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
+}
+
+#[allow(clippy::many_single_char_names)]
+#[inline(always)]
+fn rgb_difference(a: &Rgb, b: &Rgb) -> 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)
+}
diff --git a/src/main.rs b/src/main.rs
deleted file mode 100644
index f1915f8..0000000
--- a/src/main.rs
+++ /dev/null
@@ -1,175 +0,0 @@
-use std::time::Instant;
-use std::{collections::HashMap, env::args};
-
-use image::io::Reader as ImageReader;
-use image::Rgb;
-
-use ahash::RandomState;
-
-use rayon::prelude::*;
-
-const MAX_COLORS: usize = 16;
-
-const RGB_TOLERANCE: f32 = 0.25 + (1.0 - (MAX_COLORS as f32 / 256.0));
-
-fn main() {
-    let filename = args().nth(1).unwrap();
-    let outname = args().nth(2).unwrap();
-    // The percent of RGB value difference a color has to surpass to be considered unique
-
-    let imageread = ImageReader::open(&filename).expect("Failed to open image!");
-    let mut image = imageread
-        .decode()
-        .expect("Failed to decode image!")
-        .into_rgb8();
-
-    //let mem_before_sort = mallinfo().hblkhd as usize;
-    let start_sort = Instant::now();
-    let sorted_colors = unique_and_sort(image.pixels());
-    println!("Sort took {}s", start_sort.elapsed().as_secs_f32());
-
-    //let mem_before_selection = mallinfo().hblkhd as usize;
-    let start_selection = Instant::now();
-    let selected_colors = select_colors(&sorted_colors);
-    println!(
-        "Color Selection took {}s. Count {}",
-        start_selection.elapsed().as_secs_f32(),
-        selected_colors.len()
-    );
-
-    let start_array = Instant::now();
-    let mut array = vec![0usize; 256 * 256 * 256];
-    println!(
-        "Array creation took {}s",
-        start_array.elapsed().as_secs_f32()
-    );
-
-    let start_map = Instant::now();
-
-    for (sorted, _) in &sorted_colors {
-        let mut min_diff = f32::MAX;
-        let mut min_index = usize::MAX;
-
-        for (index, selected) in selected_colors.iter().enumerate() {
-            let diff = rgb_difference(sorted, selected);
-            if diff < min_diff {
-                min_diff = diff;
-                min_index = index;
-            }
-        }
-
-        array[color_index(sorted)] = min_index;
-    }
-
-    println!(
-        "Creating color map {:.2}s",
-        start_map.elapsed().as_secs_f32()
-    );
-
-    let start_fill = Instant::now();
-    // Max complexity is O(n * max_colors)
-    for color in image.pixels_mut() {
-        let index = array[color_index(color)];
-
-        *color = selected_colors[index];
-    }
-    println!(
-        "Took {:.2}s to fill in the image.\nTotal time from sort {:.2}s",
-        start_fill.elapsed().as_secs_f32(),
-        start_sort.elapsed().as_secs_f32()
-    );
-
-    image.save(outname).expect("Failed to write out");
-}
-
-fn unique_and_sort<'a, T>(pixels: T) -> Vec<(Rgb<u8>, usize)>
-where
-    T: Iterator<Item = &'a Rgb<u8>>,
-{
-    let mut colors: HashMap<Rgb<u8>, usize, RandomState> = HashMap::default();
-
-    //count pixels
-    for pixel in pixels {
-        match colors.get_mut(pixel) {
-            None => {
-                colors.insert(*pixel, 1);
-            }
-            Some(n) => *n += 1,
-        }
-    }
-
-    let mut sorted: Vec<(Rgb<u8>, usize)> = colors.into_par_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]))
-    });
-
-    sorted
-}
-
-fn select_colors(sorted: &[(Rgb<u8>, usize)]) -> Vec<Rgb<u8>> {
-    let mut selected_colors: Vec<Rgb<u8>> = Vec::with_capacity(MAX_COLORS);
-
-    for (key, _value) in sorted.iter() {
-        if selected_colors.len() >= MAX_COLORS {
-            break;
-        } else if selected_colors
-            .iter()
-            .all(|color| rgb_difference(key, color) > RGB_TOLERANCE)
-        {
-            selected_colors.push(*key);
-        }
-    }
-
-    selected_colors
-}
-
-#[inline(always)]
-fn color_index(c: &Rgb<u8>) -> usize {
-    c.0[0] as usize * (256 * 256) + c.0[1] as usize * 256 + c.0[2] as usize
-}
-
-#[allow(clippy::many_single_char_names)]
-#[inline(always)]
-fn rgb_difference(a: &Rgb<u8>, z: &Rgb<u8>) -> f32 {
-    let (a, b, c) = pixel_rgb_to_hsv(a);
-    let (d, e, f) = pixel_rgb_to_hsv(z);
-
-    (((c - f) * (c - f)) + ((a - d).abs() / 90.0) + (b - e).abs()) as f32
-}
-
-#[allow(clippy::float_cmp)]
-fn pixel_rgb_to_hsv(a: &Rgb<u8>) -> (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)
-}