initial commit

src/main.rs

@@ -0,0 +1,139 @@
+use std::{
+    collections::{HashMap, HashSet},
+    fs::File,
+    io::{BufReader, BufWriter, Write},
+};
+
+use clap::Parser;
+use image::{DynamicImage, ImageDecoder, Rgba, codecs::png::PngDecoder};
+use png::{
+    BitDepth::Eight,
+    ColorType::{self, Indexed},
+    Encoder,
+};
+
+fn encode_plte(pixels: &[[u8; 4]], width: u32, height: u32) -> Vec<u8> {
+    let mut used_colors: HashSet<[u8; 4]> = HashSet::new();
+    for pixel in pixels {
+        used_colors.insert(*pixel);
+    }
+    let mut palette = used_colors.into_iter().collect::<Vec<[u8; 4]>>();
+    palette.sort();
+
+    let palette_map: HashMap<[u8; 4], u8> = palette
+        .iter()
+        .enumerate()
+        .map(|(i, &color)| (color, i as u8))
+        .collect();
+    println!("used palette colors: {:02x?}", palette);
+    let mut w = Vec::new();
+    {
+        let mut encoder = Encoder::new(&mut w, width, height);
+        encoder.set_color(Indexed);
+        encoder.set_palette(
+            palette
+                .iter()
+                .flat_map(|color| [color[0], color[1], color[2]])
+                .collect::<Vec<_>>(),
+        );
+        encoder.set_depth(Eight);
+        encoder.set_trns(&[0x00]);
+        encoder.set_compression(png::Compression::High);
+        let mut writer = encoder.write_header().unwrap();
+        let data: Vec<u8> = pixels.iter().map(|pixel| palette_map[pixel]).collect();
+
+        writer.write_image_data(&data).unwrap();
+    }
+    w
+}
+
+fn encode_standard(pixels: &[[u8; 4]], width: u32, height: u32) -> Vec<u8> {
+    let mut w = Vec::new();
+    {
+        let mut encoder = Encoder::new(&mut w, width, height);
+        encoder.set_color(ColorType::Rgba);
+        encoder.set_depth(Eight);
+        encoder.set_compression(png::Compression::High);
+        let mut writer = encoder.write_header().unwrap();
+        let data: Vec<u8> = pixels.iter().flat_map(|&pixel| pixel).collect();
+
+        writer.write_image_data(&data).unwrap();
+    }
+    w
+}
+
+/// Simple program to optimise palette based PNGs
+#[derive(Parser, Debug)]
+#[command(version, about, long_about = None)]
+struct Args {
+    /// Path of the file to import
+    input_file: String,
+
+    /// Path of the file to export
+    output_file: String,
+
+    /// Colors to substitute into color variants, hex-formatted from -> to as rrggbb rrggbb
+    #[arg(short, long, num_args=2, value_names = ["from", "to"])]
+    subs: Vec<String>,
+}
+
+fn parse_rgb(hex: String) -> Result<[u8; 4], String> {
+    if hex.len() != 6 {
+        return Err("expected rrggbb".into());
+    }
+
+    let bytes: Vec<u8> = (0..6)
+        .step_by(2)
+        .map(|i| u8::from_str_radix(&hex[i..i + 2], 16))
+        .collect::<Result<_, _>>()
+        .map_err(|_| "invalid hex color")?;
+
+    Ok([bytes[0], bytes[1], bytes[2], 0xff])
+}
+
+fn main() {
+    let args = Args::parse();
+    let in_path = args.input_file;
+    let out_path = args.output_file;
+
+    let asubs_len = args.subs.len();
+
+    let pairs: HashMap<[u8; 4], [u8; 4]> = args
+        .subs
+        .into_iter()
+        .map(|item| parse_rgb(item).unwrap())
+        .collect::<Vec<_>>()
+        .as_chunks::<2>()
+        .0
+        .into_iter()
+        .map(|[a, b]| (*a, *b))
+        .collect();
+    if asubs_len == 0 {
+        println!("no substitutions selected, only re-encoding")
+    } else {
+        println!("substitutions to use: {:?}", pairs);
+    }
+
+    let decoder = PngDecoder::new(BufReader::new(File::open(in_path).unwrap())).unwrap();
+    let (width, height) = decoder.dimensions();
+    let img = DynamicImage::from_decoder(decoder).unwrap();
+
+    let image_data = img
+        .to_rgba8()
+        .pixels()
+        .map(|&Rgba(px)| *pairs.get(&px).or_else(|| Some(&px)).unwrap())
+        .collect::<Vec<_>>();
+
+    let plte_data = encode_plte(image_data.as_slice(), width, height);
+    let std_data = encode_standard(image_data.as_slice(), width, height);
+
+    println!("plte data has length {}", plte_data.len(),);
+    println!("std data has length {}", std_data.len(),);
+
+    let mut file = File::create(out_path).unwrap();
+    if plte_data.len() <= std_data.len() {
+        file.write_all(&plte_data).unwrap();
+    } else {
+        file.write_all(&std_data).unwrap();
+    }
+}