Introduce indexed color with the PICO-8 palette

rust/Cargo.lock

@@ -242,6 +242,7 @@ dependencies = [
  "anyhow",
  "clap",
  "image",
+ "ordered-float",
 ]
 
 [[package]]
@@ -642,6 +643,15 @@ version = "1.19.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92"
 
+[[package]]
+name = "ordered-float"
+version = "4.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4a91171844676f8c7990ce64959210cd2eaef32c2612c50f9fae9f8aaa6065a6"
+dependencies = [
+ "num-traits",
+]
+
 [[package]]
 name = "paste"
 version = "1.0.15"

rust/Cargo.toml

@@ -7,6 +7,7 @@ edition = "2021"
 anyhow = "1.0.86"
 clap = { version = "4.5.15", features = ["derive"] }
 image = "0.25.2"
+ordered-float = "4.2.2"
 
 [profile.dev.package.zune-jpeg]
 opt-level = 3

rust/src/encode.rs

@@ -1,6 +1,6 @@
 use image::{Rgb, RgbImage};
 
-use crate::redmean_sq;
+use crate::indexed::redmean_sq;
 
 pub type Color = [u8; 3]; // TODO: Convert to indexed
 

rust/src/indexed.rs

@@ -0,0 +1,165 @@
+use image::RgbImage;
+use ordered_float::NotNan;
+
+// pub fn euclid_sq(x: [u8; 3], y: [u8; 3]) -> f32 {
+//     let mut result = 0.0;
+//     for i in 0..3 {
+//         let delta = x[i] as f32 - y[i] as f32;
+//         result += delta * delta;
+//     }
+//     result
+// }
+
+pub fn redmean_sq(x: [u8; 3], y: [u8; 3]) -> f32 {
+    let mut ds: [f32; 3] = [0.0; 3];
+    for (i, d) in ds.iter_mut().enumerate() {
+        let delta = x[i] as f32 - y[i] as f32;
+        *d = delta * delta;
+    }
+    let rm = (x[0] as f32 + y[0] as f32) / 2.0;
+    let result = (2.0 + rm / 256.0) * ds[0] + 4.0 * ds[1] + (2.0 + (255.0 - rm) / 256.0) * ds[2];
+    // normalize value to the range 0.0 - 1.0
+    result / 584971.0
+}
+
+type Rgb = [u8; 3];
+
+const PICO8_PALETTE: [Rgb; 32] = [
+    // 0-15: standard palette
+    [0x00, 0x00, 0x00],
+    [0x1d, 0x2b, 0x53],
+    [0x7e, 0x25, 0x53],
+    [0x00, 0x87, 0x51],
+    [0xab, 0x52, 0x36],
+    [0x5f, 0x57, 0x4f],
+    [0xc2, 0xc3, 0xc7],
+    [0xff, 0xf1, 0xe8],
+    [0xff, 0x00, 0x4d],
+    [0xff, 0xa3, 0x00],
+    [0xff, 0xec, 0x27],
+    [0x00, 0xe4, 0x36],
+    [0x29, 0xad, 0xff],
+    [0x83, 0x76, 0x9c],
+    [0xff, 0x77, 0xa8],
+    [0xff, 0xcc, 0xaa],
+    // 16-31: secret palette
+    // In PICO-8, these indexes start at 128
+    [0x29, 0x18, 0x14],
+    [0x11, 0x1d, 0x35],
+    [0x42, 0x21, 0x36],
+    [0x12, 0x53, 0x59],
+    [0x74, 0x2f, 0x29],
+    [0x49, 0x33, 0x3b],
+    [0xa2, 0x88, 0x79],
+    [0xf3, 0xef, 0x7d],
+    [0xbe, 0x12, 0x50],
+    [0xff, 0x6c, 0x24],
+    [0xa8, 0xe7, 0x2e],
+    [0x00, 0xb5, 0x43],
+    [0x06, 0x5a, 0xb5],
+    [0x75, 0x46, 0x65],
+    [0xff, 0x6e, 0x59],
+    [0xff, 0x9d, 0x81],
+];
+
+#[derive(Clone, Copy)]
+pub struct PicoColor(u8);
+
+impl PicoColor {
+    pub fn to_rgb(self) -> Rgb {
+        let idx = ((self.0 & 0x80) >> 3) | (self.0 & 0x0F);
+        PICO8_PALETTE[idx as usize]
+    }
+}
+
+pub struct PicoImage {
+    pub palette: [PicoColor; 16],
+    pub width: u32,
+    pub height: u32,
+    pub data: Vec<u8>,
+}
+
+impl PicoImage {
+    pub fn new(width: u32, height: u32) -> Self {
+        let mut palette = [PicoColor(0); 16];
+        for (i, color) in palette.iter_mut().enumerate() {
+            *color = PicoColor(i as u8);
+        }
+        Self {
+            palette,
+            width,
+            height,
+            data: vec![0; (width * height) as usize],
+        }
+    }
+
+    pub fn to_indexed(img: &RgbImage) -> Self {
+        // Brute-force nearest-color conversion
+        // TODO begin using this everywhere
+        let mut result = Self::new(img.width(), img.height());
+        for y in 0..img.height() {
+            let offset = (y * img.width()) as usize;
+            for x in 0..img.width() {
+                let original = img.get_pixel(x, y).0;
+                let best = (0..16)
+                    .min_by_key(|&i| {
+                        let dist_sq = redmean_sq(original, PicoColor(i).to_rgb());
+                        NotNan::new(dist_sq).unwrap()
+                    })
+                    .unwrap();
+                result.data[offset + x as usize] = best;
+            }
+        }
+        result
+    }
+
+    pub fn get(&self, x: u32, y: u32) -> u8 {
+        self.data[(y * self.width + x) as usize]
+    }
+
+    pub fn to_rgb(&self) -> RgbImage {
+        let mut result = RgbImage::new(self.width, self.height);
+        for y in 0..self.height {
+            for x in 0..self.width {
+                let index = self.get(x, y);
+                let color = self.palette[index as usize];
+                result.put_pixel(x, y, image::Rgb(color.to_rgb()));
+            }
+        }
+        result
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_to_rgb() {
+        let to_rgb = |i| PicoColor(i).to_rgb();
+
+        // Regular palette
+        assert_eq!(to_rgb(0), [0, 0, 0]);
+        assert_eq!(to_rgb(1), PICO8_PALETTE[1]);
+        assert_eq!(to_rgb(2), to_rgb(0x12));
+        assert_eq!(to_rgb(3), to_rgb(0x73));
+
+        // Secret palette
+        assert_eq!(to_rgb(128), PICO8_PALETTE[16]);
+        assert_eq!(to_rgb(128 + 15), PICO8_PALETTE[31]);
+        assert_eq!(to_rgb(128 + 16), PICO8_PALETTE[16]);
+    }
+
+    #[test]
+    fn test_quantization_round_trip() {
+        for i in 0..16 {
+            let mut img = RgbImage::new(1, 1);
+            let input_color = PICO8_PALETTE[i];
+            img.put_pixel(0, 0, image::Rgb(input_color));
+            let indexed = PicoImage::to_indexed(&img);
+            let img = indexed.to_rgb();
+            let output_color = img.get_pixel(0, 0).0;
+            assert_eq!(input_color, output_color);
+        }
+    }
+}

rust/src/main.rs

@@ -1,9 +1,11 @@
 mod encode;
+mod indexed;
 
 use anyhow::{bail, Result};
 use clap::Parser;
 use encode::{choose_encoding, Region};
 use image::{DynamicImage, RgbImage};
+use indexed::PicoImage;
 
 #[derive(Debug, Parser)]
 #[command()]
@@ -32,10 +34,17 @@ fn main() -> Result<()> {
             output.put_pixel(x, y, *pixel);
         }
     }
+
+    // Lossy compression
     let error_scale_factor = 0.2; // Adjusted by feel. Q=0 is unrecognizable but not blank.
     let error_per_pixel = (1.0 - args.max_err.unwrap_or(80.0) / 100.0) * error_scale_factor;
     let sq_error_per_pixel = error_per_pixel * error_per_pixel;
     quadtree_quant(&mut output, sq_error_per_pixel * 128.0 * 128.0);
+
+    // Convert to indexed
+    let indexed = PicoImage::to_indexed(&output);
+    output = indexed.to_rgb();
+
     output.save(&args.output)?;
     println!("Wrote {}", &args.output);
     Ok(())
@@ -55,15 +64,3 @@ fn quadtree_quant(img: &mut RgbImage, max_err: f32) {
     );
     tree.apply(img, region);
 }
-
-pub fn redmean_sq(x: [u8; 3], y: [u8; 3]) -> f32 {
-    let mut ds: [f32; 3] = [0.0, 0.0, 0.0];
-    for (i, d) in ds.iter_mut().enumerate() {
-        let delta = x[i] as f32 - y[i] as f32;
-        *d = delta * delta;
-    }
-    let rm = (x[0] as f32 + y[0] as f32) / 2.0;
-    let result = (2.0 + rm / 256.0) * ds[0] + 4.0 * ds[1] + (2.0 + (255.0 - rm) / 256.0) * ds[2];
-    // normalize value to the range 0.0 - 1.0
-    result / 584971.0
-}