From e8258a570d21540723431c470e575e737eb91958 Mon Sep 17 00:00:00 2001 From: Chris Mounce Date: Wed, 14 Aug 2024 01:35:25 -0700 Subject: [PATCH] Move encoding logic into EncodeTree, encode.rs --- rust/src/encode.rs | 148 +++++++++++++++++++++++++++++++++++++++++++++ rust/src/main.rs | 61 ++++++------------- 2 files changed, 167 insertions(+), 42 deletions(-) create mode 100644 rust/src/encode.rs diff --git a/rust/src/encode.rs b/rust/src/encode.rs new file mode 100644 index 0000000..c42658a --- /dev/null +++ b/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, + } +} diff --git a/rust/src/main.rs b/rust/src/main.rs index 16561cb..c89195f 100644 --- a/rust/src/main.rs +++ b/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) { - // 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 quadtree_quant(img: &mut RgbImage, max_err: f32) { + let region = Region { + x: 0, + y: 0, + width: img.width(), + height: img.height(), + }; + let tree = choose_encoding(img, region, max_err); + println!( + "error = {}, predicted size = {} bytes", + tree.error, tree.bytes + ); + tree.apply(img, region); } -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;