Skip to main content

falcon_library/
palette.rs

1//! Dominant-colour palettes, matching Eagle's `palettes` field.
2//!
3//! Eagle runs modified median cut quantization (MMCQ, the algorithm behind
4//! color-thief) on a copy of the image at most 360 px wide, then turns box
5//! populations into percentages. This port keeps Eagle's exact control flow,
6//! quirks included, so the same pixels give the same palette:
7//!
8//! - the colour bounds are tracked with `if min .. else if max`, so a channel's
9//!   first value only ever sets its minimum;
10//! - each of the two splitting passes stops after 5 iterations, which caps a
11//!   palette at about 11 colours;
12//! - boxes are kept in a lazily sorted queue whose tie order matters.
13
14use image::{DynamicImage, RgbaImage};
15use serde_json::Number;
16
17use crate::model::Palette;
18
19const SIGBITS: u32 = 5;
20const RSHIFT: u32 = 8 - SIGBITS;
21const BINS: usize = 1 << SIGBITS;
22const MAX_ITERATIONS: u32 = 5;
23const FRACT_BY_POPULATIONS: f64 = 0.75;
24const MAX_COLORS: usize = 12;
25const ALPHA_THRESHOLD: u8 = 170;
26/// Width of the image Eagle analyses.
27pub const ANALYSIS_WIDTH: u32 = 360;
28
29/// Extracts the palette Eagle would store for `img`.
30///
31/// `img` is the image Eagle analyses: the item's thumbnail, or the original
32/// when the item has no thumbnail. It is shrunk to [`ANALYSIS_WIDTH`] first.
33pub fn extract(img: &DynamicImage) -> Vec<Palette> {
34    from_rgba(analysis_image(img).as_raw())
35}
36
37/// Extracts a palette from an RGBA buffer that is already analysis-sized.
38pub fn from_rgba(rgba: &[u8]) -> Vec<Palette> {
39    let Some(swatches) = quantize(rgba, skip_for(rgba.len()), ALPHA_THRESHOLD, MAX_COLORS) else {
40        return Vec::new();
41    };
42    to_palettes(swatches)
43}
44
45/// The image Eagle hands to the quantizer: `img` scaled down (never up) to at
46/// most [`ANALYSIS_WIDTH`] wide, keeping the aspect ratio.
47///
48/// Eagle draws onto a canvas with default smoothing, which Chromium does as a
49/// single bilinear sample per pixel; this mimics that rather than using an
50/// area filter.
51pub fn analysis_image(img: &DynamicImage) -> RgbaImage {
52    let src = img.to_rgba8();
53    let (w, h) = src.dimensions();
54    if w <= ANALYSIS_WIDTH {
55        return src;
56    }
57    let dw = ANALYSIS_WIDTH;
58    let dh = ((h as f64 / (w as f64 / dw as f64)).round() as u32).max(1);
59    bilinear(&src, dw, dh)
60}
61
62/// Every Nth pixel is sampled, depending on the buffer's size in bytes.
63fn skip_for(bytes: usize) -> usize {
64    match bytes {
65        0..100_000 => 1,
66        100_000..1_000_000 => 10,
67        1_000_000..10_000_000 => 100,
68        10_000_000..100_000_000 => 1000,
69        _ => 10000,
70    }
71}
72
73/// One MMCQ colour and the number of sampled pixels it stands for, with the
74/// total number of sampled pixels.
75struct Swatches {
76    colors: Vec<([u8; 3], u64)>,
77    pixel_count: u64,
78}
79
80/// Turns box populations into Eagle's rounded, filtered percentages.
81fn to_palettes(s: Swatches) -> Vec<Palette> {
82    let mut colors: Vec<([u8; 3], f64)> = s
83        .colors
84        .into_iter()
85        .map(|(color, count)| {
86            let ratio = count as f64 / s.pixel_count as f64 * 100.0;
87            let ratio = (ratio * 100.0).round() / 100.0;
88            // Large shares are truncated to whole percents.
89            let ratio = if ratio >= 5.0 { ratio.trunc() } else { ratio };
90            (color, ratio)
91        })
92        .filter(|&(_, ratio)| ratio >= 0.25)
93        .collect();
94    // Stable, like Array.prototype.sort, so equal ratios keep box order.
95    colors.sort_by(|a, b| b.1.total_cmp(&a.1));
96
97    if colors.len() > 4 {
98        let mut kept = Vec::new();
99        for (i, c) in colors.into_iter().enumerate() {
100            let threshold = if i < 6 { 0.1 } else { 0.3 };
101            if kept.len() < 5 || c.1 > threshold {
102                kept.push(c);
103            }
104        }
105        colors = kept;
106    }
107    colors.truncate(MAX_COLORS);
108
109    colors
110        .into_iter()
111        .map(|(color, ratio)| Palette { color, ratio: js_number(ratio), extra: Default::default() })
112        .collect()
113}
114
115/// A JSON number written the way JavaScript prints it: whole values have no
116/// fractional part.
117fn js_number(x: f64) -> Number {
118    if x.fract() == 0.0 {
119        Number::from(x as u64)
120    } else {
121        Number::from_f64(x).expect("ratios are finite")
122    }
123}
124
125/// Population histogram over 5-bit RGB.
126struct Histogram(Vec<u64>);
127
128impl Histogram {
129    fn at(&self, r: i32, g: i32, b: i32) -> u64 {
130        self.0[((r << (2 * SIGBITS)) + (g << SIGBITS) + b) as usize]
131    }
132}
133
134/// An axis of colour space.
135#[derive(Clone, Copy, PartialEq)]
136enum Axis {
137    R,
138    G,
139    B,
140}
141
142/// A box in 5-bit RGB space. Bounds are inclusive and can be inverted (empty)
143/// because of the bounds quirk in [`quantize`].
144#[derive(Clone)]
145struct VBox {
146    lo: [i32; 3],
147    hi: [i32; 3],
148    count: u64,
149}
150
151impl VBox {
152    fn new(lo: [i32; 3], hi: [i32; 3], histo: &Histogram) -> VBox {
153        let mut count = 0;
154        for r in lo[0]..=hi[0] {
155            for g in lo[1]..=hi[1] {
156                for b in lo[2]..=hi[2] {
157                    count += histo.at(r, g, b);
158                }
159            }
160        }
161        VBox { lo, hi, count }
162    }
163
164    fn width(&self, axis: Axis) -> i32 {
165        let i = axis as usize;
166        self.hi[i] - self.lo[i] + 1
167    }
168
169    fn volume(&self) -> i64 {
170        [Axis::R, Axis::G, Axis::B].iter().map(|&a| self.width(a) as i64).product()
171    }
172
173    /// Population-weighted mean colour of the bins' centres, truncated.
174    fn avg(&self, histo: &Histogram) -> [u8; 3] {
175        let mult = (1 << RSHIFT) as f64;
176        let (mut total, mut sum) = (0u64, [0f64; 3]);
177        for r in self.lo[0]..=self.hi[0] {
178            for g in self.lo[1]..=self.hi[1] {
179                for b in self.lo[2]..=self.hi[2] {
180                    let n = histo.at(r, g, b);
181                    total += n;
182                    for (s, v) in sum.iter_mut().zip([r, g, b]) {
183                        *s += n as f64 * (v as f64 + 0.5) * mult;
184                    }
185                }
186            }
187        }
188        if total > 0 {
189            sum.map(|s| (s / total as f64) as u8)
190        } else {
191            std::array::from_fn(|i| (mult * (self.lo[i] + self.hi[i] + 1) as f64 / 2.0) as u8)
192        }
193    }
194}
195
196/// The JS queue Eagle uses: pushes append, and the contents are stably sorted
197/// in ascending order only when popping after a push. `pop` takes the last
198/// (largest) box.
199struct LazyQueue {
200    boxes: Vec<VBox>,
201    sorted: bool,
202    key: fn(&VBox) -> i64,
203}
204
205impl LazyQueue {
206    fn new(key: fn(&VBox) -> i64) -> LazyQueue {
207        LazyQueue { boxes: Vec::new(), sorted: false, key }
208    }
209
210    fn push(&mut self, b: VBox) {
211        self.boxes.push(b);
212        self.sorted = false;
213    }
214
215    fn pop(&mut self) -> Option<VBox> {
216        if !self.sorted {
217            let key = self.key;
218            self.boxes.sort_by_key(key);
219            self.sorted = true;
220        }
221        self.boxes.pop()
222    }
223
224    fn len(&self) -> usize {
225        self.boxes.len()
226    }
227}
228
229fn by_count(b: &VBox) -> i64 {
230    b.count as i64
231}
232
233fn by_count_times_volume(b: &VBox) -> i64 {
234    b.count as i64 * b.volume()
235}
236
237/// Quantizes every `skip`th pixel of `rgba` with alpha at least
238/// `alpha_threshold` into at most `max_colors` boxes.
239fn quantize(rgba: &[u8], skip: usize, alpha_threshold: u8, max_colors: usize) -> Option<Swatches> {
240    let mut histo = Histogram(vec![0; BINS * BINS * BINS]);
241    let (mut lo, mut hi) = ([31i32; 3], [0i32; 3]);
242    let mut pixel_count = 0u64;
243    for px in rgba.as_chunks::<4>().0.iter().step_by(skip) {
244        if px[3] < alpha_threshold {
245            continue;
246        }
247        let v = [px[0], px[1], px[2]].map(|c| (c >> RSHIFT) as i32);
248        histo.0[((v[0] << (2 * SIGBITS)) + (v[1] << SIGBITS) + v[2]) as usize] += 1;
249        pixel_count += 1;
250        for i in 0..3 {
251            // Eagle's quirk: a value that lowers the minimum never raises the
252            // maximum.
253            if v[i] < lo[i] {
254                lo[i] = v[i];
255            } else if v[i] > hi[i] {
256                hi[i] = v[i];
257            }
258        }
259    }
260    if pixel_count == 0 {
261        return None;
262    }
263
264    let mut pq = LazyQueue::new(by_count);
265    pq.push(VBox::new(lo, hi, &histo));
266    split(&mut pq, FRACT_BY_POPULATIONS * max_colors as f64, &histo);
267
268    let mut pq2 = LazyQueue::new(by_count_times_volume);
269    while let Some(b) = pq.pop() {
270        pq2.push(b);
271    }
272    let target = max_colors as f64 - pq2.len() as f64;
273    split(&mut pq2, target, &histo);
274
275    let mut colors = Vec::new();
276    while let Some(b) = pq2.pop() {
277        colors.push((b.avg(&histo), b.count));
278    }
279    Some(Swatches { colors, pixel_count })
280}
281
282/// Splits the largest box repeatedly until `target` boxes have been added or
283/// the iteration budget runs out.
284fn split(queue: &mut LazyQueue, target: f64, histo: &Histogram) {
285    let mut ncolors = 1.0;
286    let mut iterations = 0;
287    while iterations < MAX_ITERATIONS {
288        let Some(vbox) = queue.pop() else { return };
289        if vbox.count == 0 {
290            queue.push(vbox);
291            iterations += 1;
292            continue;
293        }
294        let Some((first, second)) = median_cut(histo, &vbox) else {
295            return;
296        };
297        queue.push(first);
298        if let Some(second) = second {
299            queue.push(second);
300            ncolors += 1.0;
301        }
302        if ncolors >= target {
303            return;
304        }
305        iterations += 1;
306    }
307}
308
309/// Cuts `vbox` across its longest side near the population median.
310fn median_cut(histo: &Histogram, vbox: &VBox) -> Option<(VBox, Option<VBox>)> {
311    if vbox.count == 0 {
312        return None;
313    }
314    if vbox.count == 1 {
315        return Some((vbox.clone(), None));
316    }
317    let (rw, gw, bw) = (vbox.width(Axis::R), vbox.width(Axis::G), vbox.width(Axis::B));
318    let maxw = rw.max(gw).max(bw);
319    let axis = if maxw == rw {
320        Axis::R
321    } else if maxw == gw {
322        Axis::G
323    } else {
324        Axis::B
325    };
326    let a = axis as usize;
327    let (start, end) = (vbox.lo[a], vbox.hi[a]);
328
329    // Running population along the axis. Out-of-range slots read as 0,
330    // standing in for the `undefined` holes of the JS array.
331    let mut partial = vec![0u64; BINS];
332    let mut total = 0;
333    for i in start..=end {
334        let mut slab = vbox.clone();
335        slab.lo[a] = i;
336        slab.hi[a] = i;
337        total += VBox::new(slab.lo, slab.hi, histo).count;
338        partial[i as usize] = total;
339    }
340    let ps = |i: i32| if (start..=end).contains(&i) { partial[i as usize] } else { 0 };
341    let lookahead = |i: i32| total - ps(i);
342
343    let i = (start..=end).find(|&i| 2 * ps(i) > total)?;
344    let (left, right) = (i - start, end - i);
345    let mut d2 = if left <= right {
346        (end - 1).min((i as f64 + right as f64 / 2.0) as i32)
347    } else {
348        start.max((i as f64 - 1.0 - left as f64 / 2.0) as i32)
349    };
350    while ps(d2) == 0 {
351        d2 += 1;
352    }
353    let mut count2 = lookahead(d2);
354    while count2 == 0 && ps(d2 - 1) != 0 {
355        d2 -= 1;
356        count2 = lookahead(d2);
357    }
358
359    let mut hi1 = vbox.hi;
360    let mut lo2 = vbox.lo;
361    hi1[a] = d2;
362    lo2[a] = d2 + 1;
363    Some((VBox::new(vbox.lo, hi1, histo), Some(VBox::new(lo2, vbox.hi, histo))))
364}
365
366/// Resamples `src` to `dw` x `dh` with one bilinear sample per output pixel,
367/// sampling at pixel centres and clamping at the edges.
368fn bilinear(src: &RgbaImage, dw: u32, dh: u32) -> RgbaImage {
369    let (sw, sh) = src.dimensions();
370    let (sx, sy) = (sw as f64 / dw as f64, sh as f64 / dh as f64);
371    RgbaImage::from_fn(dw, dh, |x, y| {
372        let fx = ((x as f64 + 0.5) * sx - 0.5).clamp(0.0, (sw - 1) as f64);
373        let fy = ((y as f64 + 0.5) * sy - 0.5).clamp(0.0, (sh - 1) as f64);
374        let (x0, y0) = (fx.floor() as u32, fy.floor() as u32);
375        let (x1, y1) = ((x0 + 1).min(sw - 1), (y0 + 1).min(sh - 1));
376        let (tx, ty) = (fx - x0 as f64, fy - y0 as f64);
377        let p = |x, y| src.get_pixel(x, y).0.map(f64::from);
378        let (a, b, c, d) = (p(x0, y0), p(x1, y0), p(x0, y1), p(x1, y1));
379        image::Rgba(std::array::from_fn(|i| {
380            let top = a[i] + (b[i] - a[i]) * tx;
381            let bottom = c[i] + (d[i] - c[i]) * tx;
382            (top + (bottom - top) * ty).round() as u8
383        }))
384    })
385}
386
387/// The picture Eagle analyses when a palette is refreshed: the original
388/// when the item has `noThumbnail` (a small SVG is rendered at its own size),
389/// otherwise the thumbnail file. `None` when there is nothing to read.
390pub fn refresh_source(item: &crate::model::Item, root: &std::path::Path) -> Option<DynamicImage> {
391    use crate::thumbnail;
392    if item.no_thumbnail == Some(true) {
393        let raw = item.original_path(root);
394        if item.ext.eq_ignore_ascii_case("svg") {
395            return thumbnail::render_svg(&std::fs::read(raw).ok()?, None).ok();
396        }
397        return thumbnail::decode(&raw).ok();
398    }
399    thumbnail::decode(&item.thumbnail_path(root)?).ok()
400}
401
402impl crate::Library {
403    /// Recomputes an item's palette (Eagle's "Refresh Palette") with the
404    /// same median cut as import, from [`refresh_source`]. Only `palettes`
405    /// changes, in place, so the item's key order stays the same. Audio is
406    /// skipped, as Eagle's analyser skips it; the result is then false.
407    pub fn refresh_palette(&mut self, id: &str) -> crate::Result<bool> {
408        let mut item = self.item(id).cloned().ok_or_else(|| self.invalid(format!("no item {id}")))?;
409        if crate::media::is_audio(&item.ext) {
410            return Ok(false);
411        }
412        let img = refresh_source(&item, self.root()).ok_or_else(|| self.invalid(format!("nothing to analyse for {}.{}", item.name, item.ext)))?;
413        item.palettes = extract(&img);
414        self.update_item(item).map(|()| true)
415    }
416}
417
418#[cfg(test)]
419mod tests {
420    use super::*;
421
422    fn solid(rgba: [u8; 4], n: usize) -> Vec<u8> {
423        rgba.repeat(n)
424    }
425
426    #[test]
427    fn ratios_follow_eagle_rounding() {
428        let s = Swatches { colors: vec![([1, 2, 3], 4517), ([4, 5, 6], 222), ([7, 8, 9], 20)], pixel_count: 10000 };
429        let p = to_palettes(s);
430        let ratios: Vec<String> = p.iter().map(|c| c.ratio.to_string()).collect();
431        // 45.17 truncates to 45, 2.22 keeps its decimals, 0.2 is dropped.
432        assert_eq!(ratios, ["45", "2.22"]);
433    }
434
435    #[test]
436    fn whole_small_ratios_print_without_decimals() {
437        assert_eq!(js_number(3.0).to_string(), "3");
438        assert_eq!(js_number(0.74).to_string(), "0.74");
439    }
440
441    #[test]
442    fn transparent_pixels_are_ignored() {
443        let mut buf = solid([255, 0, 0, 255], 100);
444        buf.extend(solid([0, 0, 255, 0], 100));
445        let p = from_rgba(&buf);
446        assert_eq!(p.len(), 1);
447        assert_eq!(p[0].ratio.to_string(), "100");
448    }
449
450    #[test]
451    fn single_pixel_hits_the_bounds_quirk() {
452        // The only value lowers each minimum and never sets a maximum, so the
453        // box is empty and Eagle stores no palette.
454        assert!(from_rgba(&solid([40, 40, 40, 255], 1)).is_empty());
455        // A second pixel of the same colour raises the maximum.
456        let p = from_rgba(&solid([40, 40, 40, 255], 2));
457        assert_eq!(p.len(), 1);
458        assert_eq!((p[0].color, p[0].ratio.to_string()), ([44, 44, 44], "100".to_string()));
459    }
460
461    #[test]
462    fn two_colours_split_evenly() {
463        let mut buf = solid([250, 10, 10, 255], 60);
464        buf.extend(solid([10, 10, 250, 255], 40));
465        let p = from_rgba(&buf);
466        let got: Vec<_> = p.iter().map(|c| (c.color, c.ratio.to_string())).collect();
467        assert_eq!(got, [([252, 12, 12], "60".to_string()), ([12, 12, 252], "40".to_string())]);
468    }
469
470    #[test]
471    fn skip_tracks_buffer_size() {
472        assert_eq!(skip_for(360 * 270 * 4), 10);
473        assert_eq!(skip_for(99_999), 1);
474        assert_eq!(skip_for(100_000_000), 10000);
475    }
476}