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 PALETTE_SIZE: usize = PALETTE_SIZE;
64
65    fn build_palette(image: &RgbImage) -> Vec<Lab> {
66        let gamma: f32 = 1.0 / (GAMMA_DIV as f32);
67
68        let candidates = image.pixels().copied().map(rgb_to_lab).collect::<Vec<_>>();
69
70        let centroid = candidates.par_iter().copied().reduce(
71            || <Lab>::new(0.0, 0.0, 0.0),
72            |mut acc, color| {
73                acc.l += color.l;
74                acc.a += color.a;
75                acc.b += color.b;
76                acc
77            },
78        ) / candidates.len() as f32;
79
80        let mut palette = [centroid; PALETTE_SIZE];
81
82        let mut tree = KdTree::<_, _, 3, PALETTE_SIZE, u32>::with_capacity(PALETTE_SIZE);
83        tree.add(&[palette[0].l, palette[0].a, palette[0].b], 0);
84
85        let mut next_idx = 1;
86
87        let mut wc = [0; PALETTE_SIZE];
88
89        let candidates = (0..STEPS as u32 / 4)
90            .flat_map(|idx| {
91                let [x, y, z, w] = sample_4d(idx % (1 << 16), 0, idx / (1 << 16));
92                [
93                    candidates[(x * candidates.len() as f32) as usize],
94                    candidates[(y * candidates.len() as f32) as usize],
95                    candidates[(z * candidates.len() as f32) as usize],
96                    candidates[(w * candidates.len() as f32) as usize],
97                ]
98            })
99            .collect::<Vec<_>>();
100
101        for candidate in candidates {
102            let winner =
103                tree.nearest_one::<SquaredEuclidean>(&[candidate.l, candidate.a, candidate.b]);
104
105            tree.remove(
106                &[
107                    palette[winner.item].l,
108                    palette[winner.item].a,
109                    palette[winner.item].b,
110                ],
111                winner.item,
112            );
113
114            palette[winner.item].l += (candidate.l - palette[winner.item].l) * gamma;
115            palette[winner.item].a += (candidate.a - palette[winner.item].a) * gamma;
116            palette[winner.item].b += (candidate.b - palette[winner.item].b) * gamma;
117
118            tree.add(
119                &[
120                    palette[winner.item].l,
121                    palette[winner.item].a,
122                    palette[winner.item].b,
123                ],
124                winner.item,
125            );
126
127            wc[winner.item] += 1;
128
129            if wc[winner.item] >= THETA && next_idx < PALETTE_SIZE {
130                tree.add(&[candidate.l, candidate.a, candidate.b], next_idx);
131
132                wc[winner.item] = 0;
133                wc[next_idx] = 0;
134                next_idx += 1;
135            }
136        }
137
138        palette
139            .into_iter()
140            .map(|lab| {
141                [
142                    OrderedFloat(lab.l),
143                    OrderedFloat(lab.a),
144                    OrderedFloat(lab.b),
145                ]
146            })
147            .collect::<HashSet<_>>()
148            .into_iter()
149            .map(|[l, a, b]| Lab::new(*l, *a, *b))
150            .collect::<Vec<_>>()
151    }
152}
a_sixel/adu.rs

a_sixel/
adu.rs
