Files
color-cell/rust/src/main.rs
T

93 lines
2.7 KiB
Rust

use anyhow::{bail, Result};
use clap::Parser;
use image::{DynamicImage, Rgb, RgbImage};
#[derive(Debug, Parser)]
#[command()]
struct Args {
#[arg(short, long, value_name = "FILE")]
input: String,
#[arg(short, long, value_name = "FILE")]
output: String,
#[arg(short = 'e', long, value_name = "NUM")]
max_err: Option<f32>,
}
fn main() -> Result<()> {
let args = Args::parse();
let img = image::open(&args.input)?;
let img = match img {
DynamicImage::ImageRgb8(img) => img,
_ => bail!("Unsupported type: {:?}", img),
};
println!("Opened {} ({}x{})", &args.input, img.width(), img.height());
let mut output = RgbImage::new(128, 128);
for x in 0..=127 {
for y in 0..=127 {
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]
}