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
20pub 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
88pub 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 const KERNEL: &[(isize, isize, f32)] = &[
127 (1, 0, 5.0),
128 (2, 0, 3.0),
129 (-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 (-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 const KERNEL: &[(isize, isize, f32)] = &[
150 (1, 0, 4.0),
152 (2, 0, 3.0),
153 (-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 const KERNEL: &[(isize, isize, f32)] = &[
169 (1, 0, 2.0),
171 (-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 const KERNEL: &[(isize, isize, f32)] = &[
184 (1, 0, 7.0),
186 (-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 const KERNEL: &[(isize, isize, f32)] = &[
201 (1, 0, 7.0),
203 (2, 0, 5.0),
204 (-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 (-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 const KERNEL: &[(isize, isize, f32)] = &[
227 (1, 0, 8.0),
229 (2, 0, 4.0),
230 (-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 (-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 const KERNEL: &[(isize, isize, f32)] = &[
253 (1, 0, 1.0),
255 (2, 0, 1.0),
256 (-1, 1, 1.0),
258 (0, 1, 1.0),
259 (1, 1, 1.0),
260 (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 const KERNEL: &[(isize, isize, f32)] = &[
272 (1, 0, 8.0),
274 (2, 0, 4.0),
275 (-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
353pub 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}