139 lines
4.1 KiB
Rust
139 lines
4.1 KiB
Rust
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();
|
|
}
|
|
}
|