243 lines
8.1 KiB
Rust
243 lines
8.1 KiB
Rust
//! Palette grouping and simultaneous replacement using original RGB values.
|
|
use super::Swatch;
|
|
use std::collections::HashMap;
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub struct ColorGroup {
|
|
pub rgb: [u8; 3],
|
|
pub count: u64,
|
|
pub members: Vec<[u8; 3]>,
|
|
}
|
|
|
|
/// Most frequent colors seed groups. Every member is within `tolerance` of
|
|
/// its seed in each RGB channel; nearby shades cannot form unbounded chains.
|
|
/// Spatial buckets keep large palettes from requiring an all-pairs search.
|
|
pub fn group_palette(palette: &[Swatch], tolerance: u8) -> Vec<ColorGroup> {
|
|
if tolerance == 0 {
|
|
return palette
|
|
.iter()
|
|
.map(|s| ColorGroup {
|
|
rgb: s.rgb,
|
|
count: s.count,
|
|
members: vec![s.rgb],
|
|
})
|
|
.collect();
|
|
}
|
|
let mut sorted: Vec<_> = palette.iter().collect();
|
|
sorted.sort_by_key(|s| (std::cmp::Reverse(s.count), s.rgb));
|
|
let mut groups: Vec<ColorGroup> = Vec::new();
|
|
let mut buckets = HashMap::<[i16; 3], usize>::new();
|
|
let side = i16::from(tolerance) + 1;
|
|
for swatch in sorted {
|
|
let bucket = swatch.rgb.map(|v| i16::from(v) / side);
|
|
let mut best = None;
|
|
for x in -1..=1 {
|
|
for y in -1..=1 {
|
|
for z in -1..=1 {
|
|
if let Some(&index) =
|
|
buckets.get(&[bucket[0] + x, bucket[1] + y, bucket[2] + z])
|
|
{
|
|
let seed = groups[index].rgb;
|
|
let delta =
|
|
std::array::from_fn::<_, 3, _>(|i| seed[i].abs_diff(swatch.rgb[i]));
|
|
if delta.iter().all(|&d| d <= tolerance) {
|
|
let distance: u32 = delta.iter().map(|&d| u32::from(d).pow(2)).sum();
|
|
let candidate = (distance, index);
|
|
if best.is_none_or(|old| candidate < old) {
|
|
best = Some(candidate);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
if let Some((_, index)) = best {
|
|
groups[index].count += swatch.count;
|
|
groups[index].members.push(swatch.rgb);
|
|
} else {
|
|
buckets.insert(bucket, groups.len());
|
|
groups.push(ColorGroup {
|
|
rgb: swatch.rgb,
|
|
count: swatch.count,
|
|
members: vec![swatch.rgb],
|
|
});
|
|
}
|
|
}
|
|
groups.sort_by_key(|g| (std::cmp::Reverse(g.count), g.rgb));
|
|
groups
|
|
}
|
|
|
|
impl ColorGroup {
|
|
pub fn highlight_rgba(&self, rgba: &[u8]) -> Vec<u8> {
|
|
let members: std::collections::HashSet<_> = self.members.iter().copied().collect();
|
|
let mut output = rgba.to_vec();
|
|
for pixel in output.as_chunks_mut::<4>().0 {
|
|
if pixel[3] == 0 {
|
|
continue;
|
|
}
|
|
if members.contains(&[pixel[0], pixel[1], pixel[2]]) {
|
|
pixel[..3].copy_from_slice(&[255, 0, 255]);
|
|
} else {
|
|
for channel in &mut pixel[..3] {
|
|
*channel /= 4;
|
|
}
|
|
}
|
|
}
|
|
output
|
|
}
|
|
}
|
|
|
|
pub struct ColorMapping {
|
|
pub source: [u8; 3],
|
|
pub target: [u8; 3],
|
|
pub members: Vec<[u8; 3]>,
|
|
}
|
|
|
|
#[derive(Default)]
|
|
pub struct ColorMappings(pub Vec<ColorMapping>);
|
|
|
|
impl ColorMappings {
|
|
pub fn insert_group(&mut self, group: &ColorGroup) {
|
|
let members: std::collections::HashSet<_> = group.members.iter().copied().collect();
|
|
// The latest explicit mapping owns overlapping colors. Keep unaffected
|
|
// members of earlier mappings, and remove entries that become empty.
|
|
for entry in &mut self.0 {
|
|
entry.members.retain(|rgb| !members.contains(rgb));
|
|
}
|
|
self.0.retain(|entry| !entry.members.is_empty());
|
|
self.0.push(ColorMapping {
|
|
source: group.rgb,
|
|
target: group.rgb,
|
|
members: group.members.clone(),
|
|
});
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub fn insert(&mut self, source: [u8; 3], target: [u8; 3]) {
|
|
self.insert_group(&ColorGroup {
|
|
rgb: source,
|
|
count: 1,
|
|
members: vec![source],
|
|
});
|
|
self.0.last_mut().unwrap().target = target;
|
|
}
|
|
|
|
pub fn lookup(&self) -> HashMap<[u8; 3], [u8; 3]> {
|
|
self.0
|
|
.iter()
|
|
.flat_map(|entry| entry.members.iter().map(move |&rgb| (rgb, entry.target)))
|
|
.collect()
|
|
}
|
|
|
|
pub fn apply_rgba(&self, rgba: &[u8]) -> Vec<u8> {
|
|
let lookup = self.lookup();
|
|
let mut output = rgba.to_vec();
|
|
for pixel in output.as_chunks_mut::<4>().0 {
|
|
if pixel[3] != 0
|
|
&& let Some(target) = lookup.get(&[pixel[0], pixel[1], pixel[2]])
|
|
{
|
|
pixel[..3].copy_from_slice(target);
|
|
}
|
|
}
|
|
output
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
fn swatch(rgb: [u8; 3], count: u64) -> Swatch {
|
|
Swatch {
|
|
rgb,
|
|
count,
|
|
opaque_count: count,
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn grouping_combines_counts_without_chaining_and_keeps_exact_mode() {
|
|
let palette = vec![
|
|
swatch([100, 0, 0], 100),
|
|
swatch([105, 0, 0], 3),
|
|
swatch([110, 0, 0], 2),
|
|
swatch([0, 255, 0], 1),
|
|
];
|
|
let groups = group_palette(&palette, 5);
|
|
assert_eq!(groups.len(), 3);
|
|
assert_eq!(groups[0].members, [[100, 0, 0], [105, 0, 0]]);
|
|
assert_eq!(groups[0].count, 103);
|
|
assert_eq!(groups.iter().map(|g| g.count).sum::<u64>(), 106);
|
|
assert_eq!(group_palette(&palette, 0).len(), 4);
|
|
let mut reversed = palette;
|
|
reversed.reverse();
|
|
assert_eq!(group_palette(&reversed, 5)[0].members, groups[0].members);
|
|
}
|
|
|
|
#[test]
|
|
fn group_mapping_uses_fixed_members_and_latest_mapping_owns_overlap() {
|
|
let groups = group_palette(&[swatch([255, 0, 0], 10), swatch([252, 2, 1], 1)], 4);
|
|
let mut mappings = ColorMappings::default();
|
|
mappings.insert_group(&groups[0]);
|
|
mappings.0[0].target = [0, 0, 255];
|
|
assert_eq!(mappings.apply_rgba(&[252, 2, 1, 128]), [0, 0, 255, 128]);
|
|
mappings.insert([252, 2, 1], [0, 255, 0]);
|
|
assert_eq!(mappings.0[0].members, [[255, 0, 0]]);
|
|
assert_eq!(
|
|
mappings.apply_rgba(&[252, 2, 1, 255, 255, 0, 0, 255]),
|
|
[0, 255, 0, 255, 0, 0, 255, 255]
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn highlight_matches_group_members_and_preserves_transparency() {
|
|
let group = ColorGroup {
|
|
rgb: [250, 0, 0],
|
|
count: 2,
|
|
members: vec![[250, 0, 0], [253, 2, 0]],
|
|
};
|
|
assert_eq!(
|
|
group.highlight_rgba(&[253, 2, 0, 128, 0, 0, 252, 255, 250, 0, 0, 0]),
|
|
[255, 0, 255, 128, 0, 0, 63, 255, 250, 0, 0, 0]
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn dense_palette_preserves_every_color_and_respects_tolerance() {
|
|
let palette: Vec<_> = (0..=255)
|
|
.flat_map(|r| (0..=255).step_by(8).map(move |g| swatch([r, g, 128], 1)))
|
|
.collect();
|
|
let groups = group_palette(&palette, 8);
|
|
let members: std::collections::HashSet<_> = groups
|
|
.iter()
|
|
.flat_map(|g| g.members.iter().copied())
|
|
.collect();
|
|
assert_eq!(members.len(), palette.len());
|
|
assert_eq!(
|
|
groups.iter().map(|g| g.count).sum::<u64>(),
|
|
palette.len() as u64
|
|
);
|
|
for group in groups {
|
|
assert!(
|
|
group
|
|
.members
|
|
.iter()
|
|
.all(|rgb| (0..3).all(|i| rgb[i].abs_diff(group.rgb[i]) <= 8))
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn replacements_are_exact_simultaneous_and_preserve_alpha() {
|
|
let mut mappings = ColorMappings::default();
|
|
mappings.insert([255, 0, 0], [0, 255, 0]);
|
|
mappings.insert([0, 255, 0], [255, 0, 0]);
|
|
let source = [255, 0, 0, 255, 0, 255, 0, 128, 254, 0, 0, 255, 255, 0, 0, 0];
|
|
assert_eq!(
|
|
mappings.apply_rgba(&source),
|
|
[0, 255, 0, 255, 255, 0, 0, 128, 254, 0, 0, 255, 255, 0, 0, 0]
|
|
);
|
|
assert_eq!(ColorMappings::default().apply_rgba(&source), source);
|
|
}
|
|
}
|