From 61d593f90109cf8c26f98825156327d26fd37008 Mon Sep 17 00:00:00 2001
From: Chris Mounce <christophermounce@gmail.com>
Date: Sun, 11 Aug 2024 00:52:20 -0700
Subject: [PATCH] Implement quadtree color quantization

---
 rust/Cargo.toml  |  3 ++
 rust/src/main.rs | 75 +++++++++++++++++++++++++++++++++++++++---------
 2 files changed, 65 insertions(+), 13 deletions(-)

diff --git a/rust/Cargo.toml b/rust/Cargo.toml
index 4907ccf..512fea9 100644
--- a/rust/Cargo.toml
+++ b/rust/Cargo.toml
@@ -7,3 +7,6 @@ edition = "2021"
 anyhow = "1.0.86"
 clap = { version = "4.5.15", features = ["derive"] }
 image = "0.25.2"
+
+[profile.dev.package.zune-jpeg]
+opt-level = 3
diff --git a/rust/src/main.rs b/rust/src/main.rs
index d077a84..aece64b 100644
--- a/rust/src/main.rs
+++ b/rust/src/main.rs
@@ -9,6 +9,9 @@ struct Args {
     input: String,
     #[arg(short, long, value_name = "FILE")]
     output: String,
+
+    #[arg(short = 'e', long, value_name = "NUM")]
+    max_err: Option<f32>,
 }
 
 fn main() -> Result<()> {
@@ -22,22 +25,68 @@ fn main() -> Result<()> {
     let mut output = RgbImage::new(128, 128);
     for x in 0..=127 {
         for y in 0..=127 {
-            let [r, g, b] = img
-                .get_pixel(x * img.width() / 128, y * img.height() / 128)
-                .0;
-            let luma = r / 4 + g / 2 + b / 4;
-            output.put_pixel(
-                x,
-                y,
-                Rgb([
-                    (luma as i32 + 64 - x as i32).clamp(0, 255) as u8,
-                    luma,
-                    (luma as i32 + 64 - y as i32).clamp(0, 255) as u8,
-                ]),
-            );
+            let pixel = img.get_pixel(x * img.width() / 128, y * img.height() / 128);
+            output.put_pixel(x, y, *pixel);
         }
     }
+    quadtree_quant(
+        &mut output,
+        args.max_err.unwrap_or(2000.0) * 128.0 * 128.0,
+        [0, 0],
+        128,
+    );
     output.save(&args.output)?;
     println!("Wrote {}", &args.output);
     Ok(())
 }
+
+fn quadtree_quant(img: &mut RgbImage, max_err: f32, pos: [u32; 2], size: u32) {
+    // Calculate average color
+    let mut avg = [0.0, 0.0, 0.0];
+    for y in pos[1]..pos[1] + size {
+        for x in pos[0]..pos[0] + size {
+            let pixel = img.get_pixel(x, y);
+            for (i, component) in avg.iter_mut().enumerate() {
+                *component += (pixel.0[i] as f32).powf(2.2);
+            }
+        }
+    }
+    for component in &mut avg {
+        *component /= (size * size) as f32;
+    }
+    let avg = avg.map(|x| x.powf(0.4545).round() as u8);
+
+    // Measure squared error
+    let mut err = 0.0;
+    for y in pos[1]..pos[1] + size {
+        for x in pos[0]..pos[0] + size {
+            let pixel = img.get_pixel(x, y);
+            err += redmean_sq(avg, pixel.0);
+        }
+    }
+    if err > max_err {
+        let size = size / 2;
+        let max_err = max_err / 4.0;
+        let [x, y] = pos;
+        quadtree_quant(img, max_err, [x, y], size);
+        quadtree_quant(img, max_err, [x + size, y], size);
+        quadtree_quant(img, max_err, [x, y + size], size);
+        quadtree_quant(img, max_err, [x + size, y + size], size);
+    } else {
+        for y in pos[1]..pos[1] + size {
+            for x in pos[0]..pos[0] + size {
+                img.put_pixel(x, y, Rgb(avg));
+            }
+        }
+    }
+}
+
+fn redmean_sq(x: [u8; 3], y: [u8; 3]) -> f32 {
+    let mut ds: [f32; 3] = [0.0, 0.0, 0.0];
+    for i in 0..2 {
+        ds[i] = x[i] as f32 - y[i] as f32;
+        ds[i] *= ds[i];
+    }
+    let rm = (x[0] as f32 + y[0] as f32) / 2.0;
+    (2.0 + rm / 256.0) * ds[0] + 4.0 * ds[1] + (2.0 + (255.0 - rm) / 256.0) * ds[2]
+}