a_sixel/
bit.rs

1//! Encodes a palette by bucketing bit ranges into a power of two number of
2//! buckets. This is very fast and produces ok results for most images at larger
3//! palette sizes (e.g. 256).
4
5use 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        // Since elements to the right will get shifted off first, we put them in grb
82        // order (order of most-least significant for luminance). This probably doesn't
83        // make a huge difference, but the theory is nice.
84        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}