about summary refs log tree commit diff
path: root/src/main.rs
blob: 02c2950bba06edbec1d01e2cbf39eda498f0d6be (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
use std::{collections::HashMap, env::args};

use image::io::Reader as ImageReader;
use image::Rgb;

use rayon::prelude::*;

const MAX_COLORS: usize = 256;

const TOLERANCE: f32 = 0.025;
const RGB_TOLERANCE: f32 = 10.0 * TOLERANCE;

fn main() {
    let filename = args().nth(1).unwrap();
    let outname = args().nth(2).unwrap();
    // The percent of RGB value difference a color has to surpass to be considered unique

    let imageread = ImageReader::open(&filename).expect("Failed to open image!");
    let mut image = imageread
        .decode()
        .expect("Failed to decode image!")
        .into_rgb8();

    let selected_colors = quantize(image.pixels());

    // Max complexity is O(n * max_colors)
    image.pixels_mut().par_bridge().for_each(|color| {
        let mut min_difference = f32::MAX;
        let mut min_difference_color = *color;

        for selected_color in &selected_colors {
            let difference = rgb_difference(color, selected_color);
            if difference < min_difference {
                min_difference = difference;
                min_difference_color = *selected_color;
            }
        }
        *color = min_difference_color
    });

    image.save(outname).expect("Failed to write out");
}

fn quantize<'a, T>(pixels: T) -> Vec<Rgb<u8>>
where
    T: Iterator<Item = &'a Rgb<u8>>,
{
    let mut colors: HashMap<Rgb<u8>, usize> = HashMap::new();

    //count pixels
    for pixel in pixels {
        match colors.get_mut(pixel) {
            None => {
                colors.insert(*pixel, 1);
            }
            Some(n) => *n += 1,
        }
    }

    let mut sorted: Vec<(Rgb<u8>, usize)> = colors.into_par_iter().collect();
    sorted.sort_by(|(colour1, freq1), (colour2, freq2)| {
        freq2
            .cmp(freq1)
            .then(colour2[0].cmp(&colour1[0]))
            .then(colour2[1].cmp(&colour1[1]))
            .then(colour2[2].cmp(&colour1[2]))
    });

    let mut selected_colors: Vec<Rgb<u8>> = Vec::with_capacity(MAX_COLORS);

    for (key, _value) in sorted.iter() {
        if selected_colors.len() >= MAX_COLORS {
            break;
        } else if selected_colors
            .iter()
            .all(|color| rgb_difference(key, color) > RGB_TOLERANCE)
        {
            selected_colors.push(*key);
        }
    }

    selected_colors
}

#[allow(clippy::many_single_char_names)]
#[inline(always)]
fn rgb_difference(a: &Rgb<u8>, z: &Rgb<u8>) -> f32 {
    let (a, b, c) = pixel_rgb_to_hsv(a);
    let (d, e, f) = pixel_rgb_to_hsv(z);

    (((c - f) * (c - f)) + ((a - d).abs() / 90.0) + (b - e).abs()) as f32
}

#[allow(clippy::float_cmp)]
fn pixel_rgb_to_hsv(a: &Rgb<u8>) -> (f32, f32, f32) {
    let (r, g, b) = (
        a.0[0] as f32 / 256.0,
        a.0[1] as f32 / 256.0,
        a.0[2] as f32 / 256.0,
    );

    let value = r.max(g.max(b));
    let x_min = r.min(g.min(b));
    let chroma = value - x_min;

    let hue = if chroma == 0.0 {
        0.0
    } else if value == r {
        60.0 * ((g - b) / chroma)
    } else if value == g {
        60.0 * (2.0 + (b - r) / chroma)
    } else if value == b {
        60.0 * (4.0 + (r - g) / chroma)
    } else {
        unreachable!()
    };

    let value_saturation = if value == 0.0 { 0.0 } else { chroma / value };

    /* Rotate the color wheel counter clockwise to the negative location
          |       Keep the wheel in place and remove any full rotations
     _____V____ _____V____
    |          |          |*/
    ((hue + 360.0) % 360.0, value_saturation * 2.0, value * 2.0)
}