From 28f4ff1eee1f3e86a5569706f90f74fb1bb7aad8 Mon Sep 17 00:00:00 2001 From: mxhagen Date: Mon, 15 Sep 2025 22:43:25 +0200 Subject: [PATCH] use generic color wrapper --- Cargo.lock | 2 +- Cargo.toml | 2 +- src/cli.rs | 4 +- src/colors.rs | 133 +++++++++++++++++++++++++------------------------- src/main.rs | 87 +++++++++++++++++++-------------- 5 files changed, 121 insertions(+), 107 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 801477e..d320050 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -765,7 +765,7 @@ dependencies = [ [[package]] name = "qtizer" -version = "0.6.0" +version = "0.6.1" dependencies = [ "clap", "image", diff --git a/Cargo.toml b/Cargo.toml index 1a7f144..2a3cb7a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "qtizer" -version = "0.6.0" +version = "0.6.1" edition = "2024" description = "Quantization/palette-generation tool using k-means clustering on pixel data" readme = "README.md" diff --git a/src/cli.rs b/src/cli.rs index f2dc7eb..799a833 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -1,7 +1,7 @@ use clap::*; use image::*; -use crate::colors::ColorFormat; +use crate::colors::ColorCodeFormat; /// Quantization/palette-generation tool using k-means clustering on pixel dataI #[derive(Parser, Debug)] @@ -45,7 +45,7 @@ pub struct Args { /// Palette output format #[arg(short = 'f', long = "format", value_name = "fmt")] - pub format: Option, + pub format: Option, // TODO: add input alpha policy for opaque output images // /// Transparency policy when input has alpha but output does not diff --git a/src/colors.rs b/src/colors.rs index a21e79b..dcb39a3 100644 --- a/src/colors.rs +++ b/src/colors.rs @@ -1,94 +1,91 @@ -use image::Rgba; +use image::*; use crate::kmeans::Kmeansable; -// TODO: implement k-means for `Rgb` pixel format -// for slightly less mem-usage and faster calculation +/// marker trait for usable color types +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Color { + pub color_type: ColorType, + pub data: Vec, +} -impl Kmeansable for Rgba { - type Sum = Rgba; +impl Kmeansable for Color { + type Sum = Vec; - /// black, transparent fn zero() -> Self::Sum { - Rgba([0, 0, 0, 0]) + vec![0; 4] } - /// euclidean distance between two rgba colors fn distance(&self, other: &Self) -> f64 { - let dr = (self[0] as f64) - (other[0] as f64); - let dg = (self[1] as f64) - (other[1] as f64); - let db = (self[2] as f64) - (other[2] as f64); - let da = (self[3] as f64) - (other[3] as f64); - (dr * dr + dg * dg + db * db + da * da).sqrt() + self.data + .iter() + .zip(other.data.iter()) + .map(|(a, b)| ((*a as f64) - (*b as f64)).powi(2)) + .sum::() + .sqrt() } - /// add two rgba colors, returning a u32 sum to avoid overflow fn add(sum: &Self::Sum, other: &Self) -> Self::Sum { - Rgba([ - sum[0] + other[0] as u32, - sum[1] + other[1] as u32, - sum[2] + other[2] as u32, - sum[3] + other[3] as u32, - ]) + sum.iter() + .zip(&other.data) + .map(|(a, b)| a + *b as u32) + .collect() } - /// calculate the mean of a sum of rgba colors and a count fn div(sum: &Self::Sum, count: usize) -> Self { - if count == 0 { - // TODO: better error handling via static logger - panic!("tried to calculate mean of 0 colors (division by zero)"); - } + let data = sum + .iter() + .map(|v| (v / count as u32) as u8) + .collect::>(); - Rgba([ - (sum[0] / count as u32) as u8, - (sum[1] / count as u32) as u8, - (sum[2] / count as u32) as u8, - (sum[3] / count as u32) as u8, - ]) + Color { + color_type: match data.len() { + 3 => ColorType::Rgb8, + 4 => ColorType::Rgba8, + _ => unreachable!( + "invalid color length. only rgb or rgba colors should ever be used here." + ), + }, + data, + } } } /// calculate the rgba brightness (luminance) -pub fn brightness(&Rgba([r, g, b, _]): &Rgba) -> u32 { +pub fn brightness(color: &Color) -> u32 { + let &[r, g, b, ..] = &color.data[..] else { + unreachable!("invalid color type. only rgb or rgba colors should ever be used here."); + }; ((0.299 * r as f32) + (0.587 * g as f32) + (0.114 * b as f32)) as u32 } /// color code output format #[derive(Clone, Copy, Debug, Default, clap::ValueEnum)] -pub enum ColorFormat { +pub enum ColorCodeFormat { /// `#rrggbb` or `#rrggbbaa` #[default] Hex, /// `rgb(r, g, b)` or `rgba(r, g, b, a)` Rgb, } -use ColorFormat::*; -impl ColorFormat { +impl ColorCodeFormat { /// pretty print a color code in the format /// when writing to terminals, uses ansi escape codes for color preview - pub fn pretty_print_color_code( - format: &ColorFormat, - writer: &mut W, - color: &image::Rgba, - with_alpha: bool, - ) where + pub fn pretty_print_color_code(format: &ColorCodeFormat, writer: &mut W, color: &Color) + where W: std::io::Write, { match format { - Hex => Self::colored_with_format(writer, color, with_alpha, Self::hex_color_code), - Rgb => Self::colored_with_format(writer, color, with_alpha, Self::rgb_color_code), + ColorCodeFormat::Hex => Self::colored_with_format(writer, color, Self::hex_color_code), + ColorCodeFormat::Rgb => Self::colored_with_format(writer, color, Self::rgb_color_code), } } /// pretty print wrapper that colors output /// given a callback providing the actual color formatting - fn colored_with_format( - writer: &mut W, - color: &image::Rgba, - with_alpha: bool, - callback: fn(&mut W, &image::Rgba, bool), - ) where + fn colored_with_format(writer: &mut W, color: &Color, callback: fn(&mut W, &Color)) + where W: std::io::Write, { use std::io::IsTerminal; @@ -96,51 +93,55 @@ impl ColorFormat { if !colorize { // just print formatted color, no ansi codes - return callback(writer, color, with_alpha); + return callback(writer, color); } // ensure text has enough contrast to colored background - match brightness(&color) { + match brightness(color) { ..128 => write!(writer, "\x1b[38;2;255;255;255m"), // dark => white text _ => write!(writer, "\x1b[38;2;0;0;0m"), // light => black text } .expect("failed to write output"); + let &[r, g, b, ..] = &color.data[..] else { + unreachable!("invalid color type. only rgb or rgba colors should ever be used here."); + }; + // print ansi codes for colored background - write!(writer, "\x1b[48;2;{};{};{}m", color[0], color[1], color[2],) - .expect("failed to write output"); + write!(writer, "\x1b[48;2;{r};{g};{b}m").expect("failed to write output"); // call the actual color printing function - callback(writer, color, with_alpha); + callback(writer, color); // reset colors write!(writer, "\x1b[0m").expect("failed to write output"); } /// print uncolored hex color code, with optional alpha - fn hex_color_code(writer: &mut W, color: &image::Rgba, with_alpha: bool) + fn hex_color_code(writer: &mut W, color: &Color) where W: std::io::Write, { - let Rgba([r, g, b, a]) = color; + let c = &color.data; + write!(writer, "#{:02x}{:02x}{:02x}", c[0], c[1], c[2]).expect("failed to write output"); - write!(writer, "#{:02x}{:02x}{:02x}", r, g, b).expect("failed to write output"); - - if with_alpha { - write!(writer, "{:02x}", a).expect("failed to write output"); + if color.color_type == ColorType::Rgba8 { + write!(writer, "{:02x}", c[3]).expect("failed to write output"); } } /// print uncolored rgb color code, with optional alpha - fn rgb_color_code(writer: &mut W, color: &image::Rgba, with_alpha: bool) + fn rgb_color_code(writer: &mut W, color: &Color) where W: std::io::Write, { - let Rgba([r, g, b, a]) = color; - - match with_alpha { - true => write!(writer, "rgba({}, {}, {}, {})", r, g, b, a), - false => write!(writer, "rgb({}, {}, {})", r, g, b), + let c = &color.data; + match color.color_type { + ColorType::Rgba8 => write!(writer, "rgba({}, {}, {}, {})", c[0], c[1], c[2], c[3]), + ColorType::Rgb8 => write!(writer, "rgb({}, {}, {})", c[0], c[1], c[2]), + _ => unreachable!( + "invalid color type. only rgb or rgba colors should ever be used here." + ), } .expect("failed to write output") } diff --git a/src/main.rs b/src/main.rs index d5569cc..5e4ce1c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,7 +6,7 @@ mod cli; mod colors; mod kmeans; -use crate::colors::ColorFormat; +use crate::colors::*; fn main() { let args = cli::Args::parse(); @@ -24,17 +24,29 @@ fn main() { let mut context = kmeans::Context::new(seed); - // TODO: use rgb only, when use-alpha is false - // with this refactor, also look at storing results & necessary config in - // `Context` in order to pass it uniformly to the various handlers - // open file and parse image - let img = image::open(args.file_path) - .expect("failed to open image") - .to_rgba8(); + let img = image::open(args.file_path).expect("failed to open image"); + + let pixels = match args.alpha { + true => img + .to_rgba8() + .pixels() + .map(|p| Color { + data: p.0.to_vec(), + color_type: ColorType::Rgba8, + }) + .collect::>(), + false => img + .to_rgb8() + .pixels() + .map(|p| Color { + data: p.0.to_vec(), + color_type: ColorType::Rgb8, + }) + .collect::>(), + }; // run kmeans - let pixels = img.pixels().cloned().collect::>(); let (clusters, assignments) = context.k_means(&pixels, args.number, args.iterations); // handle output @@ -42,85 +54,86 @@ fn main() { None => palette_handler( &clusters, &mut std::io::stdout(), - args.alpha, &args.format.unwrap_or_default(), ), Some(output_file) if ImageFormat::from_path(&output_file).is_ok() => { - image_file_handler(&img, &clusters, &assignments, args.alpha, output_file) + let (width, height) = img.dimensions(); + image_file_handler(width, height, &clusters, &assignments, output_file); } Some(output_file) => { let mut file = std::fs::File::create(output_file).expect("failed to create output file"); - palette_handler( - &clusters, - &mut file, - args.alpha, - &args.format.unwrap_or_default(), - ); + palette_handler(&clusters, &mut file, &args.format.unwrap_or_default()); } } } /// handle palette output to terminal or file -fn palette_handler( - clusters: &[image::Rgba], - writer: &mut W, - alpha: bool, - format: &ColorFormat, -) where +fn palette_handler(clusters: &[Color], writer: &mut W, format: &ColorCodeFormat) +where W: std::io::Write, { // sort colors by alpha and brightness let mut clusters = clusters.to_vec(); - clusters.sort_by(|a, b| { - a[3].cmp(&b[3]) // alpha first - .then_with(|| u32::cmp(&colors::brightness(b), &colors::brightness(a))) // then brightness + clusters.sort_by(|x, y| { + let (&[r_x, g_x, b_x], &[r_y, g_y, b_y]) = (&x.data[..3], &y.data[..3]) else { + unreachable!("invalid color type. only rgb or rgba colors should ever be used here."); + }; + u32::cmp(&colors::brightness(y), &colors::brightness(x)) // descending brightness + .then_with(|| { + u32::from_be_bytes([r_x, g_x, b_x, 0]).cmp(&u32::from_be_bytes([r_y, g_y, b_y, 0])) + }) }); // output palette as hex #rrggbbaa // output with ansi escape codes for color preview in terminal for color in &clusters { - ColorFormat::pretty_print_color_code(format, writer, color, alpha); + ColorCodeFormat::pretty_print_color_code(format, writer, color); writeln!(writer).expect("failed to write color to output"); } } /// handle image output to file fn image_file_handler( - img: &ImageBuffer, Vec>, - clusters: &[image::Rgba], + width: u32, + height: u32, + clusters: &[Color], assignments: &[usize], - alpha: bool, output_file: String, ) { - let (width, height) = img.dimensions(); - // create new image by replacing each pixel with its cluster center let quantized = assignments .iter() - .flat_map(|&i| &clusters[i].0[..if alpha { 4 } else { 3 }]) + .flat_map(|&i| &clusters[i].data) .copied() .collect::>(); - let status = match alpha { - true => { + let status = match clusters.first() { + Some(c) if c.color_type == ColorType::Rgba8 => { let img = ImageBuffer::from_vec(width, height, quantized); let img: ImageBuffer, _> = img.expect("failed to create quantized image"); img.save(&output_file) } - false => { + Some(c) if c.color_type == ColorType::Rgb8 => { let img = ImageBuffer::from_vec(width, height, quantized); let img: ImageBuffer, _> = img.expect("failed to create quantized image"); img.save(&output_file) } + _ => { + cli::err_exit( + clap::error::ErrorKind::InvalidValue, + "unexpected error - this is a bug".to_string(), + ); + unreachable!("should have exited"); + } }; // TODO: better errors handling logger // save image with inferred format match status { - Ok(_) => println!("saved quantized image to {}", output_file), + Ok(_) => println!("saved quantized image to {output_file}"), Err(err) => { // errors here are unexpected, since extension alpha-capability // is validated in `cli::semantically_validate`