ft: list_image and image_image functions

This commit is contained in:
2026-06-24 15:40:44 +02:00
parent 68313eaba9
commit 9f722b1040

View File

@@ -1,21 +1,48 @@
use std::{collections::{HashMap, HashSet}, env, num::ParseIntError}; use std::{
collections::{HashMap, HashSet},
env,
num::ParseIntError,
};
use dotenv::dotenv; use dotenv::dotenv;
use image::{DynamicImage, ImageReader, Rgb}; use image::{DynamicImage, GenericImageView, ImageReader, Rgb};
// TODO (medium priority): replace hashset with another hashmap<u8, Vec<(u32, u32)>>;
// that way we can track pixel locations for colours.
type ColourSet = HashMap<u8, HashMap<u8, HashSet<u8>>>; type ColourSet = HashMap<u8, HashMap<u8, HashSet<u8>>>;
type Colour = Rgb<u8>; type Colour = Rgb<u8>;
pub fn decode_hex(s: &str) -> Result<Colour, ParseIntError> { pub fn decode_hex(s: &str) -> Result<Colour, Box<dyn std::error::Error>> {
Ok(Rgb((0..s.len()) let u8s: Vec<u8> = (0..s.len())
.step_by(2) .step_by(2)
.map(|i| u8::from_str_radix(&s[i..i + 2], 16)) .map(|i| u8::from_str_radix(&s[i..i + 2], 16))
.collect()) .collect::<Result<Vec<u8>, ParseIntError>>()?;
)
let rgb: [u8; 3] = u8s
.try_into()
.map_err(|_| <&str as Into<String>>::into("Hex string must decode to exactly 3 bytes"))?;
Ok(Rgb(rgb))
} }
fn extract_colours(img: DynamicImage) -> ColourSet { fn extract_colours(img: DynamicImage) -> ColourSet {
todo!() let mut out: ColourSet = HashMap::new();
for (_x, _y, p) in img.pixels() {
out.entry(p[0])
.or_insert_with(HashMap::new)
.entry(p[1])
.or_insert_with(HashSet::new)
.insert(p[2]);
}
out
}
fn extract_colours_set(img: DynamicImage) -> HashSet<Rgb<u8>> {
let mut out = HashSet::new();
img.pixels().for_each(|(_x, _y, p)| {
out.insert(p);
});
out.into_iter().map(|c| Rgb([c[0], c[1], c[2]])).collect()
} }
fn closest_colour(colour: Colour, all: &ColourSet) -> Option<Colour> { fn closest_colour(colour: Colour, all: &ColourSet) -> Option<Colour> {
@@ -24,44 +51,74 @@ fn closest_colour(colour: Colour, all: &ColourSet) -> Option<Colour> {
let mut closest: Option<Colour> = None; let mut closest: Option<Colour> = None;
let mut dist: f32 = f32::MAX; let mut dist: f32 = f32::MAX;
for (r, gbs) in all.iter() { for (r, gbs) in all.iter() {
let r_bar = 0.5 * (r + colour[0]) as f32; let r_bar = 0.5 * (*r as f32 + colour[0] as f32); //FIXME (low priority): i dont like the
let delta_r = ((r + colour[0]) * (r + colour[0])) as f32; //clunky as f32s everywhere
let r_squared = (*r as f32 - colour[0] as f32) * (*r as f32 - colour[0] as f32);
for (g, bs) in gbs.iter() { for (g, bs) in gbs.iter() {
let delta_g = ((g + colour[1]) * (g + colour[1])) as f32; let g_squared = (*g as f32 - colour[1] as f32) * (*g as f32 - colour[1] as f32);
for b in bs { for b in bs {
let delta_b = ((b + colour[2]) * (b + colour[2])) as f32; let b_squared = (*b as f32 - colour[2] as f32) * (*b as f32 - colour[2] as f32);
let delta = (2.0 + r_bar / 256.0) * r_squared
let delta = (2.0 + r_bar/256.0) * delta_r + 4.0 * delta_g + (2.0 + (255.0 - r_bar)/256.0) * delta_b; + 4.0 * g_squared
+ (2.0 + (255.0 - r_bar) / 256.0) * b_squared;
if delta < dist { if delta < dist {
dist = delta; dist = delta;
closest = Some(Rgb([*r,*g,*b])); closest = Some(Rgb([*r, *g, *b]));
} }
} }
} }
} }
return closest return closest;
} }
fn map_image_list(img_file: &str, args: &Vec<String>) -> HashMap<Colour, Colour> {
let img = ImageReader::open(img_file).unwrap().decode().unwrap(); // FIXME: unwraps
let all_cols = extract_colours(img);
let mut output_cols: HashMap<Colour, Colour> = HashMap::new();
args.iter().skip(2).for_each(|c_str| {
let c = decode_hex(c_str).unwrap(); // FIXME: unwraps
if let Some(col) = closest_colour(c, &all_cols) {
output_cols.insert(c, col);
}
});
output_cols
}
fn map_image_image(img_file_1: &str, img_file_2: &str) -> HashMap<Colour, Colour> {
let img_1 = ImageReader::open(img_file_1).unwrap().decode().unwrap();
let img_2 = ImageReader::open(img_file_2).unwrap().decode().unwrap();
let img_1_cols = extract_colours_set(img_1);
let img_2_cols = extract_colours(img_2);
let mut output_cols: HashMap<Colour, Colour> = HashMap::new();
img_1_cols.iter().for_each(|c| {
if let Some(col) = closest_colour(*c, &img_2_cols) {
output_cols.insert(*c, col);
}
});
output_cols
}
fn main() { fn main() {
dotenv().ok(); dotenv().ok();
pretty_env_logger::init(); pretty_env_logger::init();
// TODO (high priority): proper arg parsing
let args: Vec<String> = env::args().collect(); let args: Vec<String> = env::args().collect();
if args.len() > 1 { if args.len() > 2 {
let img_file = &args[1]; // list of colours and image
let img = ImageReader::open(img_file).unwrap().decode().unwrap(); // FIXME: unwraps // let out_cols = map_image_list(&args[1], &args);
let out_cols = map_image_image(&args[1], &args[2]);
let all_cols = extract_colours(img);
let mut output_cols: HashMap<Colour, Colour> = HashMap::new();
for c_str in args {
// TODO: parse args
let c = decode_hex(&c_str).unwrap(); // FIXME: unwraps;
if let Some(col) = closest_colour(c, &all_cols) {
output_cols.insert(c, col);
}
}
out_cols.iter().for_each(|(k, v)| {
println!(
"({}, {}, {}): ({}, {}, {})",
k[0], k[1], k[2], v[0], v[1], v[2]
)
})
} }
} }