material_colors/quantize/
quantizer_wu.rs

1#![allow(clippy::too_many_arguments)]
2
3use super::{Quantizer, QuantizerMap, QuantizerResult};
4#[cfg(all(not(feature = "std"), feature = "libm"))]
5#[allow(unused_imports)]
6use crate::utils::no_std::FloatExt;
7use crate::{
8    color::{Argb, Rgb},
9    IndexMap,
10};
11#[cfg(not(feature = "std"))]
12use alloc::{vec, vec::Vec};
13use core::fmt;
14#[cfg(feature = "std")]
15use std::{vec, vec::Vec};
16
17// A histogram of all the input colors is constructed. It has the shape of a
18//  The cube would be too large if it contained all 16 million colors:
19// historical best practice is to use 5 bits  of the 8 in each channel,
20// reducing the histogram to a volume of ~32,000.
21const INDEX_BITS: u8 = 5;
22const BITS_TO_REMOVE: u8 = 8 - INDEX_BITS;
23const SIDE_LENGTH: usize = (1 << INDEX_BITS) + 1;
24const TOTAL_SIZE: usize = SIDE_LENGTH.pow(3);
25
26pub struct QuantizerWu {
27    weights: Vec<i64>,
28    moments_r: Vec<i64>,
29    moments_g: Vec<i64>,
30    moments_b: Vec<i64>,
31    moments: Vec<f64>,
32    cubes: Vec<Cube>,
33}
34
35impl QuantizerWu {
36    fn new(max_colors: usize) -> Self {
37        Self {
38            weights: vec![0; TOTAL_SIZE],
39            moments_r: vec![0; TOTAL_SIZE],
40            moments_g: vec![0; TOTAL_SIZE],
41            moments_b: vec![0; TOTAL_SIZE],
42            moments: vec![0.0; TOTAL_SIZE],
43            cubes: vec![
44                Cube {
45                    pixels: [Rgb::default(), Rgb::default()],
46                    vol: 0
47                };
48                max_colors
49            ],
50        }
51    }
52}
53
54impl Quantizer for QuantizerWu {
55    fn quantize(pixels: &[Argb], max_colors: usize) -> QuantizerResult {
56        let mut result = QuantizerMap::quantize(pixels, max_colors);
57
58        result.color_to_count.sort_by(|_, a, _, b| a.cmp(b));
59
60        let mut quantizer = Self::new(max_colors);
61
62        quantizer.construct_histogram(result.color_to_count);
63        quantizer.compute_moments();
64
65        let create_boxes_result = quantizer.create_boxes(max_colors);
66        let color_to_count = quantizer.create_result(create_boxes_result.result_count);
67
68        QuantizerResult {
69            color_to_count,
70            input_pixel_to_cluster_pixel: IndexMap::default(),
71        }
72    }
73}
74
75impl QuantizerWu {
76    pub fn get_index<T: Into<usize>>(r: T, g: T, b: T) -> usize {
77        let r: usize = r.into();
78        let g: usize = g.into();
79        let b: usize = b.into();
80
81        (r << (INDEX_BITS * 2)) + (r << (INDEX_BITS + 1)) + (g << INDEX_BITS) + r + g + b
82    }
83
84    pub fn construct_histogram(&mut self, pixels: IndexMap<Argb, u32>) {
85        for (pixel, count) in pixels {
86            let red = pixel.red;
87            let green = pixel.green;
88            let blue = pixel.blue;
89
90            let i_r = (red >> BITS_TO_REMOVE) + 1;
91            let i_g = (green >> BITS_TO_REMOVE) + 1;
92            let i_b = (blue >> BITS_TO_REMOVE) + 1;
93
94            let index = Self::get_index(i_r, i_g, i_b);
95
96            self.weights[index] += i64::from(count);
97
98            self.moments_r[index] += i64::from(red) * i64::from(count);
99            self.moments_g[index] += i64::from(green) * i64::from(count);
100            self.moments_b[index] += i64::from(blue) * i64::from(count);
101
102            self.moments[index] += f64::from(count)
103                * f64::from(blue).mul_add(
104                    f64::from(blue),
105                    f64::from(red).mul_add(f64::from(red), f64::from(green) * f64::from(green)),
106                );
107        }
108    }
109
110    pub fn compute_moments(&mut self) {
111        for r in 1..SIDE_LENGTH {
112            let mut area = [0; SIDE_LENGTH];
113            let mut area_r = [0; SIDE_LENGTH];
114            let mut area_g = [0; SIDE_LENGTH];
115            let mut area_b = [0; SIDE_LENGTH];
116            let mut area2 = [0.0; SIDE_LENGTH];
117
118            for g in 1..SIDE_LENGTH {
119                let mut line = 0;
120                let mut line_r = 0;
121                let mut line_g = 0;
122                let mut line_b = 0;
123                let mut line2 = 0.0;
124
125                for b in 1..SIDE_LENGTH {
126                    let index = Self::get_index(r, g, b);
127
128                    line += self.weights[index];
129                    line_r += self.moments_r[index];
130                    line_g += self.moments_g[index];
131                    line_b += self.moments_b[index];
132                    line2 += self.moments[index];
133
134                    area[b] += line;
135                    area_r[b] += line_r;
136                    area_g[b] += line_g;
137                    area_b[b] += line_b;
138                    area2[b] += line2;
139
140                    let previous_index = Self::get_index(r - 1, g, b);
141
142                    self.weights[index] = self.weights[previous_index] + area[b];
143                    self.moments_r[index] = self.moments_r[previous_index] + area_r[b];
144                    self.moments_g[index] = self.moments_g[previous_index] + area_g[b];
145                    self.moments_b[index] = self.moments_b[previous_index] + area_b[b];
146                    self.moments[index] = self.moments[previous_index] + area2[b];
147                }
148            }
149        }
150    }
151
152    pub fn create_boxes(&mut self, max_color_count: usize) -> CreateBoxesResult {
153        self.cubes[0] = Cube {
154            pixels: [
155                Rgb::default(),
156                Rgb::new(
157                    SIDE_LENGTH as u8 - 1,
158                    SIDE_LENGTH as u8 - 1,
159                    SIDE_LENGTH as u8 - 1,
160                ),
161            ],
162            vol: 0,
163        };
164
165        let mut volume_variance = vec![0.0; max_color_count];
166        let mut next = 0;
167        let mut generated_color_count = max_color_count;
168        let mut i = 1;
169
170        while i < max_color_count {
171            if self.cut(next, i) {
172                volume_variance[next] = if self.cubes[next].vol > 1 {
173                    self.variance(&self.cubes[next])
174                } else {
175                    0.0
176                };
177
178                volume_variance[i] = if self.cubes[i].vol > 1 {
179                    self.variance(&self.cubes[i])
180                } else {
181                    0.0
182                };
183            } else {
184                volume_variance[next] = 0.0;
185
186                i -= 1;
187            }
188
189            next = 0;
190
191            let mut temp = volume_variance[0];
192
193            let mut j = 1;
194
195            while j <= i {
196                if volume_variance[j] > temp {
197                    temp = volume_variance[j];
198                    next = j;
199                }
200
201                j += 1;
202            }
203
204            if temp <= 0.0 {
205                generated_color_count = i + 1;
206
207                break;
208            }
209
210            i += 1;
211        }
212
213        CreateBoxesResult {
214            requested_count: max_color_count,
215            result_count: generated_color_count,
216        }
217    }
218
219    pub fn create_result(&self, color_count: usize) -> IndexMap<Argb, u32> {
220        let mut result = IndexMap::default();
221
222        for i in 0..color_count {
223            let cube = &self.cubes[i];
224            let weight = Self::volume(cube, &self.weights);
225
226            if weight > 0 {
227                let r = ((Self::volume(cube, &self.moments_r)) / weight) as u8;
228                let g = ((Self::volume(cube, &self.moments_g)) / weight) as u8;
229                let b = ((Self::volume(cube, &self.moments_b)) / weight) as u8;
230
231                let color = Rgb::new(r, g, b).into();
232
233                result.insert(color, 0);
234            }
235        }
236
237        result
238    }
239
240    pub fn variance(&self, cube: &Cube) -> f64 {
241        let dr = Self::volume(cube, &self.moments_r) as f64;
242        let dg = Self::volume(cube, &self.moments_g) as f64;
243        let db = Self::volume(cube, &self.moments_b) as f64;
244
245        let xx = self.moments[Self::get_index::<u8>(cube.r(1), cube.g(1), cube.b(1))]
246            - self.moments[Self::get_index::<u8>(cube.r(1), cube.g(1), cube.b(0))]
247            - self.moments[Self::get_index::<u8>(cube.r(1), cube.g(0), cube.b(1))]
248            + self.moments[Self::get_index::<u8>(cube.r(1), cube.g(0), cube.b(0))]
249            - self.moments[Self::get_index::<u8>(cube.r(0), cube.g(1), cube.b(1))]
250            + self.moments[Self::get_index::<u8>(cube.r(0), cube.g(1), cube.b(0))]
251            + self.moments[Self::get_index::<u8>(cube.r(0), cube.g(0), cube.b(1))]
252            - self.moments[Self::get_index::<u8>(cube.r(0), cube.g(0), cube.b(0))];
253
254        let hypotenuse = db.mul_add(db, dr.mul_add(dr, dg * dg));
255        let volume = Self::volume(cube, &self.weights) as f64;
256
257        xx - (hypotenuse / volume)
258    }
259
260    pub fn cut(&mut self, next: usize, i: usize) -> bool {
261        let (mut one, mut two) = (self.cubes[next].clone(), self.cubes[i].clone());
262
263        let whole_r = Self::volume(&one, &self.moments_r);
264        let whole_g = Self::volume(&one, &self.moments_g);
265        let whole_b = Self::volume(&one, &self.moments_b);
266        let whole_w = Self::volume(&one, &self.weights);
267
268        let max_rresult = self.maximize(
269            &one,
270            &Direction::Red,
271            one.r::<i32>(0) + 1,
272            one.r::<i32>(1),
273            whole_r,
274            whole_g,
275            whole_b,
276            whole_w,
277        );
278        let max_gresult = self.maximize(
279            &one,
280            &Direction::Green,
281            one.g::<i32>(0) + 1,
282            one.g::<i32>(1),
283            whole_r,
284            whole_g,
285            whole_b,
286            whole_w,
287        );
288        let max_bresult = self.maximize(
289            &one,
290            &Direction::Blue,
291            one.b::<i32>(0) + 1,
292            one.b::<i32>(1),
293            whole_r,
294            whole_g,
295            whole_b,
296            whole_w,
297        );
298
299        let cut_direction: Direction;
300
301        let max_r = max_rresult.maximum;
302        let max_g = max_gresult.maximum;
303        let max_b = max_bresult.maximum;
304
305        if max_r >= max_g && max_r >= max_b {
306            cut_direction = Direction::Red;
307
308            if max_rresult.cut_location < 0 {
309                return false;
310            }
311        } else if max_g >= max_r && max_g >= max_b {
312            cut_direction = Direction::Green;
313        } else {
314            cut_direction = Direction::Blue;
315        }
316
317        two.pixels[1].red = one.pixels[1].red;
318        two.pixels[1].green = one.pixels[1].green;
319        two.pixels[1].blue = one.pixels[1].blue;
320
321        match cut_direction {
322            Direction::Red => {
323                one.pixels[1].red = max_rresult.cut_location as u8;
324                two.pixels[0].red = one.pixels[1].red;
325                two.pixels[0].green = one.pixels[0].green;
326                two.pixels[0].blue = one.pixels[0].blue;
327            }
328            Direction::Green => {
329                one.pixels[1].green = max_gresult.cut_location as u8;
330                two.pixels[0].red = one.pixels[0].red;
331                two.pixels[0].green = one.pixels[1].green;
332                two.pixels[0].blue = one.pixels[0].blue;
333            }
334            Direction::Blue => {
335                one.pixels[1].blue = max_bresult.cut_location as u8;
336                two.pixels[0].red = one.pixels[0].red;
337                two.pixels[0].green = one.pixels[0].green;
338                two.pixels[0].blue = one.pixels[1].blue;
339            }
340        }
341
342        one.vol = (one.r::<i32>(1) - one.r::<i32>(0))
343            * (one.g::<i32>(1) - one.g::<i32>(0))
344            * (one.b::<i32>(1) - one.b::<i32>(0));
345        two.vol = (two.r::<i32>(1) - two.r::<i32>(0))
346            * (two.g::<i32>(1) - two.g::<i32>(0))
347            * (two.b::<i32>(1) - two.b::<i32>(0));
348
349        self.cubes[next] = one;
350        self.cubes[i] = two;
351
352        true
353    }
354
355    pub fn maximize(
356        &self,
357        cube: &Cube,
358        direction: &Direction,
359        first: i32,
360        last: i32,
361        whole_r: i64,
362        whole_g: i64,
363        whole_b: i64,
364        whole_w: i64,
365    ) -> MaximizeResult {
366        let bottom_r = Self::bottom(cube, direction, &self.moments_r) as f64;
367        let bottom_g = Self::bottom(cube, direction, &self.moments_g) as f64;
368        let bottom_b = Self::bottom(cube, direction, &self.moments_b) as f64;
369        let bottom_w = Self::bottom(cube, direction, &self.weights) as f64;
370
371        let mut max = 0.0;
372        let mut cut = -1;
373
374        for i in first..last {
375            let mut half_r = bottom_r + Self::top(cube, direction, i, &self.moments_r) as f64;
376            let mut half_g = bottom_g + Self::top(cube, direction, i, &self.moments_g) as f64;
377            let mut half_b = bottom_b + Self::top(cube, direction, i, &self.moments_b) as f64;
378            let mut half_w = bottom_w + Self::top(cube, direction, i, &self.weights) as f64;
379
380            if half_w == 0.0 {
381                continue;
382            }
383
384            let mut temp_numerator = half_b.mul_add(half_b, half_r.mul_add(half_r, half_g.powi(2)));
385            let mut temp_denominator = half_w;
386            let mut temp = temp_numerator / temp_denominator;
387
388            half_r = whole_r as f64 - half_r;
389            half_g = whole_g as f64 - half_g;
390            half_b = whole_b as f64 - half_b;
391            half_w = whole_w as f64 - half_w;
392
393            if half_w == 0.0 {
394                continue;
395            }
396
397            temp_numerator = half_b.mul_add(half_b, half_r.mul_add(half_r, half_g.powi(2)));
398            temp_denominator = half_w;
399            temp += temp_numerator / temp_denominator;
400
401            if temp > max {
402                max = temp;
403                cut = i;
404            }
405        }
406
407        MaximizeResult {
408            cut_location: cut,
409            maximum: max,
410        }
411    }
412
413    pub fn volume(cube: &Cube, moment: &[i64]) -> i64 {
414        moment[Self::get_index::<u8>(cube.r(1), cube.g(1), cube.b(1))]
415            - moment[Self::get_index::<u8>(cube.r(1), cube.g(1), cube.b(0))]
416            - moment[Self::get_index::<u8>(cube.r(1), cube.g(0), cube.b(1))]
417            + moment[Self::get_index::<u8>(cube.r(1), cube.g(0), cube.b(0))]
418            - moment[Self::get_index::<u8>(cube.r(0), cube.g(1), cube.b(1))]
419            + moment[Self::get_index::<u8>(cube.r(0), cube.g(1), cube.b(0))]
420            + moment[Self::get_index::<u8>(cube.r(0), cube.g(0), cube.b(1))]
421            - moment[Self::get_index::<u8>(cube.r(0), cube.g(0), cube.b(0))]
422    }
423
424    pub fn bottom(cube: &Cube, direction: &Direction, moment: &[i64]) -> i64 {
425        match direction {
426            Direction::Red => {
427                -moment[Self::get_index::<u8>(cube.r(0), cube.g(1), cube.b(1))]
428                    + moment[Self::get_index::<u8>(cube.r(0), cube.g(1), cube.b(0))]
429                    + moment[Self::get_index::<u8>(cube.r(0), cube.g(0), cube.b(1))]
430                    - moment[Self::get_index::<u8>(cube.r(0), cube.g(0), cube.b(0))]
431            }
432            Direction::Green => {
433                -moment[Self::get_index::<u8>(cube.r(1), cube.g(0), cube.b(1))]
434                    + moment[Self::get_index::<u8>(cube.r(1), cube.g(0), cube.b(0))]
435                    + moment[Self::get_index::<u8>(cube.r(0), cube.g(0), cube.b(1))]
436                    - moment[Self::get_index::<u8>(cube.r(0), cube.g(0), cube.b(0))]
437            }
438            Direction::Blue => {
439                -moment[Self::get_index::<u8>(cube.r(1), cube.g(1), cube.b(0))]
440                    + moment[Self::get_index::<u8>(cube.r(1), cube.g(0), cube.b(0))]
441                    + moment[Self::get_index::<u8>(cube.r(0), cube.g(1), cube.b(0))]
442                    - moment[Self::get_index::<u8>(cube.r(0), cube.g(0), cube.b(0))]
443            }
444        }
445    }
446
447    pub fn top(cube: &Cube, direction: &Direction, position: i32, moment: &[i64]) -> i64 {
448        match direction {
449            Direction::Red => {
450                moment[Self::get_index(position as usize, cube.g(1), cube.b(1))]
451                    - moment[Self::get_index(position as usize, cube.g(1), cube.b(0))]
452                    - moment[Self::get_index(position as usize, cube.g(0), cube.b(1))]
453                    + moment[Self::get_index(position as usize, cube.g(0), cube.b(0))]
454            }
455            Direction::Green => {
456                moment[Self::get_index(cube.r(1), position as usize, cube.b(1))]
457                    - moment[Self::get_index(cube.r(1), position as usize, cube.b(0))]
458                    - moment[Self::get_index(cube.r(0), position as usize, cube.b(1))]
459                    + moment[Self::get_index(cube.r(0), position as usize, cube.b(0))]
460            }
461            Direction::Blue => {
462                moment[Self::get_index(cube.r(1), cube.g(1), position as usize)]
463                    - moment[Self::get_index(cube.r(1), cube.g(0), position as usize)]
464                    - moment[Self::get_index(cube.r(0), cube.g(1), position as usize)]
465                    + moment[Self::get_index(cube.r(0), cube.g(0), position as usize)]
466            }
467        }
468    }
469}
470
471pub enum Direction {
472    Red,
473    Green,
474    Blue,
475}
476
477#[derive(Debug)]
478pub struct MaximizeResult {
479    // < 0 if cut impossible
480    pub cut_location: i32,
481    pub maximum: f64,
482}
483
484pub struct CreateBoxesResult {
485    pub requested_count: usize,
486    pub result_count: usize,
487}
488
489#[derive(Debug, Clone)]
490pub struct Cube {
491    pub pixels: [Rgb; 2],
492    pub vol: i32,
493}
494
495impl Cube {
496    pub fn r<T: From<u8>>(&self, pixel: usize) -> T {
497        self.pixels[pixel].red.into()
498    }
499
500    pub fn g<T: From<u8>>(&self, pixel: usize) -> T {
501        self.pixels[pixel].green.into()
502    }
503
504    pub fn b<T: From<u8>>(&self, pixel: usize) -> T {
505        self.pixels[pixel].blue.into()
506    }
507}
508
509impl fmt::Display for Cube {
510    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
511        write!(
512            f,
513            "Box: R {} -> {} G {} -> {} B {} -> {} VOL = {}",
514            self.pixels[0].red,
515            self.pixels[1].red,
516            self.pixels[0].green,
517            self.pixels[1].green,
518            self.pixels[0].blue,
519            self.pixels[1].blue,
520            self.vol
521        )
522    }
523}
524
525#[cfg(test)]
526mod tests {
527    use super::{Quantizer, QuantizerWu};
528    use crate::color::Argb;
529    #[cfg(not(feature = "std"))]
530    use alloc::vec::Vec;
531    #[cfg(feature = "std")]
532    use std::vec::Vec;
533
534    const RED: Argb = Argb::from_u32(0xffff0000);
535    const GREEN: Argb = Argb::from_u32(0xff00ff00);
536    const BLUE: Argb = Argb::from_u32(0xff0000ff);
537    // const WHITE: Argb = Argb::from_u32(0xffffffff);
538    // const RANDOM: Argb = Argb::from_u32(0xff426088);
539    const MAX_COLORS: usize = 256;
540
541    #[test]
542    fn test_1rando() {
543        let result = QuantizerWu::quantize(&[Argb::from_u32(0xff14_1216)], MAX_COLORS);
544        let colors = result.color_to_count.keys().collect::<Vec<_>>();
545
546        assert_eq!(colors.len(), 1);
547        assert_eq!(colors[0], &Argb::from_u32(0xff14_1216));
548    }
549
550    #[test]
551    fn test_1r() {
552        let result = QuantizerWu::quantize(&[RED], MAX_COLORS);
553        let colors = result.color_to_count.keys().collect::<Vec<_>>();
554
555        assert_eq!(colors.len(), 1);
556        assert_eq!(colors[0], &RED);
557    }
558
559    #[test]
560    fn test_1g() {
561        let result = QuantizerWu::quantize(&[GREEN], MAX_COLORS);
562        let colors = result.color_to_count.keys().collect::<Vec<_>>();
563
564        assert_eq!(colors.len(), 1);
565        assert_eq!(colors[0], &GREEN);
566    }
567
568    #[test]
569    fn test_1b() {
570        let result = QuantizerWu::quantize(&[BLUE], MAX_COLORS);
571        let colors = result.color_to_count.keys().collect::<Vec<_>>();
572
573        assert_eq!(colors.len(), 1);
574        assert_eq!(colors[0], &BLUE);
575    }
576
577    #[test]
578    fn test_5b() {
579        let result = QuantizerWu::quantize(&[BLUE, BLUE, BLUE, BLUE, BLUE], MAX_COLORS);
580        let colors = result.color_to_count.keys().collect::<Vec<_>>();
581
582        assert_eq!(colors.len(), 1);
583        assert_eq!(colors[0], &BLUE);
584    }
585
586    #[test]
587    fn test_2r_3g() {
588        let result = QuantizerWu::quantize(&[RED, RED, GREEN, GREEN, GREEN], MAX_COLORS);
589
590        assert_eq!(result.color_to_count.keys().len(), 2);
591
592        assert!(result.color_to_count.contains_key(&GREEN));
593        assert!(result.color_to_count.contains_key(&GREEN));
594    }
595
596    #[test]
597    fn test_1r_1g_1b() {
598        let result = QuantizerWu::quantize(&[RED, GREEN, BLUE], MAX_COLORS);
599
600        assert_eq!(result.color_to_count.keys().len(), 3);
601
602        assert!(result.color_to_count.contains_key(&GREEN));
603        assert!(result.color_to_count.contains_key(&RED));
604        assert!(result.color_to_count.contains_key(&BLUE));
605    }
606}
material_colors/quantize/quantizer_wu.rs

material_colors/quantize/
quantizer_wu.rs
