Rework color map creation, remove parallelation

1 files changed, 64 insertions, 24 deletions

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 {