cmounce/color-cell
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Move encoding logic into EncodeTree, encode.rs

Browse Source
This commit is contained in:
Chris Mounce
2024-08-14 01:35:25 -07:00
parent 9cd8cf3989
commit e8258a570d
2 changed files with 167 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/src/encode.rs
+148
View File
@@ -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
+19 -42
View File
@@ -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;