specify colour depth from commandline

src/main.rs

@@ -3,14 +3,16 @@ use std::{env, fs::File, io::Read};
 use image::{imageops::FilterType, DynamicImage, GenericImage, GenericImageView, Pixel, Rgb};
 
 fn main() {
-    let mut args: Vec<String> = env::args().skip(1).collect();
+    let args: Vec<String> = env::args().skip(1).collect();
     if args.is_empty() {
         println!("no input file specified");
         return;
     }
-    dbg!(&args);
 
-    let inpath = args.pop().unwrap();
+    let inpath = &args[0];
+
+    let color_depth = args.get(1).and_then(|s| s.parse().ok()).unwrap_or(6);
+
     let mut data = Vec::new();
     File::open(inpath).unwrap().read_to_end(&mut data).unwrap();
     let image = image::load_from_memory(&data).unwrap();
@@ -21,12 +23,10 @@ fn main() {
     image.save("scaled.png").unwrap();
     println!("resized");
 
-    const COLOR_COUNT: usize = 6;
+    let quantized_image = quantize_image(image.clone(), color_depth);
-
-    let quantized_image = quantize_image(image.clone(), COLOR_COUNT);
     quantized_image.save("q.png").unwrap();
 
-    let image = dither(image, COLOR_COUNT);
+    let image = dither(image, color_depth);
     image.save("out.png").unwrap();
     println!("saved");
 }
@@ -36,7 +36,7 @@ fn generate_palette(image: &DynamicImage, count: usize) -> Vec<Rgb<u8>> {
 
     let mut buckets: Vec<Vec<Rgb<u8>>> = vec![image
         .pixels()
-        .map(|(_x, _y, a)| a.to_rgb()) //
+        .map(|(_x, _y, p)| p.to_rgb()) //
         .collect()];
 
     struct Spread {
@@ -115,20 +115,51 @@ fn generate_palette(image: &DynamicImage, count: usize) -> Vec<Rgb<u8>> {
         buckets.push(new_bucket);
     }
     let mut colors = Vec::new();
-    for bucket in buckets {
+    for bucket in &buckets {
-        let mut avg = [0u128; 3];
+        let mut avg = [0u64; 3];
-        for p in &bucket {
+        for p in bucket {
             for channel in 0..3 {
-                avg[channel] += p.0[channel] as u128;
+                avg[channel] += p.0[channel] as u64;
             }
         }
-        let num = bucket.len() as u128;
+        let num = bucket.len() as u64;
         colors.push(Rgb([
             (avg[0] / num) as u8,
             (avg[1] / num) as u8,
             (avg[2] / num) as u8,
         ]))
     }
+    // for b in &mut buckets {
+    //     b.clear();
+    // }
+    // for (_x, _y, p) in image.pixels() {
+    //     let mut smallest_dist = 500.;
+    //     let mut smallest_index = 0;
+    //     for (i, &c) in colors.iter().enumerate() {
+    //         let dist = color_dist(c, p.to_rgb());
+    //         if dist < smallest_dist {
+    //             smallest_dist = dist;
+    //             smallest_index = i;
+    //         }
+    //     }
+    //     buckets[smallest_index].push(p.to_rgb());
+    // }
+    // let mut colors = Vec::new();
+    // for bucket in &buckets {
+    //     let mut avg = [0u64; 3];
+    //     for p in bucket {
+    //         for channel in 0..3 {
+    //             avg[channel] += p.0[channel] as u64;
+    //         }
+    //     }
+    //     let num = bucket.len() as u64;
+    //     colors.push(Rgb([
+    //         (avg[0] / num) as u8,
+    //         (avg[1] / num) as u8,
+    //         (avg[2] / num) as u8,
+    //     ]))
+    // }
+    // colors.push(Rgb([0, 0, 0]));
     colors
 }
 
@@ -156,16 +187,34 @@ fn quantize_image(input: DynamicImage, count: usize) -> DynamicImage {
         }
         out.put_pixel(x, y, colors[closest_index].to_rgba());
     }
-
     out
 }
 
 const BAYER_4X4: [[u8; 4]; 4] = [[0, 8, 2, 10], [12, 4, 14, 6], [3, 11, 1, 9], [15, 7, 13, 5]];
 
-fn bayer_f(x: usize, y: usize) -> f32 {
+fn bayer_4(x: u32, y: u32) -> f32 {
+    let x = x as usize % 4;
+    let y = y as usize % 4;
     (BAYER_4X4[y][x] as f32 / 16.) - (15. / 16.) / 2.
 }
 
+// const BAYER_8X8: [[u8; 8]; 8] = [
+//     [0, 32, 8, 40, 2, 34, 10, 42],
+//     [48, 16, 56, 24, 50, 18, 58, 26],
+//     [12, 44, 4, 36, 14, 46, 6, 38],
+//     [60, 28, 52, 20, 62, 30, 54, 22],
+//     [3, 35, 11, 43, 1, 33, 9, 41],
+//     [51, 19, 59, 27, 49, 17, 57, 25],
+//     [15, 47, 7, 39, 13, 45, 5, 37],
+//     [63, 31, 55, 23, 61, 29, 53, 21],
+// ];
+
+// fn bayer_8(x: u32, y: u32) -> f32 {
+//     let x = x as usize % 8;
+//     let y = y as usize % 8;
+//     (BAYER_8X8[y][x] as f32 / 64.) - (63. / 64.) / 2.
+// }
+
 fn dither(input: DynamicImage, count: usize) -> DynamicImage {
     let mut colors = generate_palette(&input, count);
     // let mut colors = vec![
@@ -191,9 +240,7 @@ fn dither(input: DynamicImage, count: usize) -> DynamicImage {
 
         let ratio = second_dist / (second_dist + best_dist) * 2. - 1.;
 
-        let bx = x as usize % 4;
+        let bayer = bayer_4(x, y);
-        let by = y as usize % 4;
-        let bayer = bayer_f(bx, by);
 
         let ratio = ratio + bayer;
         let out_col = if ratio > 0. { best_col } else { second_col };