1use std::{
6 collections::HashSet,
7 sync::atomic::{
8 AtomicU64,
9 Ordering,
10 },
11};
12
13use dilate::DilateExpand;
14use ordered_float::OrderedFloat;
15use palette::{
16 IntoColor,
17 Lab,
18 Srgb,
19};
20use rayon::iter::ParallelIterator;
21
22use crate::{
23 dither::Sierra,
24 private,
25 PaletteBuilder,
26 SixelEncoder,
27};
28
29pub type BitSixelEncoderMono<D = Sierra> = SixelEncoder<BitPaletteBuilder<2>, D>;
30pub type BitSixelEncoder4<D = Sierra> = SixelEncoder<BitPaletteBuilder<4>, D>;
31pub type BitSixelEncoder8<D = Sierra> = SixelEncoder<BitPaletteBuilder<8>, D>;
32pub type BitSixelEncoder16<D = Sierra> = SixelEncoder<BitPaletteBuilder<16>, D>;
33pub type BitSixelEncoder32<D = Sierra> = SixelEncoder<BitPaletteBuilder<32>, D>;
34pub type BitSixelEncoder64<D = Sierra> = SixelEncoder<BitPaletteBuilder<64>, D>;
35pub type BitSixelEncoder128<D = Sierra> = SixelEncoder<BitPaletteBuilder<128>, D>;
36pub type BitSixelEncoder256<D = Sierra> = SixelEncoder<BitPaletteBuilder<256>, D>;
37
38#[derive(Debug)]
39pub(crate) struct Bucket {
40 pub(crate) color: (AtomicU64, AtomicU64, AtomicU64),
41 pub(crate) count: AtomicU64,
42}
43
44#[derive(Debug)]
45pub struct BitPaletteBuilder<const PALETTE_SIZE: usize> {
46 pub(crate) buckets: Vec<Bucket>,
47}
48
49impl<const PALETTE_SIZE: usize> BitPaletteBuilder<PALETTE_SIZE> {
50 const PALETTE_DEPTH: usize = PALETTE_SIZE.ilog2() as usize;
51 const SHIFT: usize = 24 - Self::PALETTE_DEPTH;
52
53 pub(crate) fn new() -> Self {
54 BitPaletteBuilder {
55 buckets: Vec::from_iter(
56 std::iter::repeat_with(|| Bucket {
57 color: (AtomicU64::new(0), AtomicU64::new(0), AtomicU64::new(0)),
58 count: AtomicU64::new(0),
59 })
60 .take(PALETTE_SIZE),
61 ),
62 }
63 }
64
65 pub(crate) fn insert(&self, color: Srgb<u8>) {
66 let index = Self::index(color);
67 let node = &self.buckets[index];
68 node.color.0.fetch_add(color.red as u64, Ordering::Relaxed);
69 node.color
70 .1
71 .fetch_add(color.green as u64, Ordering::Relaxed);
72 node.color.2.fetch_add(color.blue as u64, Ordering::Relaxed);
73 node.count.fetch_add(1, Ordering::Relaxed);
74 }
75
76 pub(crate) fn index(color: Srgb<u8>) -> usize {
77 let r = color.red.dilate_expand::<3>().value();
78 let g = color.green.dilate_expand::<3>().value();
79 let b = color.blue.dilate_expand::<3>().value();
80
81 let rgb = g << 2 | r << 1 | b;
85
86 (rgb >> Self::SHIFT) as usize
87 }
88}
89
90impl<const PALETTE_SIZE: usize> private::Sealed for BitPaletteBuilder<PALETTE_SIZE> {}
91impl<const PALETTE_SIZE: usize> PaletteBuilder for BitPaletteBuilder<PALETTE_SIZE> {
92 const NAME: &'static str = "Bit";
93 const PALETTE_SIZE: usize = PALETTE_SIZE;
94
95 fn build_palette(image: &image::RgbImage) -> Vec<Lab> {
96 let builder = Self::new();
97
98 image.par_pixels().for_each(|pixel| {
99 builder.insert(Srgb::<u8>::new(pixel[0], pixel[1], pixel[2]));
100 });
101
102 builder
103 .buckets
104 .into_iter()
105 .filter(|node| node.count.load(Ordering::Relaxed) > 0)
106 .map(|node| {
107 let rgb = Srgb::new(
108 (node.color.0.load(Ordering::Relaxed) / node.count.load(Ordering::Relaxed))
109 as u8,
110 (node.color.1.load(Ordering::Relaxed) / node.count.load(Ordering::Relaxed))
111 as u8,
112 (node.color.2.load(Ordering::Relaxed) / node.count.load(Ordering::Relaxed))
113 as u8,
114 );
115 let lab: Lab = rgb.into_format().into_color();
116 [
117 OrderedFloat(lab.l),
118 OrderedFloat(lab.a),
119 OrderedFloat(lab.b),
120 ]
121 })
122 .collect::<HashSet<_>>()
123 .into_iter()
124 .map(|[l, a, b]| Lab::new(*l, *a, *b))
125 .collect::<Vec<_>>()
126 }
127}