1use super::{Error, Result};
2use crate::media::Rgba;
3use rayon::prelude::*;
4use std::sync::LazyLock;
5
6const LEVELS: usize = 64;
7const CELLS: usize = LEVELS * LEVELS * LEVELS;
8const ALPHA_CUT: u8 = 128;
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum Dither {
12 None,
13 Ordered,
14 FloydSteinberg,
15}
16
17#[derive(Debug, Clone, PartialEq, Eq)]
18pub struct Palette {
19 colors: Vec<[u8; 3]>,
20}
21
22impl Palette {
23 pub fn new(colors: Vec<[u8; 3]>) -> Result<Self> {
24 if colors.is_empty() || colors.len() > 256 {
25 return Err(Error::Invalid("a palette holds 1 to 256 colours"));
26 }
27 Ok(Palette { colors })
28 }
29
30 pub fn colors(&self) -> &[[u8; 3]] {
31 &self.colors
32 }
33
34 pub fn len(&self) -> usize {
35 self.colors.len()
36 }
37
38 pub fn is_empty(&self) -> bool {
39 self.colors.is_empty()
40 }
41
42 pub fn from_frames(frames: &[&Rgba], colors: usize) -> Palette {
43 let mut h = Histogram::new();
44 h.add_frames(frames);
45 h.palette(colors)
46 }
47}
48
49fn to_linear(c: f32) -> f32 {
50 if c <= 0.04045 {
51 c / 12.92
52 } else {
53 ((c + 0.055) / 1.055).powf(2.4)
54 }
55}
56
57fn to_srgb(c: f32) -> f32 {
58 let c = c.clamp(0.0, 1.0);
59 if c <= 0.0031308 {
60 c * 12.92
61 } else {
62 1.055 * c.powf(1.0 / 2.4) - 0.055
63 }
64}
65
66fn linear_to_oklab(r: f32, g: f32, b: f32) -> [f32; 3] {
67 let l = (0.412_221_46 * r + 0.536_332_55 * g + 0.051_445_995 * b).cbrt();
68 let m = (0.211_903_5 * r + 0.680_699_5 * g + 0.107_396_96 * b).cbrt();
69 let s = (0.088_302_46 * r + 0.281_718_85 * g + 0.629_978_7 * b).cbrt();
70 [
71 0.210_454_26 * l + 0.793_617_8 * m - 0.004_072_047 * s,
72 1.977_998_5 * l - 2.428_592_2 * m + 0.450_593_7 * s,
73 0.025_904_037 * l + 0.782_771_77 * m - 0.808_675_77 * s,
74 ]
75}
76
77fn oklab_to_srgb(c: [f32; 3]) -> [u8; 3] {
78 let l = c[0] + 0.396_337_78 * c[1] + 0.215_803_76 * c[2];
79 let m = c[0] - 0.105_561_346 * c[1] - 0.063_854_17 * c[2];
80 let s = c[0] - 0.089_484_18 * c[1] - 1.291_485_5 * c[2];
81 let (l, m, s) = (l * l * l, m * m * m, s * s * s);
82 let r = 4.076_741_7 * l - 3.307_711_6 * m + 0.230_969_94 * s;
83 let g = -1.268_438 * l + 2.609_757_4 * m - 0.341_319_38 * s;
84 let b = -0.004_196_086_3 * l - 0.703_418_6 * m + 1.707_614_7 * s;
85 [r, g, b].map(|v| (to_srgb(v) * 255.0).round() as u8)
86}
87
88fn srgb_f_to_oklab(c: [f32; 3]) -> [f32; 3] {
89 let [r, g, b] = c.map(|v| to_linear(v / 255.0));
90 linear_to_oklab(r, g, b)
91}
92
93#[inline(always)]
94pub(crate) fn cell(r: u8, g: u8, b: u8) -> usize {
95 ((r as usize >> 2) << 12) | ((g as usize >> 2) << 6) | (b as usize >> 2)
96}
97
98pub struct Histogram {
99 bins: Vec<[u32; 4]>,
100}
101
102impl Default for Histogram {
103 fn default() -> Self {
104 Self::new()
105 }
106}
107
108impl Histogram {
109 pub fn new() -> Self {
110 Histogram {
111 bins: vec![[0; 4]; CELLS],
112 }
113 }
114
115 pub fn clear(&mut self) {
116 self.bins.fill([0; 4]);
117 }
118
119 pub fn add(&mut self, frame: &Rgba) {
120 self.add_bytes(&frame.data);
121 }
122
123 pub fn add_frames(&mut self, frames: &[&Rgba]) {
124 let rows: Vec<&[u8]> = frames
125 .iter()
126 .flat_map(|f| f.data.chunks(f.stride().max(4) * 64))
127 .collect();
128 let parts = rayon::current_num_threads().clamp(1, 8);
129 let per = rows.len().div_ceil(parts).max(1);
130 let partial: Vec<Histogram> = rows
131 .par_chunks(per)
132 .map(|chunk| {
133 let mut h = Histogram::new();
134 for r in chunk {
135 h.add_bytes(r);
136 }
137 h
138 })
139 .collect();
140 self.merge(&partial);
141 }
142
143 fn add_bytes(&mut self, data: &[u8]) {
144 let bins = &mut self.bins[..];
145 for p in data.chunks_exact(4) {
146 if p[3] < ALPHA_CUT {
147 continue;
148 }
149 let bin = &mut bins[cell(p[0], p[1], p[2])];
150 bin[0] += 1;
151 bin[1] += (p[0] & 3) as u32;
152 bin[2] += (p[1] & 3) as u32;
153 bin[3] += (p[2] & 3) as u32;
154 }
155 }
156
157 fn merge(&mut self, parts: &[Histogram]) {
158 if parts.is_empty() {
159 return;
160 }
161 self.bins
162 .par_chunks_mut(4096)
163 .enumerate()
164 .for_each(|(i, dst)| {
165 for part in parts {
166 let src = &part.bins[i * 4096..i * 4096 + dst.len()];
167 for (d, s) in dst.iter_mut().zip(src) {
168 for k in 0..4 {
169 d[k] = d[k].saturating_add(s[k]);
170 }
171 }
172 }
173 });
174 }
175
176 pub fn palette(&self, colors: usize) -> Palette {
177 let colors = colors.clamp(1, 256);
178 let mut points: Vec<Point> = self
179 .bins
180 .iter()
181 .enumerate()
182 .filter(|(_, b)| b[0] > 0)
183 .map(|(i, b)| {
184 let n = b[0] as f32;
185 let base = [(i >> 12) & 63, (i >> 6) & 63, i & 63].map(|v| (v << 2) as f32);
186 let rgb = [
187 base[0] + b[1] as f32 / n,
188 base[1] + b[2] as f32 / n,
189 base[2] + b[3] as f32 / n,
190 ];
191 Point {
192 lab: srgb_f_to_oklab(rgb),
193 rgb,
194 weight: n,
195 }
196 })
197 .collect();
198 if points.is_empty() {
199 return Palette {
200 colors: vec![[0, 0, 0]],
201 };
202 }
203 if points.len() <= colors {
204 return Palette {
205 colors: points
206 .iter()
207 .map(|p| p.rgb.map(|v| v.round() as u8))
208 .collect(),
209 };
210 }
211 let centers = median_cut(&mut points, colors);
212 let centers = kmeans(&points, centers, 3);
213 Palette {
214 colors: centers.into_iter().map(oklab_to_srgb).collect(),
215 }
216 }
217}
218
219#[derive(Clone, Copy)]
220struct Point {
221 lab: [f32; 3],
222 rgb: [f32; 3],
223 weight: f32,
224}
225
226struct Cut {
227 start: usize,
228 end: usize,
229 score: f64,
230 axis: usize,
231}
232
233fn measure(points: &[Point], start: usize, end: usize) -> Cut {
234 let pts = &points[start..end];
235 let mut w = 0.0f64;
236 let mut s = [0.0f64; 3];
237 let mut q = [0.0f64; 3];
238 for p in pts {
239 let pw = p.weight as f64;
240 w += pw;
241 for k in 0..3 {
242 let v = p.lab[k] as f64;
243 s[k] += pw * v;
244 q[k] += pw * v * v;
245 }
246 }
247 let var: [f64; 3] = std::array::from_fn(|k| (q[k] - s[k] * s[k] / w).max(0.0));
248 let axis = (0..3)
249 .max_by(|&a, &b| var[a].total_cmp(&var[b]))
250 .unwrap_or(0);
251 let score = if end - start < 2 {
252 -1.0
253 } else {
254 var.iter().sum()
255 };
256 Cut {
257 start,
258 end,
259 score,
260 axis,
261 }
262}
263
264fn median_cut(points: &mut [Point], colors: usize) -> Vec<[f32; 3]> {
265 let mut cuts = vec![measure(points, 0, points.len())];
266 while cuts.len() < colors {
267 let Some((i, _)) = cuts
268 .iter()
269 .enumerate()
270 .filter(|(_, c)| c.score > 0.0)
271 .max_by(|a, b| a.1.score.total_cmp(&b.1.score))
272 else {
273 break;
274 };
275 let c = cuts.swap_remove(i);
276 let slice = &mut points[c.start..c.end];
277 let axis = c.axis;
278 slice.sort_unstable_by(|a, b| a.lab[axis].total_cmp(&b.lab[axis]));
279 let total: f32 = slice.iter().map(|p| p.weight).sum();
280 let mut acc = 0.0;
281 let mut mid = 1;
282 for (k, p) in slice.iter().enumerate() {
283 acc += p.weight;
284 if acc >= total / 2.0 {
285 mid = (k + 1).clamp(1, slice.len() - 1);
286 break;
287 }
288 }
289 cuts.push(measure(points, c.start, c.start + mid));
290 cuts.push(measure(points, c.start + mid, c.end));
291 }
292 cuts.iter()
293 .map(|c| {
294 let mut w = 0.0f64;
295 let mut s = [0.0f64; 3];
296 for p in &points[c.start..c.end] {
297 w += p.weight as f64;
298 for (acc, v) in s.iter_mut().zip(p.lab) {
299 *acc += p.weight as f64 * v as f64;
300 }
301 }
302 s.map(|v| (v / w) as f32)
303 })
304 .collect()
305}
306
307struct Soa {
308 l: Vec<f32>,
309 a: Vec<f32>,
310 b: Vec<f32>,
311}
312
313impl Soa {
314 fn new(centers: &[[f32; 3]]) -> Soa {
315 let n = centers.len().next_multiple_of(8);
316 let mut s = Soa {
317 l: vec![1e9; n],
318 a: vec![1e9; n],
319 b: vec![1e9; n],
320 };
321 for (i, c) in centers.iter().enumerate() {
322 s.l[i] = c[0];
323 s.a[i] = c[1];
324 s.b[i] = c[2];
325 }
326 s
327 }
328
329 #[inline(always)]
330 fn nearest(&self, p: [f32; 3]) -> usize {
331 let mut best = f32::MAX;
332 let mut at = 0;
333 for (j, ((l, a), b)) in self
334 .l
335 .chunks_exact(8)
336 .zip(self.a.chunks_exact(8))
337 .zip(self.b.chunks_exact(8))
338 .enumerate()
339 {
340 let mut d = [0.0f32; 8];
341 for k in 0..8 {
342 let x = l[k] - p[0];
343 let y = a[k] - p[1];
344 let z = b[k] - p[2];
345 d[k] = x * x + y * y + z * z;
346 }
347 let m = d.iter().copied().fold(f32::MAX, f32::min);
348 if m < best {
349 best = m;
350 at = j * 8 + d.iter().position(|&v| v == m).unwrap_or(0);
351 }
352 }
353 at
354 }
355}
356
357fn kmeans(points: &[Point], mut centers: Vec<[f32; 3]>, rounds: usize) -> Vec<[f32; 3]> {
358 for _ in 0..rounds {
359 let soa = Soa::new(¢ers);
360 let k = centers.len();
361 let parts: Vec<Vec<[f64; 4]>> = points
362 .par_chunks(4096)
363 .map(|chunk| {
364 let mut acc = vec![[0.0f64; 4]; k];
365 for p in chunk {
366 let i = nearest_lab(&soa, p.lab);
367 let w = p.weight as f64;
368 acc[i][0] += w;
369 for c in 0..3 {
370 acc[i][c + 1] += w * p.lab[c] as f64;
371 }
372 }
373 acc
374 })
375 .collect();
376 let mut sums = vec![[0.0f64; 4]; k];
377 for part in &parts {
378 for (x, y) in sums.iter_mut().zip(part) {
379 for c in 0..4 {
380 x[c] += y[c];
381 }
382 }
383 }
384 for (c, s) in centers.iter_mut().zip(&sums) {
385 if s[0] > 0.0 {
386 *c = [
387 (s[1] / s[0]) as f32,
388 (s[2] / s[0]) as f32,
389 (s[3] / s[0]) as f32,
390 ];
391 }
392 }
393 }
394 centers
395}
396
397crate::media::isa::dispatch! {
398 fn nearest_lab(soa: &Soa, p: [f32; 3]) -> usize {
399 soa.nearest(p)
400 }
401}
402
403static CELL_LAB: LazyLock<Vec<[f32; 3]>> = LazyLock::new(|| {
404 let level: Vec<f32> = (0..LEVELS)
405 .map(|v| to_linear(((v << 2) as f32 + 1.5) / 255.0))
406 .collect();
407 (0..CELLS)
408 .map(|i| linear_to_oklab(level[i >> 12], level[(i >> 6) & 63], level[i & 63]))
409 .collect()
410});
411
412pub(crate) struct Lut {
413 pub cells: Vec<u8>,
414}
415
416impl Lut {
417 pub fn new(palette: &Palette) -> Lut {
418 let centers: Vec<[f32; 3]> = palette
419 .colors
420 .iter()
421 .map(|c| c.map(|v| v as f32))
422 .map(srgb_f_to_oklab)
423 .collect();
424 let mut cells = vec![0u8; CELLS];
425 cells
426 .par_chunks_mut(SLAB)
427 .enumerate()
428 .for_each(|(br, out)| fill_slab(¢ers, br, out));
429 Lut { cells }
430 }
431}
432
433const SLAB: usize = LEVELS * LEVELS * 4;
434
435crate::media::isa::dispatch! {
436 fn fill_slab(centers: &[[f32; 3]], br: usize, out: &mut [u8]) {
437 let labs = &CELL_LAB[br * SLAB..(br + 1) * SLAB];
438 let mut candidates: Vec<u8> = Vec::with_capacity(256);
439 let mut near = vec![0f32; centers.len()];
440 for bg in 0..16 {
441 for bb in 0..16 {
442 let block = |r: usize, g: usize, b: usize| (r << 12) | ((bg * 4 + g) << 6) | (bb * 4 + b);
443 let mut lo = [f32::MAX; 3];
444 let mut hi = [f32::MIN; 3];
445 for r in 0..4 {
446 for g in 0..4 {
447 for b in 0..4 {
448 let p = labs[block(r, g, b)];
449 for k in 0..3 {
450 lo[k] = lo[k].min(p[k]);
451 hi[k] = hi[k].max(p[k]);
452 }
453 }
454 }
455 }
456 let mut bound = f32::MAX;
457 for (c, n) in centers.iter().zip(near.iter_mut()) {
458 let mut low = 0.0;
459 let mut high = 0.0;
460 for k in 0..3 {
461 let d = (lo[k] - c[k]).max(c[k] - hi[k]).max(0.0);
462 low += d * d;
463 let f = (c[k] - lo[k]).abs().max((c[k] - hi[k]).abs());
464 high += f * f;
465 }
466 *n = low;
467 bound = bound.min(high);
468 }
469 candidates.clear();
470 for (i, &n) in near.iter().enumerate() {
471 if n <= bound {
472 candidates.push(i as u8);
473 }
474 }
475 for r in 0..4 {
476 for g in 0..4 {
477 for b in 0..4 {
478 let at = block(r, g, b);
479 let p = labs[at];
480 let mut best = f32::MAX;
481 let mut pick = 0;
482 for &i in &candidates {
483 let c = centers[i as usize];
484 let x = c[0] - p[0];
485 let y = c[1] - p[1];
486 let z = c[2] - p[2];
487 let d = x * x + y * y + z * z;
488 if d < best {
489 best = d;
490 pick = i;
491 }
492 }
493 out[at] = pick;
494 }
495 }
496 }
497 }
498 }
499 }
500}
501
502const BAYER: [[u8; 8]; 8] = [
503 [0, 32, 8, 40, 2, 34, 10, 42],
504 [48, 16, 56, 24, 50, 18, 58, 26],
505 [12, 44, 4, 36, 14, 46, 6, 38],
506 [60, 28, 52, 20, 62, 30, 54, 22],
507 [3, 35, 11, 43, 1, 33, 9, 41],
508 [51, 19, 59, 27, 49, 17, 57, 25],
509 [15, 47, 7, 39, 13, 45, 5, 37],
510 [63, 31, 55, 23, 61, 29, 53, 21],
511];
512
513pub(crate) struct Mapper<'a> {
514 pub lut: &'a Lut,
515 pub palette: &'a [[u8; 3]],
516 pub dither: Dither,
517 pub transparent: u8,
518}
519
520impl Mapper<'_> {
521 pub fn map(&self, src: &Rgba, out: &mut [u8], err: &mut Vec<i16>) {
522 let w = src.width as usize;
523 if w == 0 || src.height == 0 {
524 return;
525 }
526 match self.dither {
527 Dither::FloydSteinberg => self.map_fs(src, out, err),
528 _ => {
529 let band = 16;
530 out.par_chunks_mut(w * band)
531 .zip(src.data.par_chunks(w * 4 * band))
532 .enumerate()
533 .for_each(|(i, (o, s))| self.map_rows(s, o, w, i * band));
534 }
535 }
536 }
537
538 fn map_rows(&self, src: &[u8], out: &mut [u8], w: usize, y0: usize) {
539 let t = self.transparent;
540 let cells = &self.lut.cells[..];
541 if self.dither == Dither::None {
542 for (o, p) in out.iter_mut().zip(src.chunks_exact(4)) {
543 *o = if p[3] < ALPHA_CUT {
544 t
545 } else {
546 cells[cell(p[0], p[1], p[2])]
547 };
548 }
549 return;
550 }
551 let amp = (96.0 / (self.palette.len() as f32).cbrt()).round() as i32;
552 for (y, (orow, srow)) in out.chunks_mut(w).zip(src.chunks(w * 4)).enumerate() {
553 let brow = &BAYER[(y0 + y) & 7];
554 for (x, (o, p)) in orow.iter_mut().zip(srow.chunks_exact(4)).enumerate() {
555 if p[3] < ALPHA_CUT {
556 *o = t;
557 continue;
558 }
559 let d = ((brow[x & 7] as i32 * 2 - 63) * amp) >> 7;
560 let r = (p[0] as i32 + d).clamp(0, 255) as u8;
561 let g = (p[1] as i32 + d).clamp(0, 255) as u8;
562 let b = (p[2] as i32 + d).clamp(0, 255) as u8;
563 *o = cells[cell(r, g, b)];
564 }
565 }
566 }
567
568 fn map_fs(&self, src: &Rgba, out: &mut [u8], err: &mut Vec<i16>) {
569 let w = src.width as usize;
570 let h = src.height as usize;
571 let n = (w + 2) * 3;
572 err.clear();
573 err.resize(n * 2, 0);
574 let (mut cur, mut nxt) = err.split_at_mut(n);
575 let cells = &self.lut.cells[..];
576 let pal = self.palette;
577 for y in 0..h {
578 nxt.fill(0);
579 let srow = &src.data[y * w * 4..(y + 1) * w * 4];
580 let orow = &mut out[y * w..(y + 1) * w];
581 let rev = y & 1 == 1;
582 for i in 0..w {
583 let x = if rev { w - 1 - i } else { i };
584 let p = &srow[x * 4..x * 4 + 4];
585 if p[3] < ALPHA_CUT {
586 orow[x] = self.transparent;
587 continue;
588 }
589 let e = (x + 1) * 3;
590 let want = [
591 (p[0] as i32 + cur[e] as i32).clamp(0, 255),
592 (p[1] as i32 + cur[e + 1] as i32).clamp(0, 255),
593 (p[2] as i32 + cur[e + 2] as i32).clamp(0, 255),
594 ];
595 let idx = cells[cell(want[0] as u8, want[1] as u8, want[2] as u8)];
596 orow[x] = idx;
597 let got = pal[idx as usize];
598 let (fwd, back) = if rev { (e - 3, e + 3) } else { (e + 3, e - 3) };
599 for c in 0..3 {
600 let d = want[c] - got[c] as i32;
601 cur[fwd + c] = (cur[fwd + c] as i32 + d * 7 / 16).clamp(-255, 255) as i16;
602 nxt[back + c] += (d * 3 / 16) as i16;
603 nxt[e + c] += (d * 5 / 16) as i16;
604 nxt[fwd + c] += (d / 16) as i16;
605 }
606 }
607 std::mem::swap(&mut cur, &mut nxt);
608 }
609 }
610}