factor out the quantizer

1 files changed, 62 insertions, 72 deletions

diff --git a/src/main.rs b/src/main.rs
index c421f48..b72e912 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,16 +1,18 @@
 use std::{collections::HashMap, env::args, time::Instant};
 
 use image::io::Reader as ImageReader;
-use image::Rgb;
+use image::{buffer::Pixels, Rgb};
+
+const MAX_COLORS: usize = 256;
+
+const TOLERANCE: f32 = 0.6;
+const RGB_TOLERANCE: f32 = 10.0 * TOLERANCE;
 
 fn main() {
     let before = Instant::now();
     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 tolerance = 0.6;
-    let rgb_tolerance = 10.0 * tolerance;
-    let max_colors = 256;
 
     println!("File is {}", &filename);
 
@@ -23,80 +25,23 @@ fn main() {
     println!("Decoded!");
     let before_algo = Instant::now();
 
-    let mut colors: HashMap<Rgb<u8>, usize> = HashMap::new();
-
-    //count pixels
-    for pixel in image.pixels() {
-        match colors.get_mut(pixel) {
-            None => {
-                colors.insert(*pixel, 1);
-            }
-            Some(n) => *n += 1,
-        }
-    }
-
-    println!(
-        "{} has {} colors in it. Sorting most occuring to least...",
-        filename,
-        colors.len()
-    );
-
-    let mut sorted: Vec<(Rgb<u8>, 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]))
-    });
-
-    println!("Sorted! Selecting colors...");
-
-    for (color, count) in sorted.iter().take(10) {
-        println!("{:?} count {}", color, count);
-    }
-
-    for (color, count) in sorted.iter().rev().take(10) {
-        println!("rev {:?} count {}", color, count);
-    }
-
-    let mut sorted_iter = sorted.iter();
-
-    let mut selected_colors: Vec<Rgb<u8>> = Vec::with_capacity(max_colors);
-    selected_colors.push(sorted_iter.next().unwrap().0);
-
-    for (key, _value) in sorted_iter {
-        if selected_colors.len() < max_colors {
-            for selected_color in selected_colors.iter() {
-                if rgb_difference(key, selected_color) > rgb_tolerance {
-                    selected_colors.push(*key);
-                    break;
-                }
-            }
-        } else {
-            break;
-        }
-    }
-
-    for color in selected_colors.iter().take(10) {
-        println!("selected {:?}", color);
-    }
+    let selected_colors = quantize(image.pixels());
 
     println!("Selected {} colors! Creating map...", selected_colors.len());
 
-    let mut color_map: HashMap<Rgb<u8>, Rgb<u8>> = HashMap::with_capacity(sorted.len());
+    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);
     }
 
     // Max complexity is O(n * max_colors)
-    'sorted_colors: for (key, _value) in sorted.iter() {
-        let mut min_difference = f64::MAX;
-        let mut min_difference_color = *key;
+    for color in image.pixels() {
+        let mut min_difference = f32::MAX;
+        let mut min_difference_color = *color;
 
         for index in 0..selected_colors.len() {
-            let difference = rgb_difference(key, unsafe { selected_colors.get_unchecked(index) });
+            let difference = rgb_difference(color, unsafe { selected_colors.get_unchecked(index) });
             /*if difference == 0.0 {
                 continue 'sorted_colors;
             }*/
@@ -106,7 +51,7 @@ fn main() {
             }
         }
 
-        color_map.insert(*key, min_difference_color);
+        color_map.insert(*color, min_difference_color);
     }
 
     println!("Mapped! Filling in image...");
@@ -120,7 +65,7 @@ fn main() {
         Instant::now().duration_since(before_algo).as_millis()
     );
 
-    let mut recounted_colors = Vec::with_capacity(max_colors);
+    let mut recounted_colors = Vec::with_capacity(MAX_COLORS);
     // Recount colors
     for pixel in image.pixels() {
         if !recounted_colors.contains(pixel) {
@@ -131,7 +76,7 @@ fn main() {
 
     println!(
         "Aiming for a max of {} colors, got {}",
-        max_colors,
+        MAX_COLORS,
         recounted_colors.len()
     );
 
@@ -142,7 +87,52 @@ fn main() {
     );
 }
 
-fn rgb_difference(a: &Rgb<u8>, z: &Rgb<u8>) -> f64 {
+fn quantize(pixels: Pixels<Rgb<u8>>) -> Vec<Rgb<u8>> {
+    let mut colors: HashMap<Rgb<u8>, usize> = HashMap::new();
+
+    //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_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]))
+    });
+
+    println!("Sorted! Selecting colors...");
+
+    let mut sorted_iter = sorted.iter();
+
+    let mut selected_colors: Vec<Rgb<u8>> = Vec::with_capacity(MAX_COLORS);
+    selected_colors.push(sorted_iter.next().unwrap().0);
+
+    for (key, _value) in sorted_iter {
+        if selected_colors.len() < MAX_COLORS {
+            for selected_color in selected_colors.iter() {
+                if rgb_difference(key, selected_color) > RGB_TOLERANCE {
+                    selected_colors.push(*key);
+                    break;
+                }
+            }
+        } else {
+            break;
+        }
+    }
+
+    selected_colors
+}
+
+fn rgb_difference(a: &Rgb<u8>, z: &Rgb<u8>) -> f32 {
     //((a.0[0] as i16 - b.0[0] as i16).abs() + (a.0[1] as i16 - b.0[1] as i16).abs() +(a.0[2] as i16 - b.0[2] as i16).abs()) as u16
     //(a.0[0] as i16 - b.0[0] as i16).abs().max((a.0[1] as i16 - b.0[1] as i16).abs().max(a.0[2] as i16 - b.0[2] as i16).abs()) as u16
     //(a.0[0] as i16 - b.0[0] as i16).abs().max((a.0[1] as i16 - b.0[1] as i16).abs()).max((a.0[2] as i16 - b.0[2] as i16).abs()) as u16
@@ -153,7 +143,7 @@ fn rgb_difference(a: &Rgb<u8>, z: &Rgb<u8>) -> f64 {
     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 f64
+    (((c - f) * (c - f)) + ((a - d).abs() / 90.0) + (b - e).abs()) as f32
 }
 
 #[warn(clippy::float_cmp)]