a_sixel/
adu.rs

1//! Use Adaptive Distributive Units to "learn" the image's color properties and
2//! select palette entries.
3//!
4//! See <httpe://faculty.uca.edu/ecelebi/documents/ISJ_2014.pdf> for the
5//! original paper on this algorithm.
6//!
7//! The parameters from the paper for a 256 color palette are:
8//! - THETA = (400 * 256^0.5) = 6400
9//! - STEPS = (2 * 256 - 3) * THETA = 3257600
10//! - GAMMA = 0.015 or GAMMA_DIV ~= 64
11//!
12//! as is specified by the default arguments to this struct.
13//!
14//! See the code for the type aliases (e.g. [`ADUSixelEncoder16`]) for more
15//! default parameters.
16
17use std::collections::HashSet;
18
19use image::RgbImage;
20use kiddo::{
21    float::kdtree::KdTree,
22    SquaredEuclidean,
23};
24use ordered_float::OrderedFloat;
25use palette::Lab;
26use rayon::iter::{
27    IntoParallelRefIterator,
28    ParallelIterator,
29};
30use sobol_burley::sample_4d;
31
32use crate::{
33    dither::Sierra,
34    private,
35    rgb_to_lab,
36    PaletteBuilder,
37    SixelEncoder,
38};
39
40pub type ADUSixelEncoderMono<D = Sierra> = SixelEncoder<ADUPaletteBuilder<2, 400, 400>, D>;
41pub type ADUSixelEncoder4<D = Sierra> = SixelEncoder<ADUPaletteBuilder<4, 800, 4000>, D>;
42pub type ADUSixelEncoder8<D = Sierra> = SixelEncoder<ADUPaletteBuilder<8, 1131, 14703>, D>;
43pub type ADUSixelEncoder16<D = Sierra> = SixelEncoder<ADUPaletteBuilder<16, 1600, 46400>, D>;
44pub type ADUSixelEncoder32<D = Sierra> = SixelEncoder<ADUPaletteBuilder<32, 2262, 137982>, D>;
45pub type ADUSixelEncoder64<D = Sierra> = SixelEncoder<ADUPaletteBuilder<64, 3200, 400000>, D>;
46pub type ADUSixelEncoder128<D = Sierra> = SixelEncoder<ADUPaletteBuilder<128, 4525, 1144947>, D>;
47pub type ADUSixelEncoder256<D = Sierra> = SixelEncoder<ADUPaletteBuilder<256>, D>;
48
49pub struct ADUPaletteBuilder<
50    const PALETTE_SIZE: usize = 256,
51    const THETA: usize = 6400,
52    const STEPS: usize = 3257600,
53    const GAMMA_DIV: usize = 64,
54>;
55
56impl<const PALETTE_SIZE: usize, const THETA: usize, const STEPS: usize, const GAMMA_DIV: usize>
57    private::Sealed for ADUPaletteBuilder<PALETTE_SIZE, THETA, STEPS, GAMMA_DIV>
58{
59}
60impl<const PALETTE_SIZE: usize, const THETA: usize, const STEPS: usize, const GAMMA_DIV: usize>
61    PaletteBuilder for ADUPaletteBuilder<PALETTE_SIZE, THETA, STEPS, GAMMA_DIV>
62{
63    const NAME: &'static str = "ADU";
64    const PALETTE_SIZE: usize = PALETTE_SIZE;
65
66    fn build_palette(image: &RgbImage) -> Vec<Lab> {
67        let gamma: f32 = 1.0 / (GAMMA_DIV as f32);
68
69        let candidates = image.pixels().copied().map(rgb_to_lab).collect::<Vec<_>>();
70
71        let centroid = candidates.par_iter().copied().reduce(
72            || <Lab>::new(0.0, 0.0, 0.0),
73            |mut acc, color| {
74                acc.l += color.l;
75                acc.a += color.a;
76                acc.b += color.b;
77                acc
78            },
79        ) / candidates.len() as f32;
80
81        let mut palette = [centroid; PALETTE_SIZE];
82
83        let mut tree = KdTree::<_, _, 3, PALETTE_SIZE, u32>::with_capacity(PALETTE_SIZE);
84        tree.add(&[palette[0].l, palette[0].a, palette[0].b], 0);
85
86        let mut next_idx = 1;
87
88        let mut wc = [0; PALETTE_SIZE];
89
90        let candidates = (0..STEPS as u32 / 4)
91            .flat_map(|idx| {
92                let [x, y, z, w] = sample_4d(idx % (1 << 16), 0, idx / (1 << 16));
93                [
94                    candidates[(x * candidates.len() as f32) as usize],
95                    candidates[(y * candidates.len() as f32) as usize],
96                    candidates[(z * candidates.len() as f32) as usize],
97                    candidates[(w * candidates.len() as f32) as usize],
98                ]
99            })
100            .collect::<Vec<_>>();
101
102        for candidate in candidates {
103            let winner =
104                tree.nearest_one::<SquaredEuclidean>(&[candidate.l, candidate.a, candidate.b]);
105
106            tree.remove(
107                &[
108                    palette[winner.item].l,
109                    palette[winner.item].a,
110                    palette[winner.item].b,
111                ],
112                winner.item,
113            );
114
115            palette[winner.item].l += (candidate.l - palette[winner.item].l) * gamma;
116            palette[winner.item].a += (candidate.a - palette[winner.item].a) * gamma;
117            palette[winner.item].b += (candidate.b - palette[winner.item].b) * gamma;
118
119            tree.add(
120                &[
121                    palette[winner.item].l,
122                    palette[winner.item].a,
123                    palette[winner.item].b,
124                ],
125                winner.item,
126            );
127
128            wc[winner.item] += 1;
129
130            if wc[winner.item] >= THETA && next_idx < PALETTE_SIZE {
131                tree.add(&[candidate.l, candidate.a, candidate.b], next_idx);
132
133                wc[winner.item] = 0;
134                wc[next_idx] = 0;
135                next_idx += 1;
136            }
137        }
138
139        palette
140            .into_iter()
141            .map(|lab| {
142                [
143                    OrderedFloat(lab.l),
144                    OrderedFloat(lab.a),
145                    OrderedFloat(lab.b),
146                ]
147            })
148            .collect::<HashSet<_>>()
149            .into_iter()
150            .map(|[l, a, b]| Lab::new(*l, *a, *b))
151            .collect::<Vec<_>>()
152    }
153}
a_sixel/adu.rs

a_sixel/
adu.rs
