1use 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;
26pub const ANALYSIS_WIDTH: u32 = 360;
28
29pub fn extract(img: &DynamicImage) -> Vec<Palette> {
34 from_rgba(analysis_image(img).as_raw())
35}
36
37pub 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
45pub 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
62fn 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
73struct Swatches {
76 colors: Vec<([u8; 3], u64)>,
77 pixel_count: u64,
78}
79
80fn 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 let ratio = if ratio >= 5.0 { ratio.trunc() } else { ratio };
90 (color, ratio)
91 })
92 .filter(|&(_, ratio)| ratio >= 0.25)
93 .collect();
94 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
115fn 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
125struct 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#[derive(Clone, Copy, PartialEq)]
136enum Axis {
137 R,
138 G,
139 B,
140}
141
142#[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 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
196struct 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
237fn 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 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
282fn 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
309fn 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 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
366fn 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
387pub 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 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 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 assert!(from_rgba(&solid([40, 40, 40, 255], 1)).is_empty());
455 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}