a_sixel/
dither.rs

1use dilate::DilateExpand;
2use image::RgbImage;
3use kiddo::float::kdtree::KdTree;
4use palette::Lab;
5use rayon::{
6    iter::{
7        IndexedParallelIterator,
8        IntoParallelIterator,
9        ParallelIterator,
10    },
11    slice::ParallelSliceMut,
12};
13use sobol_burley::sample_4d;
14
15use crate::{
16    private,
17    rgb_to_lab,
18};
19
20/// Each struct in this module implements this trait and can be combined with
21/// the [`SixelEncoder`](crate::SixelEncoder) struct to dither the result.
22///
23/// Most of the dithering algorithms are based on the ones described in
24/// [this article](https://tannerhelland.com/2012/12/28/dithering-eleven-algorithms-source-code.html).
25/// There is also an implementation of the Bayer matrix dithering algorithm
26/// and a Sobol sequence based dithering algorithm.
27pub trait Dither: private::Sealed {
28    const KERNEL: &[(isize, isize, f32)];
29    const DIV: f32;
30
31    fn dither_and_palettize(
32        image: &RgbImage,
33        in_palette: &[Lab],
34        target_palette_size: usize,
35    ) -> Vec<usize> {
36        let _ = target_palette_size;
37
38        let pixels = image.pixels().copied().map(rgb_to_lab).collect::<Vec<_>>();
39
40        let mut palette = KdTree::<_, _, 3, 257, u32>::with_capacity(in_palette.len());
41
42        for (idx, color) in in_palette.iter().enumerate() {
43            palette.add(color.as_ref(), idx);
44        }
45
46        let mut result = vec![0; image.width() as usize * image.height() as usize];
47
48        let mut spills = vec![[0.0; 3]; image.width() as usize * image.height() as usize];
49
50        for (idx, p) in pixels.iter().copied().enumerate() {
51            let pixel = <Lab>::new(
52                p.l + spills[idx][0],
53                p.a + spills[idx][1],
54                p.b + spills[idx][2],
55            );
56
57            let palette_idx =
58                palette.nearest_one::<kiddo::SquaredEuclidean>(&[pixel.l, pixel.a, pixel.b]);
59
60            let error0 = pixel.l - in_palette[palette_idx.item].l;
61            let error1 = pixel.a - in_palette[palette_idx.item].a;
62            let error2 = pixel.b - in_palette[palette_idx.item].b;
63            let spill = [error0, error1, error2];
64
65            result[idx] = palette_idx.item;
66
67            for (dx, dy, m) in Self::KERNEL {
68                let x = idx as isize % image.width() as isize + dx;
69                let y = idx as isize / image.width() as isize + dy;
70                if x < 0 || y < 0 || x >= image.width() as isize || y >= image.height() as isize {
71                    continue;
72                }
73
74                let jdx = y * image.width() as isize + x;
75                let target = &mut spills[jdx as usize];
76                *target = [
77                    target[0] + (spill[0] * m) / Self::DIV / 2.0,
78                    target[1] + (spill[1] * m) / Self::DIV / 2.0,
79                    target[2] + (spill[2] * m) / Self::DIV / 2.0,
80                ];
81            }
82        }
83
84        result
85    }
86}
87
88/// Do not perform dithering. This is the fastest possible palette encoding
89/// option.
90pub struct NoDither;
91impl private::Sealed for NoDither {}
92impl Dither for NoDither {
93    const DIV: f32 = 1.0;
94    const KERNEL: &[(isize, isize, f32)] = &[];
95
96    fn dither_and_palettize(image: &RgbImage, in_palette: &[Lab], _: usize) -> Vec<usize> {
97        let pixels = image.pixels().copied().map(rgb_to_lab).collect::<Vec<_>>();
98
99        let mut palette = KdTree::<_, _, 3, 32, u32>::with_capacity(in_palette.len());
100
101        for (idx, color) in in_palette.iter().enumerate() {
102            palette.add(color.as_ref(), idx);
103        }
104
105        let mut result = vec![0; image.width() as usize * image.height() as usize];
106        pixels
107            .into_par_iter()
108            .zip(&mut result)
109            .for_each(|(p, dest)| {
110                *dest = palette
111                    .nearest_one::<kiddo::SquaredEuclidean>(&[p.l, p.a, p.b])
112                    .item;
113            });
114
115        result
116    }
117}
118
119pub struct Sierra;
120impl private::Sealed for Sierra {}
121impl Dither for Sierra {
122    const DIV: f32 = 32.0;
123    // _ _ X 5 3
124    // 2 4 5 4 2
125    // _ 2 3 2 _
126    const KERNEL: &[(isize, isize, f32)] = &[
127        (1, 0, 5.0),
128        (2, 0, 3.0),
129        //
130        (-2, 1, 2.0),
131        (-1, 1, 4.0),
132        (0, 1, 5.0),
133        (1, 1, 4.0),
134        (2, 1, 2.0),
135        //
136        (-1, 2, 2.0),
137        (0, 2, 3.0),
138        (1, 2, 2.0),
139    ];
140}
141
142pub struct Sierra2;
143impl private::Sealed for Sierra2 {}
144impl Dither for Sierra2 {
145    const DIV: f32 = 16.0;
146    // _ _ X 4 3
147    // 1 2 3 2 1
148    // _ 1 2 1 _
149    const KERNEL: &[(isize, isize, f32)] = &[
150        //x y  m
151        (1, 0, 4.0),
152        (2, 0, 3.0),
153        //
154        (-2, 1, 1.0),
155        (-1, 1, 2.0),
156        (0, 1, 3.0),
157        (1, 1, 2.0),
158        (2, 1, 1.0),
159    ];
160}
161
162pub struct SierraLite;
163impl private::Sealed for SierraLite {}
164impl Dither for SierraLite {
165    const DIV: f32 = 4.0;
166    //  _ X 2
167    //  1 1 _
168    const KERNEL: &[(isize, isize, f32)] = &[
169        //x y  m
170        (1, 0, 2.0),
171        //
172        (-1, 1, 1.0),
173        (0, 1, 1.0),
174    ];
175}
176
177pub struct FloydSteinberg;
178impl private::Sealed for FloydSteinberg {}
179impl Dither for FloydSteinberg {
180    const DIV: f32 = 16.0;
181    // _ X 7
182    // 3 5 3
183    const KERNEL: &[(isize, isize, f32)] = &[
184        //x y  m
185        (1, 0, 7.0),
186        //
187        (-1, 1, 3.0),
188        (0, 1, 5.0),
189        (1, 1, 1.0),
190    ];
191}
192
193pub struct JJN;
194impl private::Sealed for JJN {}
195impl Dither for JJN {
196    const DIV: f32 = 48.0;
197    // _ _ X 7 5
198    // 3 5 7 5 3
199    // 1 2 5 3 1
200    const KERNEL: &[(isize, isize, f32)] = &[
201        //x y  m
202        (1, 0, 7.0),
203        (2, 0, 5.0),
204        //
205        (-2, 1, 3.0),
206        (-1, 1, 5.0),
207        (0, 1, 7.0),
208        (1, 1, 5.0),
209        (2, 1, 3.0),
210        //
211        (-2, 2, 1.0),
212        (-1, 2, 2.0),
213        (0, 2, 5.0),
214        (1, 2, 3.0),
215        (2, 2, 1.0),
216    ];
217}
218
219pub struct Stucki;
220impl private::Sealed for Stucki {}
221impl Dither for Stucki {
222    const DIV: f32 = 42.0;
223    // _ _ X 8 4
224    // 2 4 8 4 2
225    // 1 2 4 2 1
226    const KERNEL: &[(isize, isize, f32)] = &[
227        //x y  m
228        (1, 0, 8.0),
229        (2, 0, 4.0),
230        //
231        (-2, 1, 2.0),
232        (-1, 1, 4.0),
233        (0, 1, 8.0),
234        (1, 1, 4.0),
235        (2, 1, 2.0),
236        //
237        (-2, 2, 1.0),
238        (-1, 2, 2.0),
239        (0, 2, 4.0),
240        (1, 2, 2.0),
241        (2, 2, 1.0),
242    ];
243}
244
245pub struct Atkinson;
246impl private::Sealed for Atkinson {}
247impl Dither for Atkinson {
248    const DIV: f32 = 8.0;
249    // _ X 1 1
250    // 1 1 1 _
251    // _ 1 _ _
252    const KERNEL: &[(isize, isize, f32)] = &[
253        //x y  m
254        (1, 0, 1.0),
255        (2, 0, 1.0),
256        //
257        (-1, 1, 1.0),
258        (0, 1, 1.0),
259        (1, 1, 1.0),
260        //
261        (0, 2, 1.0),
262    ];
263}
264
265pub struct Burkes;
266impl private::Sealed for Burkes {}
267impl Dither for Burkes {
268    const DIV: f32 = 32.0;
269    // _ _ X 8 4
270    // 2 4 8 4 2
271    const KERNEL: &[(isize, isize, f32)] = &[
272        //x y  m
273        (1, 0, 8.0),
274        (2, 0, 4.0),
275        //
276        (-2, 1, 2.0),
277        (-1, 1, 4.0),
278        (0, 1, 8.0),
279        (1, 1, 4.0),
280        (2, 1, 2.0),
281    ];
282}
283
284pub struct Bayer;
285impl private::Sealed for Bayer {}
286impl Dither for Bayer {
287    const DIV: f32 = 0.0;
288    const KERNEL: &[(isize, isize, f32)] = &[];
289
290    fn dither_and_palettize(
291        image: &RgbImage,
292        in_palette: &[Lab],
293        target_palette_size: usize,
294    ) -> Vec<usize> {
295        let r = (1usize << 8).div_ceil(target_palette_size.div_ceil(3));
296        let matrix_size = r.next_power_of_two();
297        let r = r as f32;
298        let mut l_matrix = vec![0.0; matrix_size * matrix_size];
299        let mut a_matrix = vec![0.0; matrix_size * matrix_size];
300        let mut b_matrix = vec![0.0; matrix_size * matrix_size];
301
302        (0..matrix_size as u32)
303            .into_par_iter()
304            .zip(l_matrix.par_chunks_mut(matrix_size))
305            .zip(a_matrix.par_chunks_mut(matrix_size))
306            .zip(b_matrix.par_chunks_mut(matrix_size))
307            .for_each(|(((y, row_l), row_a), row_b)| {
308                for (x, ((cell_l, cell_a), cell_b)) in
309                    row_l.iter_mut().zip(row_a).zip(row_b).enumerate()
310                {
311                    let x = x as u32;
312                    let bits = (x ^ y).dilate_expand::<2>().value()
313                        | (y.dilate_expand::<2>().value() << 1);
314                    let bits = bits.reverse_bits() >> bits.leading_zeros();
315                    let partial_thresh = bits as f32 / matrix_size.pow(2) as f32;
316
317                    *cell_l = partial_thresh - 0.5;
318                    *cell_a = partial_thresh - 0.5;
319                    *cell_b = partial_thresh - 0.5;
320                }
321            });
322
323        let mut palette = KdTree::<_, _, 3, 257, u32>::with_capacity(in_palette.len());
324        for (idx, color) in in_palette.iter().enumerate() {
325            palette.add(color.as_ref(), idx);
326        }
327
328        let mut result = vec![0; image.width() as usize * image.height() as usize];
329        let pixels = image.pixels().copied().map(rgb_to_lab).collect::<Vec<_>>();
330
331        pixels
332            .into_par_iter()
333            .enumerate()
334            .zip(&mut result)
335            .for_each(|((idx, p), dest)| {
336                let x = (idx % image.width() as usize) % matrix_size;
337                let y = (idx / image.width() as usize) % matrix_size;
338
339                let m_idx = x + y * matrix_size;
340                let l = p.l + r * l_matrix[m_idx];
341                let a = p.a + r * a_matrix[m_idx];
342                let b = p.b + r * b_matrix[m_idx];
343
344                *dest = palette
345                    .nearest_one::<kiddo::SquaredEuclidean>(&[l, a, b])
346                    .item;
347            });
348
349        result
350    }
351}
352
353/// Uses the Sobol sequence to perturb the colors in a quasi-random manner. This
354/// is somewhere between ordered dithering and stochastic dithering. The results
355/// are generally decent, although they may appear textured like paper.
356pub struct Sobol;
357impl private::Sealed for Sobol {}
358impl Dither for Sobol {
359    const DIV: f32 = 0.0;
360    const KERNEL: &[(isize, isize, f32)] = &[];
361
362    fn dither_and_palettize(
363        image: &RgbImage,
364        in_palette: &[Lab],
365        target_palette_size: usize,
366    ) -> Vec<usize> {
367        let r = (1usize << 8).div_ceil(target_palette_size * target_palette_size.ilog2() as usize)
368            as f32;
369
370        let mut palette = KdTree::<_, _, 3, 257, u32>::with_capacity(in_palette.len());
371        for (idx, color) in in_palette.iter().enumerate() {
372            palette.add(color.as_ref(), idx);
373        }
374
375        let mut result = vec![0; image.width() as usize * image.height() as usize];
376        let pixels = image.pixels().copied().map(rgb_to_lab).collect::<Vec<_>>();
377
378        pixels
379            .into_par_iter()
380            .enumerate()
381            .zip(&mut result)
382            .for_each(|((idx, p), dest)| {
383                let [l, a, b, _] = sample_4d(idx as u32 % (1 << 16), 0, idx as u32 / (1 << 16));
384
385                let l = p.l + (l - 0.5) * 2.0 * r;
386                let a = p.a + (a - 0.5) * 2.0 * r;
387                let b = p.b + (b - 0.5) * 2.0 * r;
388
389                *dest = palette
390                    .nearest_one::<kiddo::SquaredEuclidean>(&[l, a, b])
391                    .item;
392            });
393
394        result
395    }
396}