color_theme/
lib.rs

1use image::{RgbImage, Rgb};
2
3/// Sort the pixels by the color channel with the most range
4fn median_cut_sort_bucket(pixel_bucket: &mut [Rgb<u8>]) {
5    let mut r_minmax = (u8::MAX, 0);
6    let mut g_minmax = (u8::MAX, 0);
7    let mut b_minmax = (u8::MAX, 0);
8
9    for pixel in pixel_bucket.iter() {
10        let r_val = pixel.0[0];
11        let g_val = pixel.0[1];
12        let b_val = pixel.0[2];
13
14        r_minmax.0 = r_minmax.0.min(r_val);
15        r_minmax.1 = r_minmax.1.max(r_val);
16
17        g_minmax.0 = g_minmax.0.min(g_val);
18        g_minmax.1 = g_minmax.1.max(g_val);
19
20        b_minmax.0 = b_minmax.0.min(b_val);
21        b_minmax.1 = b_minmax.1.max(b_val);
22    }
23
24
25    let ranges = [
26        r_minmax.1 - r_minmax.0,
27        g_minmax.1 - g_minmax.0,
28        b_minmax.1 - b_minmax.0
29    ];
30
31    let largest_range_channel = ranges.iter()
32        .enumerate()
33        .max_by_key(|x| x.1)
34        .expect("Guaranteed to be a max")
35        .0;
36    
37    pixel_bucket.sort_by_key(|x| x.0[largest_range_channel]);
38}
39
40fn average_pixels(pixels: &[Rgb<u8>]) -> Rgb<u8> {
41    let pixels_size = pixels.len();
42
43    let mut total_pixel_values = (0, 0, 0);
44    for pixel in pixels {
45        total_pixel_values.0 += pixel.0[0] as usize;
46        total_pixel_values.1 += pixel.0[1] as usize;
47        total_pixel_values.2 += pixel.0[2] as usize;
48    }
49
50    let pixel_values_average = [
51        (total_pixel_values.0/pixels_size) as u8,
52        (total_pixel_values.1/pixels_size) as u8,
53        (total_pixel_values.2/pixels_size) as u8
54    ];
55
56    Rgb::from(pixel_values_average)
57}
58
59
60fn bucket_from_image_mut(image: &mut [u8], bucket_size: usize, bucket_index: usize) -> &mut [Rgb<u8>] {
61    let slice_start = bucket_size*bucket_index*3;
62    let slice_end = bucket_size*(bucket_index+1)*3;
63    let bucket = image.get_mut(slice_start..slice_end).expect("Guaranteed to be in bounds");
64
65    // SAFETY: `bucket` should consist of chunks of 3 contiguous `u8`s,
66    // which is the same representation as `Rgb<u8>` ([u8; 3]).
67    //
68    // The slice is guaranteed to be within `image`, as `bucket.len()/3` rounds down.
69    // e.g. if `bucket.len() == 8`, the returned slice, will have a length of 2 (6 bytes).
70    unsafe { std::slice::from_raw_parts_mut(bucket.as_mut_ptr().cast(), bucket.len()/3) }
71}
72
73/// https://en.wikipedia.org/wiki/Median_cut
74/// 
75/// Some pixels at the end of the image may not be included.
76pub fn median_cut_palette(rgb_image: &mut RgbImage, palette_n: u8) -> Vec<Rgb<u8>> {
77    let pixel_count = rgb_image.len() / 3;
78
79    let iterations = (palette_n as f32).log2().ceil() as u32 + 1;
80    let mut bucket_count = 1;
81    let mut bucket_size = pixel_count;
82    for i in 0..iterations {
83        bucket_count = 2_u32.pow(i) as usize;
84        bucket_size = pixel_count / bucket_count;
85
86        for bucket_index in 0..bucket_count {
87            let bucket_pixels = bucket_from_image_mut(rgb_image, bucket_size, bucket_index);
88            median_cut_sort_bucket(bucket_pixels);
89        }
90    }
91
92    debug_assert!(bucket_count >= palette_n as usize);
93    let mut colors = Vec::with_capacity(palette_n as usize);
94    for bucket_index in 0..palette_n as usize {
95        let bucket_pixels = bucket_from_image_mut(rgb_image, bucket_size, bucket_index);
96        colors.push(average_pixels(&bucket_pixels));
97    }
98
99    colors
100}
101
102fn calculate_saturation(pixel: Rgb<u8>) -> u8 {
103    let max = *pixel.0.iter().max().expect("Impossible as color array is not empty");
104    let min = *pixel.0.iter().min().expect("Impossible as color array is not empty");
105    (255.0 * (1.0 - min as f32/max as f32)) as u8
106}
107
108fn change_brightness(pixel: Rgb<u8>, target_brightness: u8) -> Rgb<u8> {
109    let max_brightness = *pixel.0.iter().max().expect("Impossible as color array is not empty");
110    let multiplier = target_brightness as f32 / max_brightness as f32;
111
112    let new_values = [
113        (pixel.0[0] as f32 * multiplier).min(255.0) as u8,
114        (pixel.0[1] as f32 * multiplier).min(255.0) as u8,
115        (pixel.0[2] as f32 * multiplier).min(255.0) as u8
116    ];
117
118    Rgb::from(new_values)
119}
120
121/// Choose most saturated color in palette,
122/// and adjust brightness.
123pub fn get_theme_color(palette: &[Rgb<u8>], brightness: Option<u8>) -> Option<Rgb<u8>> {
124    let theme = *palette.iter().max_by_key(|&&x| calculate_saturation(x))?;
125
126    Some(match brightness {
127        Some(b) => change_brightness(theme, b),
128        None => theme
129    })
130}
color_theme/lib.rs

color_theme/
lib.rs
