mirror of
https://github.com/storytold/vectorcraft.git
synced 2026-11-09 16:38:07 +00:00
Merge pull request #5467 from vangel76/perf/trace-denoise-union-find
Image Trace: noise removal keeps its regions in u32 arrays (#5462)
This commit is contained in:
BFlat
@@ -13,94 +13,286 @@ pub fn denoise(labels: &mut [u16], w: usize, h: usize, min_area: usize) {
|
||||
/// clears specks, not a colour's only regions. Smaller regions of it go as any other.
|
||||
pub(crate) fn denoise_keeping(labels: &mut [u16], w: usize, h: usize, min_area: usize, keep: &[bool]) {
|
||||
use std::cmp::Reverse;
|
||||
use std::collections::{BTreeMap, BinaryHeap};
|
||||
use std::collections::BinaryHeap;
|
||||
|
||||
if min_area <= 1 || w == 0 || h == 0 || w.checked_mul(h) != Some(labels.len()) {
|
||||
// Every index below is a `u32`, so a pixel costs 8 bytes and a region 18 instead of a `Vec<usize>` per region. An image
|
||||
// of `u32::MAX` pixels or more (far above the decoder's cap) is left as it is.
|
||||
const NONE: u32 = u32::MAX;
|
||||
if min_area <= 1 || w == 0 || h == 0 || w.checked_mul(h) != Some(labels.len()) || labels.len() >= NONE as usize {
|
||||
return;
|
||||
}
|
||||
struct Region {
|
||||
label: u16,
|
||||
pixels: Vec<usize>,
|
||||
}
|
||||
// Region membership is kept current after every merge. In particular, no stale `seen`
|
||||
// flags can exclude previously recoloured pixels from a subsequent connected component.
|
||||
let mut at = vec![usize::MAX; labels.len()];
|
||||
let mut regions: Vec<Region> = Vec::new();
|
||||
// Region of each pixel as first found (`NONE` for transparent). A merge does not touch it: the region a pixel
|
||||
// belongs to now is `find(at[i])`.
|
||||
let mut at = vec![NONE; labels.len()];
|
||||
// Pixels of a region form a singly linked list through `next`, so absorbing a region is one pointer write.
|
||||
let mut next = vec![NONE; labels.len()];
|
||||
let mut head: Vec<u32> = Vec::new();
|
||||
let mut tail: Vec<u32> = Vec::new();
|
||||
let mut size: Vec<u32> = Vec::new();
|
||||
let mut colour: Vec<u16> = Vec::new();
|
||||
// Union-find over region ids: a region absorbed into another points at it.
|
||||
let mut parent: Vec<u32> = Vec::new();
|
||||
let neighbours = |i: usize| {
|
||||
let (x, y) = (i % w, i / w);
|
||||
[(x > 0).then(|| i - 1), (x + 1 < w).then(|| i + 1), (y > 0).then(|| i - w), (y + 1 < h).then(|| i + w)]
|
||||
};
|
||||
for start in 0..labels.len() {
|
||||
if at[start] != usize::MAX || labels[start] == TRANSPARENT {
|
||||
if at[start] != NONE || labels[start] == TRANSPARENT {
|
||||
continue;
|
||||
}
|
||||
let id = regions.len();
|
||||
let id = head.len() as u32;
|
||||
let label = labels[start];
|
||||
let mut pixels = vec![start];
|
||||
at[start] = id;
|
||||
let mut cursor = 0;
|
||||
while let Some(&i) = pixels.get(cursor) {
|
||||
for j in neighbours(i).into_iter().flatten() {
|
||||
if at[j] == usize::MAX && labels[j] == label {
|
||||
let mut last = start;
|
||||
let mut count = 1u32;
|
||||
let mut cursor = start;
|
||||
loop {
|
||||
for j in neighbours(cursor).into_iter().flatten() {
|
||||
if at[j] == NONE && labels[j] == label {
|
||||
at[j] = id;
|
||||
pixels.push(j);
|
||||
next[last] = j as u32;
|
||||
last = j;
|
||||
count += 1;
|
||||
}
|
||||
}
|
||||
cursor += 1;
|
||||
match next[cursor] {
|
||||
NONE => break,
|
||||
n => cursor = n as usize,
|
||||
}
|
||||
}
|
||||
regions.push(Region { label, pixels });
|
||||
head.push(start as u32);
|
||||
tail.push(last as u32);
|
||||
size.push(count);
|
||||
colour.push(label);
|
||||
parent.push(id);
|
||||
}
|
||||
let mut pending: BinaryHeap<_> = regions.iter().enumerate().map(|(id, r)| Reverse((r.pixels.len(), id))).collect();
|
||||
while let Some(Reverse((area, id))) = pending.pop() {
|
||||
if area >= min_area {
|
||||
break;
|
||||
fn find(parent: &mut [u32], mut a: u32) -> u32 {
|
||||
while parent[a as usize] != a {
|
||||
let up = parent[parent[a as usize] as usize];
|
||||
parent[a as usize] = up;
|
||||
a = up;
|
||||
}
|
||||
if regions[id].pixels.len() != area {
|
||||
a
|
||||
}
|
||||
// Only regions below `min_area` are ever taken from the queue; a larger one ends the loop.
|
||||
let mut pending: BinaryHeap<_> =
|
||||
(0..size.len() as u32).filter(|&id| (size[id as usize] as usize) < min_area).map(|id| Reverse((size[id as usize], id))).collect();
|
||||
// Scratch for one merge: `(other region, shared pixel edges)` and `(colour, shared pixel edges)`.
|
||||
let mut borders: Vec<(u32, usize)> = Vec::new();
|
||||
let mut by_colour: Vec<(u16, usize)> = Vec::new();
|
||||
while let Some(Reverse((area, id))) = pending.pop() {
|
||||
if size[id as usize] != area {
|
||||
continue; // absorbed region or an old heap entry for a region that grew
|
||||
}
|
||||
if area >= RARE_ISLAND && keep.get(regions[id].label as usize).copied().unwrap_or(false) {
|
||||
if area as usize >= RARE_ISLAND && keep.get(colour[id as usize] as usize).copied().unwrap_or(false) {
|
||||
continue;
|
||||
}
|
||||
let mut borders: BTreeMap<usize, usize> = BTreeMap::new();
|
||||
for &i in ®ions[id].pixels {
|
||||
for j in neighbours(i).into_iter().flatten() {
|
||||
let other = at[j];
|
||||
if other != id && other != usize::MAX {
|
||||
*borders.entry(other).or_default() += 1;
|
||||
borders.clear();
|
||||
// Pairs are folded together whenever the list doubles, so a big region keeps it short.
|
||||
let mut limit = 4096;
|
||||
let mut pixel = head[id as usize];
|
||||
while pixel != NONE {
|
||||
for j in neighbours(pixel as usize).into_iter().flatten() {
|
||||
if at[j] != NONE {
|
||||
let other = find(&mut parent, at[j]);
|
||||
if other != id {
|
||||
borders.push((other, 1));
|
||||
}
|
||||
}
|
||||
}
|
||||
pixel = next[pixel as usize];
|
||||
if borders.len() >= limit {
|
||||
merge_counts(&mut borders);
|
||||
limit = limit.max(2 * borders.len());
|
||||
}
|
||||
}
|
||||
let mut colours: BTreeMap<u16, usize> = BTreeMap::new();
|
||||
for (&other, &border) in &borders {
|
||||
*colours.entry(regions[other].label).or_default() += border;
|
||||
}
|
||||
let Some((&colour, _)) = colours.iter().max_by_key(|(label, border)| (**border, Reverse(**label))) else {
|
||||
merge_counts(&mut borders);
|
||||
by_colour.clear();
|
||||
by_colour.extend(borders.iter().map(|&(other, border)| (colour[other as usize], border)));
|
||||
merge_counts(&mut by_colour);
|
||||
let Some(&(chosen, _)) = by_colour.iter().max_by_key(|(label, border)| (*border, Reverse(*label))) else {
|
||||
continue; // an opaque island surrounded only by transparency stays intact
|
||||
};
|
||||
// All adjacent regions of the chosen colour become connected by this merge. Absorb
|
||||
// them together, keeping the largest one's storage to avoid repeatedly moving it.
|
||||
// them together, keeping the largest one as the survivor.
|
||||
let Some(target) = borders
|
||||
.keys()
|
||||
.copied()
|
||||
.filter(|&other| regions[other].label == colour)
|
||||
.max_by_key(|&other| (regions[other].pixels.len(), Reverse(other)))
|
||||
.iter()
|
||||
.map(|&(other, _)| other)
|
||||
.filter(|&other| colour[other as usize] == chosen)
|
||||
.max_by_key(|&other| (size[other as usize], Reverse(other)))
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
let sources: Vec<_> =
|
||||
std::iter::once(id).chain(borders.keys().copied().filter(|&other| other != target && regions[other].label == colour)).collect();
|
||||
for source in sources {
|
||||
let pixels = std::mem::take(&mut regions[source].pixels);
|
||||
for &i in &pixels {
|
||||
at[i] = target;
|
||||
}
|
||||
regions[target].pixels.extend(pixels);
|
||||
for source in
|
||||
std::iter::once(id).chain(borders.iter().map(|&(other, _)| other).filter(|&other| other != target && colour[other as usize] == chosen))
|
||||
{
|
||||
next[tail[target as usize] as usize] = head[source as usize];
|
||||
tail[target as usize] = tail[source as usize];
|
||||
size[target as usize] += size[source as usize];
|
||||
size[source as usize] = 0;
|
||||
parent[source as usize] = target;
|
||||
}
|
||||
if (size[target as usize] as usize) < min_area {
|
||||
pending.push(Reverse((size[target as usize], target)));
|
||||
}
|
||||
pending.push(Reverse((regions[target].pixels.len(), target)));
|
||||
}
|
||||
for region in regions {
|
||||
for i in region.pixels {
|
||||
labels[i] = region.label;
|
||||
for (i, ®ion) in at.iter().enumerate() {
|
||||
if region != NONE {
|
||||
labels[i] = colour[find(&mut parent, region) as usize];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Sort `(key, count)` pairs and add up the counts of equal keys.
|
||||
fn merge_counts<K: Ord + Copy>(pairs: &mut Vec<(K, usize)>) {
|
||||
pairs.sort_unstable_by_key(|p| p.0);
|
||||
pairs.dedup_by(|b, a| {
|
||||
let same = a.0 == b.0;
|
||||
if same {
|
||||
a.1 += b.1;
|
||||
}
|
||||
same
|
||||
});
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// The first implementation (a pixel `Vec` per region, `BTreeMap` borders), kept to check that the compact one merges
|
||||
/// in exactly the same order.
|
||||
fn denoise_reference(labels: &mut [u16], w: usize, h: usize, min_area: usize, keep: &[bool]) {
|
||||
use std::cmp::Reverse;
|
||||
use std::collections::{BTreeMap, BinaryHeap};
|
||||
|
||||
if min_area <= 1 || w == 0 || h == 0 || w.checked_mul(h) != Some(labels.len()) {
|
||||
return;
|
||||
}
|
||||
struct Region {
|
||||
label: u16,
|
||||
pixels: Vec<usize>,
|
||||
}
|
||||
// Region membership is kept current after every merge. In particular, no stale `seen`
|
||||
// flags can exclude previously recoloured pixels from a subsequent connected component.
|
||||
let mut at = vec![usize::MAX; labels.len()];
|
||||
let mut regions: Vec<Region> = Vec::new();
|
||||
let neighbours = |i: usize| {
|
||||
let (x, y) = (i % w, i / w);
|
||||
[(x > 0).then(|| i - 1), (x + 1 < w).then(|| i + 1), (y > 0).then(|| i - w), (y + 1 < h).then(|| i + w)]
|
||||
};
|
||||
for start in 0..labels.len() {
|
||||
if at[start] != usize::MAX || labels[start] == TRANSPARENT {
|
||||
continue;
|
||||
}
|
||||
let id = regions.len();
|
||||
let label = labels[start];
|
||||
let mut pixels = vec![start];
|
||||
at[start] = id;
|
||||
let mut cursor = 0;
|
||||
while let Some(&i) = pixels.get(cursor) {
|
||||
for j in neighbours(i).into_iter().flatten() {
|
||||
if at[j] == usize::MAX && labels[j] == label {
|
||||
at[j] = id;
|
||||
pixels.push(j);
|
||||
}
|
||||
}
|
||||
cursor += 1;
|
||||
}
|
||||
regions.push(Region { label, pixels });
|
||||
}
|
||||
let mut pending: BinaryHeap<_> = regions.iter().enumerate().map(|(id, r)| Reverse((r.pixels.len(), id))).collect();
|
||||
while let Some(Reverse((area, id))) = pending.pop() {
|
||||
if area >= min_area {
|
||||
break;
|
||||
}
|
||||
if regions[id].pixels.len() != area {
|
||||
continue; // absorbed region or an old heap entry for a region that grew
|
||||
}
|
||||
if area >= RARE_ISLAND && keep.get(regions[id].label as usize).copied().unwrap_or(false) {
|
||||
continue;
|
||||
}
|
||||
let mut borders: BTreeMap<usize, usize> = BTreeMap::new();
|
||||
for &i in ®ions[id].pixels {
|
||||
for j in neighbours(i).into_iter().flatten() {
|
||||
let other = at[j];
|
||||
if other != id && other != usize::MAX {
|
||||
*borders.entry(other).or_default() += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
let mut colours: BTreeMap<u16, usize> = BTreeMap::new();
|
||||
for (&other, &border) in &borders {
|
||||
*colours.entry(regions[other].label).or_default() += border;
|
||||
}
|
||||
let Some((&colour, _)) = colours.iter().max_by_key(|(label, border)| (**border, Reverse(**label))) else {
|
||||
continue; // an opaque island surrounded only by transparency stays intact
|
||||
};
|
||||
// All adjacent regions of the chosen colour become connected by this merge. Absorb
|
||||
// them together, keeping the largest one's storage to avoid repeatedly moving it.
|
||||
let Some(target) = borders
|
||||
.keys()
|
||||
.copied()
|
||||
.filter(|&other| regions[other].label == colour)
|
||||
.max_by_key(|&other| (regions[other].pixels.len(), Reverse(other)))
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
let sources: Vec<_> =
|
||||
std::iter::once(id).chain(borders.keys().copied().filter(|&other| other != target && regions[other].label == colour)).collect();
|
||||
for source in sources {
|
||||
let pixels = std::mem::take(&mut regions[source].pixels);
|
||||
for &i in &pixels {
|
||||
at[i] = target;
|
||||
}
|
||||
regions[target].pixels.extend(pixels);
|
||||
}
|
||||
pending.push(Reverse((regions[target].pixels.len(), target)));
|
||||
}
|
||||
for region in regions {
|
||||
for i in region.pixels {
|
||||
labels[i] = region.label;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn noisy(w: usize, h: usize, colours: u64, transparent_one_in: u64, seed: u64) -> Vec<u16> {
|
||||
let mut state = seed;
|
||||
let mut next = move || {
|
||||
state ^= state << 13;
|
||||
state ^= state >> 7;
|
||||
state ^= state << 17;
|
||||
state
|
||||
};
|
||||
(0..w * h).map(|_| if transparent_one_in > 0 && next() % transparent_one_in == 0 { TRANSPARENT } else { (next() % colours) as u16 }).collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compact_denoise_matches_the_reference_merge_order() {
|
||||
for (seed, (w, h, colours, transparent)) in
|
||||
[(31, 17, 3, 0), (64, 64, 2, 0), (50, 70, 5, 9), (97, 41, 8, 0), (120, 90, 4, 5), (7, 300, 3, 0), (1, 40, 3, 0), (40, 1, 3, 0)]
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
{
|
||||
let base = noisy(w, h, colours, transparent, seed as u64 + 1);
|
||||
for min_area in [2, 3, 5, 9, 25, 80, 400, 100_000] {
|
||||
for keep in [&[][..], &[false, true, false, true][..]] {
|
||||
let (mut a, mut b) = (base.clone(), base.clone());
|
||||
denoise_reference(&mut a, w, h, min_area, keep);
|
||||
denoise_keeping(&mut b, w, h, min_area, keep);
|
||||
assert_eq!(a, b, "{w}x{h}, {colours} colours, min_area {min_area}, keep {keep:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A minimum so large that regions keep merging into big ones.
|
||||
#[test]
|
||||
fn compact_denoise_matches_the_reference_with_a_large_min_area() {
|
||||
let (w, h) = (320, 320);
|
||||
let base = noisy(w, h, 6, 0, 99);
|
||||
let (mut a, mut b) = (base.clone(), base);
|
||||
denoise_reference(&mut a, w, h, 50_000, &[]);
|
||||
denoise_keeping(&mut b, w, h, 50_000, &[]);
|
||||
assert_eq!(a, b);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user