Move encoding logic into EncodeTree, encode.rs

rust/src/encode.rs

@@ -0,0 +1,148 @@
+use image::{Rgb, RgbImage};
+
+use crate::redmean_sq;
+
+pub type Color = [u8; 3]; // TODO: Convert to indexed
+
+#[derive(Debug)]
+pub struct EncodeTree {
+    pub error: f32,
+    pub bytes: f32,
+    pub strategy: EncodeStrategy,
+}
+
+#[derive(Debug)]
+pub enum EncodeStrategy {
+    Solid(Color),
+    // TODO: add Bitmap([Color; 2]),
+    QuadSplit(Box<[EncodeTree; 4]>),
+}
+
+impl EncodeTree {
+    pub fn apply(&self, img: &mut RgbImage, region: Region) {
+        match &self.strategy {
+            EncodeStrategy::Solid(color) => {
+                for y in region.y..region.y2() {
+                    for x in region.x..region.x2() {
+                        img.put_pixel(x, y, Rgb(*color))
+                    }
+                }
+            }
+            EncodeStrategy::QuadSplit(subtrees) => {
+                let width = region.width / 2;
+                let height = region.height / 2;
+                let coords = [
+                    [region.x, region.y],
+                    [region.x + width, region.y],
+                    [region.x, region.y + height],
+                    [region.x + width, region.y + height],
+                ];
+                for (tree, &[x, y]) in subtrees.iter().zip(coords.iter()) {
+                    tree.apply(
+                        img,
+                        Region {
+                            x,
+                            y,
+                            width,
+                            height,
+                        },
+                    );
+                }
+            }
+        }
+    }
+}
+
+#[derive(Clone, Copy)]
+pub struct Region {
+    pub x: u32,
+    pub y: u32,
+    pub width: u32,
+    pub height: u32,
+}
+
+impl Region {
+    pub fn x2(&self) -> u32 {
+        self.x + self.width
+    }
+
+    pub fn y2(&self) -> u32 {
+        self.y + self.height
+    }
+
+    pub fn area(&self) -> u32 {
+        self.width * self.height
+    }
+}
+
+pub fn choose_encoding(img: &RgbImage, region: Region, max_error: f32) -> EncodeTree {
+    // Attempt to quantize as solid color
+    let info = analyze_solid(img, region);
+    if info.error <= max_error {
+        return EncodeTree {
+            error: info.error,
+            bytes: 1.0,
+            strategy: EncodeStrategy::Solid(info.color),
+        };
+    }
+
+    // Error is too high, so we have to subdivide
+    let mut subinfo = vec![];
+    let width = region.width / 2;
+    let height = region.height / 2;
+    for y in [region.y, region.y + height] {
+        for x in [region.x, region.x + width] {
+            subinfo.push(choose_encoding(
+                img,
+                Region {
+                    x,
+                    y,
+                    width,
+                    height,
+                },
+                max_error / 4.0,
+            ));
+        }
+    }
+    EncodeTree {
+        error: subinfo.iter().fold(0.0, |acc, info| acc + info.error),
+        bytes: 1.0 + subinfo.iter().fold(0.0, |acc, info| acc + info.bytes),
+        strategy: EncodeStrategy::QuadSplit(Box::new(subinfo.try_into().unwrap())),
+    }
+}
+
+struct SolidInfo {
+    color: Color,
+    error: f32,
+}
+
+fn analyze_solid(img: &RgbImage, region: Region) -> SolidInfo {
+    // Calculate average color
+    let mut avg = [0.0, 0.0, 0.0];
+    for y in region.y..region.y2() {
+        for x in region.x..region.x2() {
+            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 /= region.area() 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 region.y..region.y2() {
+        for x in region.x..region.x2() {
+            let pixel = img.get_pixel(x, y);
+            err += redmean_sq(avg, pixel.0);
+        }
+    }
+
+    SolidInfo {
+        color: avg,
+        error: err,
+    }
+}

rust/src/main.rs

@@ -1,6 +1,9 @@
+mod encode;
+
 use anyhow::{bail, Result};
 use clap::Parser;
-use image::{DynamicImage, Rgb, RgbImage};
+use encode::{choose_encoding, Region};
+use image::{DynamicImage, RgbImage};
 
 #[derive(Debug, Parser)]
 #[command()]
@@ -32,54 +35,28 @@ fn main() -> Result<()> {
     let error_scale_factor = 0.2; // Adjusted by feel. Q=0 is unrecognizable but not blank.
     let error_per_pixel = (1.0 - args.max_err.unwrap_or(80.0) / 100.0) * error_scale_factor;
     let sq_error_per_pixel = error_per_pixel * error_per_pixel;
-    quadtree_quant(&mut output, sq_error_per_pixel * 128.0 * 128.0, [0, 0], 128);
+    quadtree_quant(&mut output, sq_error_per_pixel * 128.0 * 128.0);
     output.save(&args.output)?;
     println!("Wrote {}", &args.output);
     Ok(())
 }
 
-fn quadtree_quant(img: &mut RgbImage, max_err: f32, pos: [u32; 2], size: u32) {
+fn quadtree_quant(img: &mut RgbImage, max_err: f32) {
-    // Calculate average color
+    let region = Region {
-    let mut avg = [0.0, 0.0, 0.0];
+        x: 0,
-    for y in pos[1]..pos[1] + size {
+        y: 0,
-        for x in pos[0]..pos[0] + size {
+        width: img.width(),
-            let pixel = img.get_pixel(x, y);
+        height: img.height(),
-            for (i, component) in avg.iter_mut().enumerate() {
+    };
-                *component += (pixel.0[i] as f32).powf(2.2);
+    let tree = choose_encoding(img, region, max_err);
-            }
+    println!(
-        }
+        "error = {}, predicted size = {} bytes",
-    }
+        tree.error, tree.bytes
-    for component in &mut avg {
+    );
-        *component /= (size * size) as f32;
+    tree.apply(img, region);
-    }
-    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 {
+pub fn redmean_sq(x: [u8; 3], y: [u8; 3]) -> f32 {
     let mut ds: [f32; 3] = [0.0, 0.0, 0.0];
     for (i, d) in ds.iter_mut().enumerate() {
         let delta = x[i] as f32 - y[i] as f32;