Add color mapping
This commit is contained in:
Schluffe
@@ -0,0 +1,239 @@
|
||||
//! 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 apply_rgba(&self, rgba: &[u8]) -> Vec<u8> {
|
||||
let lookup: HashMap<_, _> = self
|
||||
.0
|
||||
.iter()
|
||||
.flat_map(|entry| entry.members.iter().map(move |&rgb| (rgb, entry.target)))
|
||||
.collect();
|
||||
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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user