version 0.7: significant performance update, better progress

.gitignore

@@ -1,4 +1 @@
 /target
-
-# test images
-*.png

Cargo.lock

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

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "qtizer"
-version = "0.5.0"
+version = "0.7.0"
 edition = "2024"
 description = "Quantization/palette-generation tool using k-means clustering on pixel data"
 readme = "README.md"

README.md

@@ -5,6 +5,8 @@ command line quantization/palette-generation tool using k-means clustering on pi
 
 - [features](#features)
 - [usage](#usage)
+  - [example: palette creation in rgb format with alpha](#example-palette-creation-in-rgb-format-with-alpha-output-is-colored-accordingly-in-terminals)
+  - [example: image quantization to reduced palette](#example-image-quantization-to-reduced-palette-file-formats-inferred-based-on-extension)
 - [installation](#installation)
 
 
@@ -31,12 +33,13 @@ Options:
   -s, --seed <number>    Optional RNG seed for reproducible results
   -o, --output <output>  Output file path
                          - If not provided, outputs to stdout
-  -f, --format <fmt>     Palette output format [default: hex] [possible values: hex, rgb]
+                         - With image file extensions, outputs an image file
+  -f, --format <fmt>     Palette output format [possible values: hex, rgb]
   -h, --help             Print help (see more with '--help')
   -V, --version          Print version
 ```
 
-Example (output is colored accordingly in terminals):
+#### Example: palette creation in rgb format with alpha (output is colored accordingly in terminals):
 ```sh
 $ qtizer wallpaper.png -k 3 -af rgb
 
@@ -45,6 +48,11 @@ rgba(191, 150, 132, 254)
 rgba(48, 45, 51, 254)
 ```
 
+#### Example: image quantization to reduced palette (file formats inferred based on extension):
+```sh
+$ qtizer wallpaper.png -k 8 quantized.png
+```
+
 
 ## installation
 

src/cli.rs

@@ -1,5 +1,7 @@
-use crate::colors::ColorFormat;
 use clap::*;
+use image::*;
+
+use crate::colors::ColorCodeFormat;
 
 /// Quantization/palette-generation tool using k-means clustering on pixel dataI
 #[derive(Parser, Debug)]
@@ -27,8 +29,7 @@ pub struct Args {
 
     /// Output file path
     /// - If not provided, outputs to stdout
-    // TODO: implement image output
+    /// - With image file extensions, outputs an image file
-    // /// - With image file extensions, outputs an image file
     #[arg(
         short = 'o',
         long = "output",
@@ -43,13 +44,14 @@ pub struct Args {
     pub output_positional: Option<String>,
 
     /// Palette output format
-    #[arg(
+    #[arg(short = 'f', long = "format", value_name = "fmt")]
-        short = 'f',
+    pub format: Option<ColorCodeFormat>,
-        long = "format",
+
-        default_value = "hex",
+    // TODO: add input alpha policy for opaque output images
-        value_name = "fmt"
+    // /// Transparency policy when input has alpha but output does not
-    )]
+    // #[arg(short = 'p', long = "alpha-policy", value_name = "policy",)]
-    pub format: Option<ColorFormat>,
+    // pub alpha_policy: AlphaPolicy,
+
     // TODO: implement multi-threading
     // /// Number of workers to use [default: core count]
     // #[arg(short = 'j', long = "jobs",
@@ -59,35 +61,43 @@ pub struct Args {
 }
 
 /// semantic validation of arguments
+/// - `--format` cannot be specified when outputting an image file
+/// - some image formats do not support alpha (eg. jpg)
 pub fn semantically_validate(args: &Args) {
+    // check if `--format` is specified AND output has image file extension
     if args.format.is_some()
         && (args.output.clone())
             .or(args.output_positional.clone())
-            .is_some_and(is_image_file_path)
+            .is_some_and(|p| ImageFormat::from_path(p).is_ok())
     {
-        Args::command()
+        err_exit(
-            .error(
+            clap::error::ErrorKind::ArgumentConflict,
+            "cannot specify color-code format when outputting an image file.",
+        );
+    }
+
+    // check if output image format supports alpha channel
+    let output_opt = args.output.clone().or(args.output_positional.clone());
+    if args.alpha && output_opt.is_some() {
+        let output_file = output_opt.unwrap();
+        let filetype = ImageFormat::from_path(&output_file);
+
+        use ImageFormat::*;
+        if matches!(filetype, Ok(Jpeg | Bmp | Pnm | Tiff)) {
+            err_exit(
                 clap::error::ErrorKind::ArgumentConflict,
-                "cannot specify color format when outputting an image file.",
+                format!(
-            )
+                    "the `{:?}` image format does not support alpha.",
-            .exit();
+                    filetype.unwrap(),
+                ),
+            );
+        }
     }
 }
 
-// TODO: implement image output
+/// shorthand for `Args::command().error(...).exit()`
-pub fn is_image_file_path<P>(file: P) -> bool
+pub fn err_exit(kind: clap::error::ErrorKind, message: impl std::fmt::Display) {
-where
+    Args::command().error(kind, message).exit()
-    P: AsRef<std::path::Path>,
-{
-    file.as_ref()
-        .extension()
-        .and_then(|e| e.to_str())
-        .is_some_and(|e| {
-            matches!(
-                e.to_lowercase().as_str(),
-                "png" | "jpg" | "jpeg" | "bmp" | "tiff" | "ppm"
-            )
-        })
 }
 
 // TODO: implement static logger functionality for parsed arguments

src/colors.rs

@@ -1,93 +1,92 @@
+use image::*;
+
 use crate::kmeans::Kmeansable;
-use image::Rgba;
 
-// TODO: implement k-means for `Rgb<u8>` pixel format
+/// marker trait for usable color types
-//       for slightly less mem-usage and faster calculation
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct Color {
+    pub color_type: ColorType,
+    pub data: Vec<u8>,
+}
 
-impl Kmeansable for Rgba<u8> {
+impl Kmeansable for Color {
-    type Sum = Rgba<u32>;
+    type Sum = Vec<u32>;
 
-    /// black, transparent
     fn zero() -> Self::Sum {
-        Rgba([0, 0, 0, 0])
+        vec![0; 4]
     }
 
-    /// euclidean distance between two rgba colors
+    /// distance function, according to which clustering is performed
+    /// (impl avoids sqrt for performance -- uses squared distance)
     fn distance(&self, other: &Self) -> f64 {
-        let dr = (self[0] as f64) - (other[0] as f64);
+        self.data
-        let dg = (self[1] as f64) - (other[1] as f64);
+            .iter()
-        let db = (self[2] as f64) - (other[2] as f64);
+            .zip(other.data.iter())
-        let da = (self[3] as f64) - (other[3] as f64);
+            .map(|(a, b)| ((*a as f64) - (*b as f64)).powi(2))
-        (dr * dr + dg * dg + db * db + da * da).sqrt()
+            .sum::<f64>()
     }
 
-    /// add two rgba colors, returning a u32 sum to avoid overflow
     fn add(sum: &Self::Sum, other: &Self) -> Self::Sum {
-        Rgba([
+        sum.iter()
-            sum[0] + other[0] as u32,
+            .zip(&other.data)
-            sum[1] + other[1] as u32,
+            .map(|(a, b)| a + *b as u32)
-            sum[2] + other[2] as u32,
+            .collect()
-            sum[3] + other[3] as u32,
-        ])
     }
 
-    /// calculate the mean of a sum of rgba colors and a count
     fn div(sum: &Self::Sum, count: usize) -> Self {
-        if count == 0 {
+        let data = sum
-            // TODO: better error handling via static logger
+            .iter()
-            panic!("tried to calculate mean of 0 colors (division by zero)");
+            .map(|v| (v / count as u32) as u8)
-        }
+            .collect::<Vec<u8>>();
 
-        Rgba([
+        Color {
-            (sum[0] / count as u32) as u8,
+            color_type: match data.len() {
-            (sum[1] / count as u32) as u8,
+                3 => ColorType::Rgb8,
-            (sum[2] / count as u32) as u8,
+                4 => ColorType::Rgba8,
-            (sum[3] / count as u32) as u8,
+                _ => 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>(
+    pub fn pretty_print_color_code<W>(format: &ColorCodeFormat, writer: &mut W, color: &Color)
-        format: &ColorFormat,
+    where
-        writer: &mut W,
-        color: &image::Rgba<u8>,
-        with_alpha: bool,
-    ) where
         W: std::io::Write,
     {
         match format {
-            Hex => Self::colored_with_format(writer, color, with_alpha, Self::hex_color_code),
+            ColorCodeFormat::Hex => Self::colored_with_format(writer, color, Self::hex_color_code),
-            Rgb => Self::colored_with_format(writer, color, with_alpha, Self::rgb_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>(
+    fn colored_with_format<W>(writer: &mut W, color: &Color, callback: fn(&mut W, &Color))
-        writer: &mut W,
+    where
-        color: &image::Rgba<u8>,
-        with_alpha: bool,
-        callback: fn(&mut W, &image::Rgba<u8>, bool),
-    ) where
         W: std::io::Write,
     {
         use std::io::IsTerminal;
@@ -95,51 +94,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],)
+        write!(writer, "\x1b[48;2;{r};{g};{b}m").expect("failed to write output");
-            .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 color.color_type == ColorType::Rgba8 {
-
+            write!(writer, "{:02x}", c[3]).expect("failed to write output");
-        if with_alpha {
-            write!(writer, "{:02x}", a).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;
+        let c = &color.data;
-
+        match color.color_type {
-        match with_alpha {
+            ColorType::Rgba8 => write!(writer, "rgba({}, {}, {}, {})", c[0], c[1], c[2], c[3]),
-            true => write!(writer, "rgba({}, {}, {}, {})", r, g, b, a),
+            ColorType::Rgb8 => write!(writer, "rgb({}, {}, {})", c[0], c[1], c[2]),
-            false => write!(writer, "rgb({}, {}, {})", r, g, b),
+            _ => unreachable!(
+                "invalid color type. only rgb or rgba colors should ever be used here."
+            ),
         }
         .expect("failed to write output")
     }

src/kmeans.rs

@@ -44,22 +44,70 @@ impl Context<SmallRng> {
             .cloned()
             .collect::<Vec<_>>();
 
+        // make cursor invisible
+        eprint!("\x1b[?25l");
+
         for i in 0..iterations {
             // TODO: implement static logger functionality for progress
-            eprintln!("processing k-means iteration {}/{}...", i + 1, iterations);
+            //       once implemented, replace other eprint(ln)! calls too
+            eprintln!(
+                "processing k-means iteration: [ {:>9} / {:>9} ]...",
+                i + 1,
+                iterations
+            );
+
+            // precompute cluster distances to skip some distance calculations later
+            // only set for i < j -- note: dist[i][j] == dist[j][i]
+            let mut cluster_distances = vec![vec![0.0; k]; k];
+            for i in 0..k {
+                for j in (i + 1)..k {
+                    let dist = clusters[i].distance(&clusters[j]);
+                    cluster_distances[i][j] = dist; // i < j case only
+                }
+            }
 
             // assign each point to the nearest cluster
             for (i, point) in data.iter().enumerate() {
-                let min_idx = clusters
+                let mut closest_idx = 0;
-                    .iter()
+                let mut closest_dist = clusters[0].distance(point);
-                    .enumerate()
-                    .min_by(|(_, a), (_, b)| f64::total_cmp(&a.distance(point), &b.distance(point)))
-                    .unwrap()
-                    .0;
 
-                assignments[i] = min_idx;
+                // print progress every 500 points
+                if i % 500 == 0 {
+                    if i > 0 {
+                        // restore cursor position (write over previous status)
+                        eprint!("\x1b[1F");
+                    }
+                    let label_len = "processing k-means iteration".len();
+                    eprintln!(
+                        "{:>label_len$}: [ {:>9} / {:>9} ]...",
+                        "assigning point",
+                        i,
+                        data.len()
+                    );
+                }
+
+                for j in 1..k {
+                    // skip distance calculation if the cluster is too far away
+                    let (a, b) = (closest_idx.min(j), closest_idx.max(j));
+                    if cluster_distances[a][b] >= 2.0 * closest_dist {
+                        // d(c_j, c_min) >= 2 * d(p, c_min)
+                        // d(p,   c_j  ) >=     d(p, c_min)
+                        continue;
+                    }
+
+                    let dist = clusters[j].distance(point);
+                    if dist < closest_dist {
+                        closest_dist = dist;
+                        closest_idx = j;
+                    }
+                }
+
+                assignments[i] = closest_idx;
             }
 
+            // restore cursor position (write over previous status)
+            eprint!("\x1b[2F");
+
             // move cluster to mean of its assigned points
             let mut counts: Vec<usize> = vec![0; k];
             let mut sums = vec![T::zero(); k];
@@ -77,6 +125,9 @@ impl Context<SmallRng> {
             }
         }
 
+        // make cursor visible again
+        eprint!("\x1b[?25h");
+
         (clusters, assignments)
     }
 

src/main.rs

@@ -1,12 +1,13 @@
-use clap::{CommandFactory, Parser};
+use clap::*;
+use image::*;
 use std::time::{SystemTime, UNIX_EPOCH};
 
-use crate::colors::ColorFormat;
-
 mod cli;
 mod colors;
 mod kmeans;
 
+use crate::colors::*;
+
 fn main() {
     let args = cli::Args::parse();
     cli::semantically_validate(&args);
@@ -23,62 +24,125 @@ 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)
+    let img = image::open(args.file_path).expect("failed to open image");
-        .expect("failed to open image")
+
-        .to_rgba8();
+    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);
-    let (clusters, _assignments) = context.k_means(&pixels, args.number, args.iterations);
 
     // handle output
     match args.output.or(args.output_positional) {
         None => palette_handler(
             &clusters,
             &mut std::io::stdout(),
-            args.alpha,
             &args.format.unwrap_or_default(),
         ),
 
-        // TODO: implement image output
+        Some(output_file) if ImageFormat::from_path(&output_file).is_ok() => {
-        // Some(output_file) if is_image_file_path(&output_file) => image_file_handler(&clusters, &assignments, 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(
+            palette_handler(&clusters, &mut file, &args.format.unwrap_or_default());
-                &clusters,
-                &mut file,
-                args.alpha,
-                &args.format.unwrap_or_default(),
-            );
         }
     }
 }
 
-fn palette_handler<W>(
+/// handle palette output to terminal or file
-    clusters: &[image::Rgba<u8>],
+fn palette_handler<W>(clusters: &[Color], writer: &mut W, format: &ColorCodeFormat)
-    writer: &mut W,
+where
-    alpha: bool,
-    format: &ColorFormat,
-) where
     W: std::io::Write,
 {
     // sort colors by alpha and brightness
     let mut clusters = clusters.to_vec();
-    clusters.sort_by(|a, b| {
+    clusters.sort_by(|x, y| {
-        a[3].cmp(&b[3]) // alpha first
+        let (&[r_x, g_x, b_x], &[r_y, g_y, b_y]) = (&x.data[..3], &y.data[..3]) else {
-            .then_with(|| u32::cmp(&colors::brightness(b), &colors::brightness(a))) // then brightness
+            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(
+    width: u32,
+    height: u32,
+    clusters: &[Color],
+    assignments: &[usize],
+    output_file: String,
+) {
+    // create new image by replacing each pixel with its cluster center
+    let quantized = assignments
+        .iter()
+        .flat_map(|&i| &clusters[i].data)
+        .copied()
+        .collect::<Vec<_>>();
+
+    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)
+        }
+        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}"),
+        Err(err) => {
+            // errors here are unexpected, since extension alpha-capability
+            // is validated in `cli::semantically_validate`
+            cli::err_exit(
+                clap::error::ErrorKind::InvalidValue,
+                "unexpectedly failed to save quantized image.\n".to_string()
+                    + "try checking the output file format. (does it support alpha?)\n"
+                    + &format!("    ({err})"),
+            );
+        }
+    }
+}