diff --git a/src/colours.rs b/src/colours.rs new file mode 100644 index 0000000..a200727 --- /dev/null +++ b/src/colours.rs @@ -0,0 +1,106 @@ +use std::collections::{HashMap, HashSet}; + +use image::{DynamicImage, GenericImageView, Rgb}; +use log::info; +use std::fmt::Debug; + +use crate::encode_hex; + +pub type ColourMap = HashMap>>>; +pub type Colour = Rgb; + +#[derive(Clone)] +pub struct Point { + x: u32, + y: u32, +} + +impl Debug for Point { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(&format!("{}, {}", &self.x, &self.y)) + } +} + +pub struct Match { + pub colour: Colour, + pub positions: Vec, +} + +pub fn extract_colours(img: DynamicImage) -> ColourMap { + let mut out: ColourMap = HashMap::new(); + for (x, y, p) in img.pixels() { + out.entry(p[0]) + .or_insert_with(HashMap::new) + .entry(p[1]) + .or_insert_with(HashMap::new) + .entry(p[2]) + .or_insert_with(Vec::new) + .push(Point { x, y }); + } + out +} + +fn extract_colours_set(img: DynamicImage) -> HashSet> { + 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: &ColourMap) -> Option { + // use redmean to calculate the distance between colours + let mut closest: Option = None; + let mut dist: f32 = f32::MAX; + + let (colour_r, colour_g, colour_b) = (colour[0] as f32, colour[1] as f32, colour[2] as f32); + for (r, gbs) in all.iter() { + let r_bar = 0.5 * (*r as f32 + colour_r); + let r_squared = (*r as f32 - colour_r) * (*r as f32 - colour_r); + for (g, bs) in gbs.iter() { + let g_squared = (*g as f32 - colour_g) * (*g as f32 - colour_g); + for (b, vec) in bs.iter() { + let b_squared = (*b as f32 - colour_b) * (*b as f32 - colour_b); + let delta = (2.0 + r_bar / 256.0) * r_squared + + 4.0 * g_squared + + (2.0 + (255.0 - r_bar) / 256.0) * b_squared; + if delta < dist { + dist = delta; + closest = Some(Match { + colour: Rgb([*r, *g, *b]), + positions: vec.to_vec(), // FIXME (low priority): would be more efficient if + // i can make these a borrow instead, but that + // leads to lifetime issues + }); + } + } + } + } + return closest; +} + +pub fn map_image_list(all_cols: &ColourMap, cols: Vec) -> HashMap { + let mut output_cols: HashMap = HashMap::new(); + + for c in cols { + info!("Matching colour {}", encode_hex(&c)); + if let Some(col) = closest_colour(&c, &all_cols) { + output_cols.insert(c, col); + } + } + output_cols +} + +pub fn map_image_image(img_1_cols: &ColourMap, img_2: DynamicImage) -> HashMap { + let img_2_cols = extract_colours_set(img_2); + + let mut output_cols: HashMap = HashMap::new(); + img_2_cols.iter().for_each(|c| { + info!("Matching colour {}", encode_hex(&c)); + if let Some(col) = closest_colour(c, &img_1_cols) { + output_cols.insert(*c, col); + } + }); + + output_cols +} diff --git a/src/main.rs b/src/main.rs index 79ad5ed..3a7120b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,95 +1,27 @@ mod cli; +mod colours; use clap::Parser; -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use dotenv::dotenv; -use image::{DynamicImage, GenericImageView, Rgb}; use crate::cli::Cli; - -// TODO (medium priority): replace hashset with another hashmap>; -// that way we can track pixel locations for colours. -type ColourSet = HashMap>>; -type Colour = Rgb; - -fn extract_colours(img: DynamicImage) -> ColourSet { - 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> { - 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 { - // use redmean to calculate the distance between colours - let mut closest: Option = None; - let mut dist: f32 = f32::MAX; - - let (colour_r, colour_g, colour_b) = (colour[0] as f32, colour[1] as f32, colour[2] as f32); - for (r, gbs) in all.iter() { - let r_bar = 0.5 * (*r as f32 + colour_r); - let r_squared = (*r as f32 - colour_r) * (*r as f32 - colour_r); - for (g, bs) in gbs.iter() { - let g_squared = (*g as f32 - colour_g) * (*g as f32 - colour_g); - for b in bs { - let b_squared = (*b as f32 - colour_b) * (*b as f32 - colour_b); - let delta = (2.0 + r_bar / 256.0) * r_squared - + 4.0 * g_squared - + (2.0 + (255.0 - r_bar) / 256.0) * b_squared; - if delta < dist { - dist = delta; - closest = Some(Rgb([*r, *g, *b])); - } - } - } - } - return closest; -} - -fn map_image_list(all_cols: &ColourSet, cols: Vec) -> HashMap { - let mut output_cols: HashMap = HashMap::new(); - - for c in cols { - if let Some(col) = closest_colour(&c, &all_cols) { - output_cols.insert(c, col); - } - } - output_cols -} - -fn map_image_image(img_1_cols: &ColourSet, img_2: DynamicImage) -> HashMap { - let img_2_cols = extract_colours_set(img_2); - - let mut output_cols: HashMap = HashMap::new(); - img_2_cols.iter().for_each(|c| { - if let Some(col) = closest_colour(c, &img_1_cols) { - output_cols.insert(*c, col); - } - }); - - output_cols -} +use crate::colours::{Colour, Match, extract_colours, map_image_image, map_image_list}; fn encode_hex(c: &Colour) -> String { format!("#{:02x}{:02x}{:02x}", c[0], c[1], c[2]) } -fn pretty_print(h: HashMap) { - h.iter() - .for_each(|(k, v)| println!("{}: {}", encode_hex(k), encode_hex(v))) +fn pretty_print(h: HashMap) { + h.iter().for_each(|(k, v)| { + println!( + "{}: {} ({:?})", + encode_hex(k), + encode_hex(&v.colour), + v.positions + ) + }) } fn main() {