iris_lib/
color_bucket.rs

1use crate::utils::mean;
2use crate::color::{Color, ColorChannel};
3
4#[cfg(feature = "image")]
5extern crate image;
6#[cfg(feature = "image")]
7use image::GenericImageView;
8
9/// Struct holding an `Vec<Color>`.
10/// Implements helpful functions for the median cut algorithm.
11///
12#[derive(Debug, PartialEq)]
13pub struct ColorBucket {
14    colors: Vec<Color>,
15}
16
17impl ColorBucket {
18    /// Creates a `ColorBucket` based on the colors passed. Returns `None` if passed an empty vector.
19    ///
20    /// # Arguments
21    ///
22    /// * `colors` - Color vector from which the mean color is created.
23    ///
24    /// # Examples
25    ///
26    /// ```
27    /// use iris_lib::color::Color;
28    /// use iris_lib::color_bucket::ColorBucket;
29    ///
30    /// let data = vec![Color { r: 15, g: 131, b: 0, a: 255 }, Color { r: 221, g: 11, b: 22, a: 130 }, Color { r: 81, g: 11, b: 16, a: 0 }];
31    /// let result = ColorBucket::from_pixels(data.clone()).expect("Passed empty color vector to test.");
32    /// ```
33    ///
34    pub fn from_pixels(pixels: Vec<Color>) -> Option<Self> {
35        if pixels.is_empty() {
36            None
37        } else {
38            Some(Self { colors: pixels })
39        }
40    }
41
42    #[cfg(feature = "image")]
43    /// Creates a `ColorBucket` based from image files. Returns `None` if unable to load the image.
44    /// 
45    /// # Arguments
46    ///
47    /// * `image_path` - Path to the target image.
48    ///
49    /// # Examples
50    ///
51    /// ```
52    /// use iris_lib::color_bucket::ColorBucket;
53    ///
54    /// if let Some(mut result) = ColorBucket::from_image("peppers.png") {
55    ///     // ...
56    /// }
57    /// ```
58    ///
59    pub fn from_image(image_path: &str) -> Option<Self> {
60        if let Ok(img) = image::open(image_path) {
61            let mut colors = Vec::new();
62
63            for p in img.pixels() {
64                let color = Color {
65                    r: p.2.0[0],
66                    g: p.2.0[1],
67                    b: p.2.0[2],
68                    a: p.2.0[3]
69                };
70
71                colors.push(color);
72            }
73
74            Some (Self { colors })
75        } else {
76            println!("Unable to load image at: {}", image_path);
77            None
78        }
79    }
80
81    /// Recursively performs the median cut algorithm on self if iteration has not reached 0 yet.
82    /// Creates two new buckets based on own colors. One bucket with values above and one bucket with value below the median, then performs the algorithm on them again.
83    ///
84    /// If iteration has reached 0 the color mean for self is pushed to the result vector.
85    ///
86    /// # Arguments
87    ///
88    /// * `iter_count` - Iteration index is used as termination criteria. Recursion stop when 0 is reached.
89    /// * `result` - Vector holding color means for each bucket in the iteration.
90    ///
91    fn recurse(&mut self, iter_count: u8, result: &mut Vec<Color>) {
92        if iter_count == 0 {
93            result.push(self.color_mean())
94        } else {
95            let new_buckets = self.median_cut();
96            if let Some(mut bucket) = new_buckets.0 {
97                bucket.recurse(iter_count - 1, result);
98            }
99            if let Some(mut bucket) = new_buckets.1 {
100                bucket.recurse(iter_count - 1, result);
101            }
102        }
103    }
104
105    /// Creates a color palette from own pixels.
106    ///
107    /// # Arguments
108    ///
109    /// * `iter_count` - number of iterations to be performed on the bucket.
110    ///
111    /// # Example
112    ///
113    /// ```
114    /// use iris_lib::color_bucket::ColorBucket;
115    /// use iris_lib::color::Color;
116    ///
117    /// let data = vec![Color { r: 15, g: 131, b: 0, a: 255 }, Color { r: 221, g: 11, b: 22, a: 130 }, Color { r: 81, g: 11, b: 16, a: 0 }];
118    /// let mut bucket = ColorBucket::from_pixels(data.clone()).expect("Passed empty color vector to test.");
119    /// let result = bucket.make_palette(3);
120    /// ```
121    ///
122    pub fn make_palette(&mut self, iter_count: u8) -> Vec<Color> {
123        let mut result = vec![];
124        self.recurse(iter_count, &mut result);
125        result
126    }
127
128    /// Performs the median cut on a own vector (bucket) of `Color`.
129    /// Returns two `Color` vectors representing the colors above and colors below median value.
130    ///
131    fn median_cut(&mut self) -> (Option<ColorBucket>, Option<ColorBucket>) {
132        let highest_range_channel = self.highest_range_channel();
133        let median = self.color_median(highest_range_channel);
134        let mut above_median = vec![];
135        let mut below_median = vec![];
136        for color in &self.colors {
137            if color[highest_range_channel] > median {
138                above_median.push(*color);
139            } else {
140                below_median.push(*color);
141            }
142        }
143
144        (ColorBucket::from_pixels(above_median), ColorBucket::from_pixels(below_median))
145    }
146
147    /// Returns the color channel with the highest range.
148    /// IMPORTANT: Ignores alpha channel!
149    ///
150    fn highest_range_channel(&self) -> ColorChannel {
151        let ranges = self.color_ranges();
152        let mut highest_range_channel = ColorChannel::R;
153        let mut highest_value = ranges.r;
154
155        if ranges.g > highest_value {
156            highest_range_channel = ColorChannel::G;
157            highest_value = ranges.g;
158        }
159
160        if ranges.b > highest_value {
161            highest_range_channel = ColorChannel::B;
162        }
163
164        highest_range_channel
165    }
166
167    /// Returns the ranges for each color channel.
168    ///
169    fn color_ranges(&self) -> Color {
170        // Unwrap is ok here, because `max_by_key` only returns `None` for empty vectors
171        Color {
172            r: self.colors.iter().max_by_key(|c| c.r).unwrap().r - self.colors.iter().min_by_key(|c| c.r).unwrap().r,
173            g: self.colors.iter().max_by_key(|c| c.g).unwrap().g - self.colors.iter().min_by_key(|c| c.g).unwrap().g,
174            b: self.colors.iter().max_by_key(|c| c.b).unwrap().b - self.colors.iter().min_by_key(|c| c.b).unwrap().b,
175            a: self.colors.iter().max_by_key(|c| c.a).unwrap().a - self.colors.iter().min_by_key(|c| c.a).unwrap().a,
176        }
177    }
178
179    /// Sort a colors for a specific channel.
180    ///
181    /// # Arguments
182    ///
183    /// * `channel` - Target channel. The sorting is performed based on this value.
184    ///
185    fn sort_colors(&mut self, channel: ColorChannel) {
186        self.colors.sort_by_key(|x| x[channel])
187    }
188
189    /// Returns median value for a specific `ColorChannel`.
190    ///
191    /// # Arguments
192    ///
193    /// * `channel` - Target channel for which the median is calculated.
194    ///
195    fn color_median(&mut self, channel: ColorChannel) -> u8 {
196        self.sort_colors(channel);
197
198        let mid = self.colors.len() / 2;
199        if self.colors.len() % 2 == 0 {
200            let bucket = ColorBucket::from_pixels(vec![self.colors[mid - 1], self.colors[mid]]).unwrap();
201            bucket.channel_mean(channel)
202        } else {
203            self.channel_value_by_index(mid, channel)
204        }
205    }
206
207    /// Returns a color value based on the provided channel and index parameters.
208    ///
209    /// # Arguments
210    ///
211    /// * `index` - Index of the target color in the vector.
212    /// * `channel` - Color channel of the searched value.
213    ///
214    fn channel_value_by_index(&self, index: usize, channel: ColorChannel) -> u8 {
215        self.colors[index][channel]
216    }
217
218    /// Calculate the mean value for a specific color channel on own vector of `Color`.
219    ///
220    /// # Arguments
221    ///
222    /// * `channel` - Target channel for which the mean is calculated.
223    ///
224    ///
225    fn channel_mean(&self, channel: ColorChannel) -> u8 {
226        mean(self.colors.iter().map(|x| x[channel]))
227    }
228
229    /// Returns the mean color value based on own colors.
230    ///
231    fn color_mean(&self) -> Color {
232        let r = mean(self.colors.iter().map(|c| c.r));
233        let g = mean(self.colors.iter().map(|c| c.g));
234        let b = mean(self.colors.iter().map(|c| c.b));
235        let a = mean(self.colors.iter().map(|c| c.a));
236
237        Color { r, g, b, a }
238    }
239}
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244
245    #[test]
246    fn from_pixels_ut() {
247        let bucket = ColorBucket::from_pixels(vec![]);
248        assert_eq!(bucket, None);
249
250        let data = vec![Color { r: 15, g: 131, b: 0, a: 255 }, Color { r: 221, g: 11, b: 22, a: 130 }, Color { r: 81, g: 11, b: 16, a: 0 }];
251        let bucket = ColorBucket::from_pixels(data.clone()).expect("Passed empty color vector to test.");
252        assert_eq!(bucket.colors, data);
253    }
254
255    #[test]
256    fn recurse_ut() {
257        let pixels = vec![Color { r: 255, g: 0, b: 0, a: 255 }, Color { r: 0, g: 255, b: 0, a: 255 }];
258        let mut bucket = ColorBucket::from_pixels(pixels.clone()).expect("Passed empty color vector to test.");
259        let mut result = vec![];
260        bucket.recurse(1, &mut result);
261        assert_eq!(result, pixels);
262    }
263
264    #[test]
265    fn make_palette_ut() {
266        let pixels = vec![Color { r: 100, g: 120, b: 120, a: 0 }, Color { r: 150, g: 150, b: 150, a: 0 }, Color { r: 255, g: 255, b: 255, a: 0 }];
267        let mut bucket = ColorBucket::from_pixels(pixels.clone()).expect("Passed empty color vector to test.");
268
269        let colors = bucket.make_palette(3);
270        let expected = vec![Color { r: 255, g: 255, b: 255, a: 0 }, Color { r: 150, g: 150, b: 150, a: 0 }, Color { r: 100, g: 120, b: 120, a: 0 }];
271        assert_eq!(colors, expected);
272    }
273
274    #[test]
275    pub fn sort_colors_ut() {
276        let colors = generate_unsorted_colors();
277        let mut bucket = ColorBucket::from_pixels(colors.clone()).expect("Passed empty color vector to test");
278        bucket.sort_colors(ColorChannel::R);
279
280        assert_eq!(bucket.colors[0], Color { r: 0, g: 2, b: 1, a: 20 });
281        assert_eq!(bucket.colors[1], Color { r: 1, g: 23, b: 16, a: 20 });
282        assert_eq!(bucket.colors[2], Color { r: 3, g: 4, b: 15, a: 2 });
283        assert_eq!(bucket.colors[3], Color { r: 55, g: 17, b: 0, a: 118 });
284    }
285
286    #[test]
287    pub fn color_median_ut() {
288        let colors = generate_unsorted_colors();
289        let mut bucket = ColorBucket::from_pixels(colors.clone()).expect("Passed empty color vector to test");
290        let result = bucket.color_median(ColorChannel::R);
291        assert_eq!(result, 2);
292    }
293
294    #[test]
295    fn channel_value_by_index_ut() {
296        let colors = vec![
297            Color { r: 100, g: 22, b: 12, a: 0 },
298            Color { r: 126, g: 175, b: 137, a: 1 },
299            Color { r: 221, g: 225, b: 0, a: 113 },
300            Color { r: 13, g: 226, b: 0, a: 17 },
301        ];
302
303        let bucket = ColorBucket::from_pixels(colors).expect("Passing empty color vector to test");
304
305        assert_eq!(100, bucket.channel_value_by_index(0, ColorChannel::R));
306        assert_eq!(22, bucket.channel_value_by_index(0, ColorChannel::G));
307        assert_eq!(12, bucket.channel_value_by_index(0, ColorChannel::B));
308        assert_eq!(0, bucket.channel_value_by_index(0, ColorChannel::A));
309
310        assert_eq!(126, bucket.channel_value_by_index(1, ColorChannel::R));
311        assert_eq!(175, bucket.channel_value_by_index(1, ColorChannel::G));
312        assert_eq!(137, bucket.channel_value_by_index(1, ColorChannel::B));
313        assert_eq!(1, bucket.channel_value_by_index(1, ColorChannel::A));
314
315        assert_eq!(221, bucket.channel_value_by_index(2, ColorChannel::R));
316        assert_eq!(225, bucket.channel_value_by_index(2, ColorChannel::G));
317        assert_eq!(0, bucket.channel_value_by_index(2, ColorChannel::B));
318        assert_eq!(113, bucket.channel_value_by_index(2, ColorChannel::A));
319
320        assert_eq!(13, bucket.channel_value_by_index(3, ColorChannel::R));
321        assert_eq!(226, bucket.channel_value_by_index(3, ColorChannel::G));
322        assert_eq!(0, bucket.channel_value_by_index(3, ColorChannel::B));
323        assert_eq!(17, bucket.channel_value_by_index(3, ColorChannel::A));
324    }
325
326    #[test]
327    fn channel_mean_ut() {
328        let colors = vec![
329            Color { r: 100, g: 50, b: 12, a: 255 },
330            Color { r: 100, g: 50, b: 12, a: 255 },
331            Color { r: 100, g: 50, b: 12, a: 255 },
332            Color { r: 100, g: 50, b: 12, a: 255 },
333        ];
334
335        let bucket = ColorBucket::from_pixels(colors).expect("Passed empty color vector to test.");
336        let mut result = bucket.channel_mean(ColorChannel::R);
337        assert_eq!(100, result);
338        result = bucket.channel_mean(ColorChannel::G);
339        assert_eq!(50, result);
340        result = bucket.channel_mean(ColorChannel::B);
341        assert_eq!(12, result);
342        result = bucket.channel_mean(ColorChannel::A);
343        assert_eq!(255, result);
344
345        // More precise check
346        let colors = vec![
347            Color { r: 100, g: 22, b: 12, a: 0 },
348            Color { r: 126, g: 175, b: 137, a: 1 },
349            Color { r: 221, g: 225, b: 0, a: 113 },
350            Color { r: 13, g: 226, b: 0, a: 17 },
351        ];
352
353        let bucket = ColorBucket::from_pixels(colors).expect("Passed empty color vector to test.");
354
355        result = bucket.channel_mean(ColorChannel::R);
356        assert_eq!(115, result);
357        result = bucket.channel_mean(ColorChannel::G);
358        assert_eq!(162, result);
359        result = bucket.channel_mean(ColorChannel::B);
360        assert_eq!(37, result);
361        result = bucket.channel_mean(ColorChannel::A);
362        assert_eq!(32, result);
363    }
364
365    #[test]
366    pub fn ut_color_mean() {
367        let colors = generate_unsorted_colors();
368        let bucket = ColorBucket::from_pixels(colors).expect("Passed empty color vector to test.");
369
370        let result = bucket.color_mean();
371        let expected = Color { r: 14, g: 11, b: 8, a: 40 };
372        assert_eq!(expected, result);
373    }
374
375    #[test]
376    fn median_cut_ut() {
377        let mut bucket = ColorBucket::from_pixels(generate_unsorted_colors()).expect("Passed empty color vector to test.");
378        let result = bucket.median_cut();
379        assert_eq!(
380            result.0,
381            Some(ColorBucket::from_pixels(vec![Color { r: 3, g: 4, b: 15, a: 2 }, Color { r: 55, g: 17, b: 0, a: 118 }]).unwrap())
382        );
383        assert_eq!(
384            result.1,
385            Some(ColorBucket::from_pixels(vec![Color { r: 0, g: 2, b: 1, a: 20 }, Color { r: 1, g: 23, b: 16, a: 20 }]).unwrap())
386        );
387
388        let mut bucket = ColorBucket::from_pixels(vec![Color { r: 0, g: 0, b: 0, a: 0 }]).expect("Passed empty color vector to test.");
389        let result = bucket.median_cut();
390        assert_eq!(result.0, None);
391        assert_eq!(result.1, Some(ColorBucket::from_pixels(vec![Color { r: 0, g: 0, b: 0, a: 0 }])).unwrap());
392    }
393
394    #[test]
395    fn highest_range_channel_ut() {
396        let bucket = ColorBucket::from_pixels(generate_unsorted_colors()).expect("Passed empty color vector to test");
397        assert_eq!(ColorChannel::R, bucket.highest_range_channel());
398        assert_ne!(ColorChannel::G, bucket.highest_range_channel());
399        assert_ne!(ColorChannel::B, bucket.highest_range_channel());
400        assert_ne!(ColorChannel::A, bucket.highest_range_channel());
401    }
402
403    #[test]
404    fn color_ranges_ut() {
405        let bucket = ColorBucket::from_pixels(generate_unsorted_colors()).expect("Passed empty color vector to test");
406        let expected = Color { r: 55, g: 21, b: 16, a: 116 };
407        assert_eq!(expected, bucket.color_ranges());
408    }
409
410    fn generate_unsorted_colors() -> Vec<Color> {
411        vec![
412            Color { r: 55, g: 17, b: 0, a: 118 },
413            Color { r: 0, g: 2, b: 1, a: 20 },
414            Color { r: 3, g: 4, b: 15, a: 2 },
415            Color { r: 1, g: 23, b: 16, a: 20 },
416        ]
417    }
418}
iris_lib/color_bucket.rs

iris_lib/
color_bucket.rs
