Rework color map creation, remove parallelation

4 files changed, 114 insertions, 27 deletions

diff --git a/.gitignore b/.gitignore
index adc528b..4c7d6ed 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,8 @@
 /target
 *.png
+*.PNG
+*.jpg
+*.JPG
 flamegraph.svg
 perf.data
 perf.data.old
diff --git a/Cargo.lock b/Cargo.lock
index 7b23b4d..0c011f8 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -15,6 +15,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "aae1277d39aeec15cb388266ecc24b11c80469deae6067e17a1a7aa9e5c1f234"
 
 [[package]]
+name = "ahash"
+version = "0.7.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "43bb833f0bf979d8475d38fbf09ed3b8a55e1885fe93ad3f93239fc6a4f17b98"
+dependencies = [
+ "getrandom 0.2.3",
+ "once_cell",
+ "version_check",
+]
+
+[[package]]
 name = "autocfg"
 version = "1.0.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -54,8 +65,10 @@ checksum = "3d7b894f5411737b7867f4827955924d7c254fc9f4d91a6aad6b097804b1018b"
 name = "colors"
 version = "0.1.0"
 dependencies = [
+ "ahash",
  "image",
  "kmeans_colors",
+ "libc",
  "rayon",
 ]
 
@@ -136,7 +149,18 @@ checksum = "8fc3cb4d91f53b50155bdcfd23f6a4c39ae1969c2ae85982b135750cccaf5fce"
 dependencies = [
  "cfg-if",
  "libc",
- "wasi",
+ "wasi 0.9.0+wasi-snapshot-preview1",
+]
+
+[[package]]
+name = "getrandom"
+version = "0.2.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7fcd999463524c52659517fe2cea98493cfe485d10565e7b0fb07dbba7ad2753"
+dependencies = [
+ "cfg-if",
+ "libc",
+ "wasi 0.10.2+wasi-snapshot-preview1",
 ]
 
 [[package]]
@@ -288,6 +312,12 @@ dependencies = [
 ]
 
 [[package]]
+name = "once_cell"
+version = "1.8.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "692fcb63b64b1758029e0a96ee63e049ce8c5948587f2f7208df04625e5f6b56"
+
+[[package]]
 name = "png"
 version = "0.16.8"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -311,7 +341,7 @@ version = "0.7.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6a6b1679d49b24bbfe0c803429aa1874472f50d9b363131f0e89fc356b544d03"
 dependencies = [
- "getrandom",
+ "getrandom 0.1.16",
  "libc",
  "rand_chacha",
  "rand_core",
@@ -334,7 +364,7 @@ version = "0.5.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "90bde5296fc891b0cef12a6d03ddccc162ce7b2aff54160af9338f8d40df6d19"
 dependencies = [
- "getrandom",
+ "getrandom 0.1.16",
 ]
 
 [[package]]
@@ -395,12 +425,24 @@ dependencies = [
 ]
 
 [[package]]
+name = "version_check"
+version = "0.9.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5fecdca9a5291cc2b8dcf7dc02453fee791a280f3743cb0905f8822ae463b3fe"
+
+[[package]]
 name = "wasi"
 version = "0.9.0+wasi-snapshot-preview1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "cccddf32554fecc6acb585f82a32a72e28b48f8c4c1883ddfeeeaa96f7d8e519"
 
 [[package]]
+name = "wasi"
+version = "0.10.2+wasi-snapshot-preview1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fd6fbd9a79829dd1ad0cc20627bf1ed606756a7f77edff7b66b7064f9cb327c6"
+
+[[package]]
 name = "weezl"
 version = "0.1.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/Cargo.toml b/Cargo.toml
index a116a39..9a2c9ad 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -8,6 +8,8 @@ edition = "2021"
 
 [dependencies]
 image = "0.23.14"
+ahash = "0.7.4"
+libc = "0.2.103"
 rayon = "*"
 
 [dependencies.kmeans_colors]
diff --git a/src/main.rs b/src/main.rs
index 7b2ae62..f1915f8 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,14 +1,16 @@
+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 = 256;
+const MAX_COLORS: usize = 16;
 
-const TOLERANCE: f32 = 0.025;
-const RGB_TOLERANCE: f32 = 10.0 * TOLERANCE;
+const RGB_TOLERANCE: f32 = 0.25 + (1.0 - (MAX_COLORS as f32 / 256.0));
 
 fn main() {
     let filename = args().nth(1).unwrap();
@@ -21,41 +23,70 @@ fn main() {
         .expect("Failed to decode image!")
         .into_rgb8();
 
-    let selected_colors = quantize(image.pixels());
+    //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;
+            }
+        }
 
-    let mut color_map: HashMap<Rgb<u8>, Rgb<u8>> = HashMap::with_capacity(image.len() / 2);
-    // Selected colors are themselves
-    for color in selected_colors.iter() {
-        color_map.insert(*color, *color);
+        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 quantized = color_map.entry(*color).or_insert({
-            let mut min_difference = f32::MAX;
-            let mut min_difference_color = *color;
-
-            for selected_color in &selected_colors {
-                let difference = rgb_difference(color, selected_color);
-                if difference < min_difference {
-                    min_difference = difference;
-                    min_difference_color = *selected_color;
-                }
-            }
-            min_difference_color
-        });
+        let index = array[color_index(color)];
 
-        *color = *quantized;
+        *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 quantize<'a, T>(pixels: T) -> Vec<Rgb<u8>>
+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> = HashMap::new();
+    let mut colors: HashMap<Rgb<u8>, usize, RandomState> = HashMap::default();
 
     //count pixels
     for pixel in pixels {
@@ -76,6 +107,10 @@ where
             .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() {
@@ -92,6 +127,11 @@ where
     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 {