Skip to main content

crabshell_render/
colors.rs

1//! The main colors of an image, by median cut, as for a theme taken from a
2//! wallpaper.
3
4use std::cmp::Reverse;
5use std::ops::Range;
6use std::path::Path;
7
8use crabshell_ui::{Color, Fit};
9
10use crate::icons::halve;
11use crate::images::{self, Picture};
12
13/// The deepest cut: 2^8 boxes, 256 colors.
14const MAX_DEPTH: u8 = 8;
15
16/// The main colors of the image file at `path`, most common first.
17///
18/// The file is decoded as [`Image`](crabshell_ui::Image) decodes it, under
19/// the same caps, then halved until its longer side is at most `size`
20/// pixels; 0 keeps its size. Its pixels are split by median cut into up to
21/// 2^`depth` boxes, with `depth` clamped to 0..=8: each time, the box with
22/// the widest range on one channel is split at its median on that channel.
23/// Each box gives the mean of its pixels, opaque. Pixels with alpha under
24/// 128 are skipped. Boxes with as many pixels as each other go by their
25/// color's bytes, so a result is the same on every run. An image with
26/// fewer distinct colors than boxes gives each of them once.
27///
28/// Decoding a wallpaper takes tens of milliseconds, so call it off the
29/// loop's thread.
30///
31/// # Errors
32///
33/// If the file is missing or isn't a regular file, isn't a PNG, JPEG,
34/// GIF, WebP or SVG, is past the caps, or doesn't decode.
35pub fn quantize_colors(path: &Path, depth: u8, size: u32) -> Result<Vec<Color>, String> {
36  let image = match images::decode(path)? {
37    Picture::Pixels(mut image) => {
38      while size > 0 && image.width().max(image.height()) > size {
39        image = halve(image.as_ref())?;
40      }
41      image
42    }
43    Picture::Svg(tree) => {
44      let (width, height) = (tree.size().width(), tree.size().height());
45      let shrink = match size {
46        0 => 1.0,
47        size => (size as f32 / width.max(height)).min(1.0),
48      };
49      let (width, height) = (
50        (width * shrink).ceil() as u32,
51        (height * shrink).ceil() as u32,
52      );
53      images::check(width, height)?;
54      images::svg(&tree, width, height, Fit::Fill)?
55    }
56  };
57  let pixels = image
58    .pixels()
59    .iter()
60    .filter(|pixel| pixel.alpha() >= 128)
61    .map(|pixel| {
62      let color = pixel.demultiply();
63      [color.red(), color.green(), color.blue()]
64    })
65    .collect();
66  Ok(median_cut(pixels, depth))
67}
68
69/// A box of the cut: a run of the pixels, and its widest channel.
70struct Cut {
71  run: Range<usize>,
72  channel: usize,
73  /// How far that channel's values reach: 0 for a box of one color.
74  spread: u8,
75}
76
77impl Cut {
78  fn new(pixels: &[[u8; 3]], run: Range<usize>) -> Cut {
79    let (mut low, mut high) = ([u8::MAX; 3], [0; 3]);
80    for pixel in &pixels[run.clone()] {
81      for channel in 0..3 {
82        low[channel] = low[channel].min(pixel[channel]);
83        high[channel] = high[channel].max(pixel[channel]);
84      }
85    }
86    let spread = |channel: usize| high[channel].saturating_sub(low[channel]);
87    let channel = (0..3).max_by_key(|&channel| spread(channel)).unwrap_or(0);
88    Cut {
89      run,
90      channel,
91      spread: spread(channel),
92    }
93  }
94}
95
96fn median_cut(mut pixels: Vec<[u8; 3]>, depth: u8) -> Vec<Color> {
97  let mut cuts = Vec::new();
98  if !pixels.is_empty() {
99    cuts.push(Cut::new(&pixels, 0..pixels.len()));
100  }
101  while cuts.len() < 1 << depth.min(MAX_DEPTH) {
102    // None left to split means each box holds one color.
103    let Some(index) = (0..cuts.len())
104      .filter(|&index| cuts[index].spread > 0)
105      .max_by_key(|&index| cuts[index].spread)
106    else {
107      break;
108    };
109    let (run, channel) = (cuts[index].run.clone(), cuts[index].channel);
110    let values = &mut pixels[run.clone()];
111    values.sort_unstable_by_key(|pixel| pixel[channel]);
112    let median = values[values.len() / 2][channel];
113    // Equal values stay on one side, so no color is split between boxes.
114    let mut split = values.partition_point(|pixel| pixel[channel] < median);
115    if split == 0 {
116      split = values.partition_point(|pixel| pixel[channel] <= median);
117    }
118    let middle = run.start + split;
119    cuts[index] = Cut::new(&pixels, run.start..middle);
120    cuts.push(Cut::new(&pixels, middle..run.end));
121  }
122  let mut colors: Vec<(usize, Color)> = cuts
123    .iter()
124    .map(|cut| {
125      let values = &pixels[cut.run.clone()];
126      let count = values.len() as u64;
127      let mut sums = [0_u64; 3];
128      for pixel in values {
129        for (sum, &value) in sums.iter_mut().zip(pixel) {
130          *sum += u64::from(value);
131        }
132      }
133      let [r, g, b] = sums.map(|sum| ((sum + count / 2) / count) as u8);
134      (values.len(), Color::rgb(r, g, b))
135    })
136    .collect();
137  colors.sort_by_key(|&(count, color)| (Reverse(count), [color.r, color.g, color.b]));
138  colors.into_iter().map(|(_, color)| color).collect()
139}