ft: list_image and image_image functions

src/main.rs

@@ -1,21 +1,48 @@
-use std::{collections::{HashMap, HashSet}, env, num::ParseIntError};
+use std::{
+    collections::{HashMap, HashSet},
+    env,
+    num::ParseIntError,
+};
 
 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 Colour = Rgb<u8>;
 
-pub fn decode_hex(s: &str) -> Result<Colour, ParseIntError> {
-    Ok(Rgb((0..s.len())
+pub fn decode_hex(s: &str) -> Result<Colour, Box<dyn std::error::Error>> {
+    let u8s: Vec<u8> = (0..s.len())
         .step_by(2)
         .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 {
-    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> {
@@ -24,44 +51,74 @@ fn closest_colour(colour: Colour, all: &ColourSet) -> Option<Colour> {
     let mut closest: Option<Colour> = None;
     let mut dist: f32 = f32::MAX;
     for (r, gbs) in all.iter() {
-        let r_bar = 0.5 * (r + colour[0]) as f32;
-        let delta_r = ((r + colour[0]) * (r + colour[0])) as f32;
+        let r_bar = 0.5 * (*r as f32 + colour[0] as f32); //FIXME (low priority): i dont like the
+        //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() {
-            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 {
-                let delta_b = ((b + colour[2]) * (b + colour[2])) as f32;
-
-                let delta = (2.0 + r_bar/256.0) * delta_r + 4.0 * delta_g + (2.0 + (255.0 - r_bar)/256.0) * delta_b;
+                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
+                    + 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]));
+                    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() {
     dotenv().ok();
     pretty_env_logger::init();
 
+    // TODO (high priority): proper arg parsing
     let args: Vec<String> = env::args().collect();
-    if args.len() > 1 {
-        let img_file = &args[1];
-        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();
-
-        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);
-            }
-        }
+    if args.len() > 2 {
+        // list of colours and image
+        // let out_cols = map_image_list(&args[1], &args);
+        let out_cols = map_image_image(&args[1], &args[2]);
 
+        out_cols.iter().for_each(|(k, v)| {
+            println!(
+                "({}, {}, {}): ({}, {}, {})",
+                k[0], k[1], k[2], v[0], v[1], v[2]
+            )
+        })
     }
 }