version 0.7: significant performance update, better progress

Cargo.lock

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

Cargo.toml

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

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

src/colors.rs

@@ -1,94 +1,92 @@
-use image::Rgba;
+use image::*;
 
 use crate::kmeans::Kmeansable;
 
-// 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;
@@ -96,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

@@ -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)
+    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);
 
     // 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(
+            palette_handler(&clusters, &mut file, &args.format.unwrap_or_default());
-                &clusters,
-                &mut file,
-                args.alpha,
-                &args.format.unwrap_or_default(),
-            );
         }
     }
 }
 
 /// handle palette output to terminal or file
-fn palette_handler<W>(
+fn palette_handler<W>(clusters: &[Color], writer: &mut W, format: &ColorCodeFormat)
-    clusters: &[image::Rgba<u8>],
+where
-    writer: &mut W,
-    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(
-    img: &ImageBuffer<Rgba<u8>, Vec<u8>>,
+    width: u32,
-    clusters: &[image::Rgba<u8>],
+    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 {
+    let status = match clusters.first() {
-        true => {
+        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`