hs_hackathon_vision/raw/
colors.rs

1use crate::raw::{
2    bounding_box::BoundingBox, BLUE, GREEN, LAB_BLUE, LAB_GREEN, LAB_RED, LAB_WHITE, RED,
3    WHITE_BRIGHTNESS_THRESHOLD, WHITE_SATURATION_THRESHOLD, WHITE_VALUE_THRESHOLD,
4};
5use eyre::ContextCompat;
6use image::{DynamicImage, GenericImageView, Rgb, Rgba};
7use palette::{
8    color_difference::EuclideanDistance,
9    encoding::{Linear, Srgb as GammaSrgb},
10    white_point::D65,
11    FromColor, Hsv, Lab, LinSrgb, Srgb, Srgba,
12};
13use std::collections::HashMap;
14
15#[derive(PartialEq, Clone, Copy, Eq, Hash, Debug)]
16pub enum Color {
17    Red,
18    Blue,
19    Green,
20    White,
21    Unknown,
22}
23
24impl Color {
25    pub fn lab_rgb(&self) -> Lab {
26        match *self {
27            Color::Red => Lab::from_components(LAB_RED),
28            Color::Blue => Lab::from_components(LAB_BLUE),
29            Color::Green => Lab::from_components(LAB_GREEN),
30            Color::White => Lab::from_components(LAB_WHITE),
31            Color::Unknown => {
32                panic!("Unknow colour, current detected colours are White/Red/Blue/Green")
33            }
34        }
35    }
36}
37
38pub const fn colors() -> &'static [Color] {
39    &[Color::Red, Color::Blue, Color::Green]
40}
41
42fn is_white(pixel: Rgba<u8>) -> bool {
43    let max_pixel = u8::max(pixel[0], u8::max(pixel[1], pixel[2]));
44    let linear_rgb = to_srgb(pixel).into_linear();
45
46    let hsv = Hsv::from_color(linear_rgb);
47    (hsv.saturation <= WHITE_SATURATION_THRESHOLD && hsv.value >= WHITE_VALUE_THRESHOLD)
48        || (max_pixel > WHITE_BRIGHTNESS_THRESHOLD)
49}
50pub fn detect_color(img: &DynamicImage, bbox: &BoundingBox) -> Color {
51    let mut color_count = HashMap::new();
52
53    for (x, y) in bbox.iter() {
54        let color: Color = get_color_lab(img.get_pixel(x, y));
55        *color_count.entry(color).or_insert(0) += 1;
56    }
57
58    let max = color_count
59        .into_iter()
60        .max_by_key(|color| color.1)
61        .expect("No maximum color found!");
62    max.0
63}
64
65fn get_color_lab(pixel: Rgba<u8>) -> Color {
66    let distances: Vec<(i32, Color)> = colors()
67        .iter()
68        .map(|color| {
69            let rgb = color.lab_rgb();
70            let s_pixel = to_srgb(pixel);
71
72            let lab: Lab = Lab::from_color(s_pixel);
73
74            // Compute euclidean distance
75            let distance = lab.distance_squared(rgb) as i32;
76            (distance, *color)
77        })
78        .collect();
79    // Minimum distance is the detected color
80    let min_distance = distances
81        .iter()
82        .min_by_key(|(distance, _)| *distance)
83        .expect("No maximum color found!");
84
85    min_distance.1
86}
87
88fn to_srgb(pixel: Rgba<u8>) -> palette::rgb::Rgb {
89    Srgb::new(
90        pixel.0[0] as f32 / 255.0,
91        pixel.0[1] as f32 / 255.0,
92        pixel.0[2] as f32 / 255.0,
93    )
94}
hs_hackathon_vision/raw/colors.rs

hs_hackathon_vision/raw/
colors.rs
