mxhagen/qtizer
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

use generic color wrapper

Browse Source
This commit is contained in:
mxhagen 2025-09-15 22:43:25 +02:00
parent 0fb6e337c2
commit 28f4ff1eee
5 changed files with 121 additions and 107 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
Cargo.lock generated
View File

@ -765,7 +765,7 @@ dependencies = [
[[package]]
name = "qtizer"
version = "0.6.0"
version = "0.6.1"
dependencies = [
"clap",
"image",

2
Cargo.toml
View File

@ -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"

4
src/cli.rs
View File

@ -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<ColorFormat>,
pub format: Option<ColorCodeFormat>,
// TODO: add input alpha policy for opaque output images
// /// Transparency policy when input has alpha but output does not

133
src/colors.rs
View File

@ -1,94 +1,91 @@
use image::Rgba;
use image::*;
use crate::kmeans::Kmeansable;
// TODO: implement k-means for `Rgb<u8>` 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<u8>,
}
impl Kmeansable for Rgba<u8> {
type Sum = Rgba<u32>;
impl Kmeansable for Color {
type Sum = Vec<u32>;
/// 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::<f64>()
.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::<Vec<u8>>();
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<u8>) -> 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<W>(
format: &ColorFormat,
writer: &mut W,
color: &image::Rgba<u8>,
with_alpha: bool,
) where
pub fn pretty_print_color_code<W>(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<W>(
writer: &mut W,
color: &image::Rgba<u8>,
with_alpha: bool,
callback: fn(&mut W, &image::Rgba<u8>, bool),
) where
fn colored_with_format<W>(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<W>(writer: &mut W, color: &image::Rgba<u8>, with_alpha: bool)
fn hex_color_code<W>(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<W>(writer: &mut W, color: &image::Rgba<u8>, with_alpha: bool)
fn rgb_color_code<W>(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")
}

87
src/main.rs
View File

@ -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::<Vec<_>>(),
false => img
.to_rgb8()
.pixels()
.map(|p| Color {
data: p.0.to_vec(),
color_type: ColorType::Rgb8,
})
.collect::<Vec<_>>(),
};
// run kmeans
let pixels = img.pixels().cloned().collect::<Vec<_>>();
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<W>(
clusters: &[image::Rgba<u8>],
writer: &mut W,
alpha: bool,
format: &ColorFormat,
) where
fn palette_handler<W>(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<Rgba<u8>, Vec<u8>>,
clusters: &[image::Rgba<u8>],
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::<Vec<_>>();
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<Rgba<u8>, _> = 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<Rgb<u8>, _> = 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`