From f547a21f005a07fdc0d1f98e121ebf11b6d241dd Mon Sep 17 00:00:00 2001
From: CrispyPin <crispin@tasa.se>
Date: Wed, 4 Dec 2024 15:20:41 +0100
Subject: [PATCH] initial dithering implemetation

---
 src/main.rs | 59 ++++++++++++++++++++++++++++++++++++++++++++---------
 1 file changed, 49 insertions(+), 10 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 2e545f4..af48a55 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -18,19 +18,15 @@ fn main() {
 
     let (w, h) = image.dimensions();
     let image = image.resize(w / 8, h / 8, FilterType::CatmullRom);
+    image.save("scaled.png").unwrap();
     println!("resized");
-    let image = dither_limit(image, 2);
+    // let image = quantize_image(image, 6);
+    let image = dither(image, 12);
     image.save("out.png").unwrap();
     println!("saved");
 }
 
-const BAYER_4x4: [[u8; 4]; 4] = [[0, 8, 2, 10], [12, 4, 14, 6], [3, 11, 1, 9], [15, 7, 13, 5]];
-
-fn bayer_f(x: usize, y: usize) -> f32 {
-    (BAYER_4x4[y][x] as f32 / 16.) // - (15./16.)/2.
-}
-
-fn quantize(image: &DynamicImage, count: usize) -> Vec<Rgb<u8>> {
+fn generate_palette(image: &DynamicImage, count: usize) -> Vec<Rgb<u8>> {
     assert!(count > 0);
 
     let mut buckets: Vec<Vec<Rgb<u8>>> = vec![image
@@ -116,8 +112,8 @@ fn color_dist(a: Rgb<u8>, b: Rgb<u8>) -> f32 {
     (r * r + g * g + b * b).sqrt()
 }
 
-fn dither_limit(input: DynamicImage, count: usize) -> DynamicImage {
-    let colors = quantize(&input, count);
+fn quantize_image(input: DynamicImage, count: usize) -> DynamicImage {
+    let colors = generate_palette(&input, count);
     let mut out = input.clone();
 
     for (x, y, color) in input.pixels() {
@@ -136,3 +132,46 @@ fn dither_limit(input: DynamicImage, count: usize) -> DynamicImage {
 
     out
 }
+
+const BAYER_4X4: [[u8; 4]; 4] = [[0, 8, 2, 10], [12, 4, 14, 6], [3, 11, 1, 9], [15, 7, 13, 5]];
+
+fn bayer_f(x: usize, y: usize) -> f32 {
+    (BAYER_4X4[y][x] as f32 / 16.) - (15. / 16.) / 2.
+}
+
+fn dither(input: DynamicImage, count: usize) -> DynamicImage {
+    let mut colors = generate_palette(&input, count);
+    // let mut colors = vec![
+    //     Rgb([177, 138, 129]),
+    //     Rgb([155, 68, 52]),
+    //     Rgb([234, 130, 70]),
+    //     Rgb([110, 75, 72]),
+    //     Rgb([60, 43, 41]),
+    //     Rgb([0, 0, 0]),
+    // ];
+    let (w, h) = input.dimensions();
+    let mut out = DynamicImage::new(w, h, input.color());
+
+    for (x, y, color) in input.pixels() {
+        let real_col = color.to_rgb();
+        colors.sort_unstable_by_key(|&c| (color_dist(c, real_col) * 100.) as u32);
+
+        let best_col = colors[0];
+        let second_col = colors[1];
+
+        let best_dist = color_dist(real_col, best_col);
+        let second_dist = color_dist(real_col, second_col);
+
+        let ratio = second_dist / (second_dist + best_dist) * 2. - 1.;
+
+        let bx = x as usize % 4;
+        let by = y as usize % 4;
+        let bayer = bayer_f(bx, by);
+
+        let ratio = ratio + bayer * 1.5;
+        let out_col = if ratio > 0. { best_col } else { second_col };
+
+        out.put_pixel(x, y, out_col.to_rgba());
+    }
+    out
+}