refactor: code organisation
This commit is contained in:
106
src/colours.rs
Normal file
106
src/colours.rs
Normal file
@@ -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<u8, HashMap<u8, HashMap<u8, Vec<Point>>>>;
|
||||||
|
pub type Colour = Rgb<u8>;
|
||||||
|
|
||||||
|
#[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<Point>,
|
||||||
|
}
|
||||||
|
|
||||||
|
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<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: &ColourMap) -> Option<Match> {
|
||||||
|
// use redmean to calculate the distance between colours
|
||||||
|
let mut closest: Option<Match> = 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<Colour>) -> HashMap<Colour, Match> {
|
||||||
|
let mut output_cols: HashMap<Colour, Match> = 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<Colour, Match> {
|
||||||
|
let img_2_cols = extract_colours_set(img_2);
|
||||||
|
|
||||||
|
let mut output_cols: HashMap<Colour, Match> = 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
|
||||||
|
}
|
||||||
92
src/main.rs
92
src/main.rs
@@ -1,95 +1,27 @@
|
|||||||
mod cli;
|
mod cli;
|
||||||
|
mod colours;
|
||||||
|
|
||||||
use clap::Parser;
|
use clap::Parser;
|
||||||
use std::collections::{HashMap, HashSet};
|
use std::collections::HashMap;
|
||||||
|
|
||||||
use dotenv::dotenv;
|
use dotenv::dotenv;
|
||||||
use image::{DynamicImage, GenericImageView, Rgb};
|
|
||||||
|
|
||||||
use crate::cli::Cli;
|
use crate::cli::Cli;
|
||||||
|
use crate::colours::{Colour, Match, extract_colours, map_image_image, map_image_list};
|
||||||
// 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>;
|
|
||||||
|
|
||||||
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<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> {
|
|
||||||
// use redmean to calculate the distance between colours
|
|
||||||
let mut closest: Option<Colour> = 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<Colour>) -> HashMap<Colour, Colour> {
|
|
||||||
let mut output_cols: HashMap<Colour, Colour> = 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<Colour, Colour> {
|
|
||||||
let img_2_cols = extract_colours_set(img_2);
|
|
||||||
|
|
||||||
let mut output_cols: HashMap<Colour, Colour> = 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
|
|
||||||
}
|
|
||||||
|
|
||||||
fn encode_hex(c: &Colour) -> String {
|
fn encode_hex(c: &Colour) -> String {
|
||||||
format!("#{:02x}{:02x}{:02x}", c[0], c[1], c[2])
|
format!("#{:02x}{:02x}{:02x}", c[0], c[1], c[2])
|
||||||
}
|
}
|
||||||
|
|
||||||
fn pretty_print(h: HashMap<Colour, Colour>) {
|
fn pretty_print(h: HashMap<Colour, Match>) {
|
||||||
h.iter()
|
h.iter().for_each(|(k, v)| {
|
||||||
.for_each(|(k, v)| println!("{}: {}", encode_hex(k), encode_hex(v)))
|
println!(
|
||||||
|
"{}: {} ({:?})",
|
||||||
|
encode_hex(k),
|
||||||
|
encode_hex(&v.colour),
|
||||||
|
v.positions
|
||||||
|
)
|
||||||
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
fn main() {
|
fn main() {
|
||||||
|
|||||||
Reference in New Issue
Block a user