vh/vtracer
1
0
Fork 0
mirror of https://github.com/visioncortex/vtracer.git synced 2026-08-30 17:05:57 -07:00
Code Issues Actions Packages Projects Releases Wiki Activity

Watershed: hierarchy stacking, native cutout, and cached re-cuts

Browse Source 
Three refinements that make the watershed frontend a first-class citizen
of both compositing modes and of interactive tuning:

Stacked mode now stacks for real. Instead of one full-canvas background
plus disjoint regions, the cut emits the merge tree itself: the root
(whole canvas, mean color) first, then progressively finer ancestor
regions, then the final regions on top — the same principle as the color
clustering frontend, just with watershed-born clusters. Sub-pixel gaps
between abutting regions therefore show their common ancestor's color
rather than an unrelated backdrop, and overdraw stays seam-free. A
painted-area budget (3x canvas) keeps pathological persistence chains
from ballooning the stack; the root and final regions are always emitted
so coverage never depends on it.

Cutout is native. The watershed hierarchy already decided every merge, so
the flattened partition reaches the mosaic untouched: merge_diff is 0 for
watershed (the gradient-step re-merge still applies to the color path).
Faces are exactly the cut regions.

Re-cuts are cached. WatershedHierarchy is now public and split into
build(img) — Kruskal, BPT, volume persistence; depends only on the image
— and cut(detail, min_area), which is near-linear: region formation,
graph-level small-basin absorption (region adjacencies, not pixel
sweeps), then the merge tree. Session builds the hierarchy lazily on the
first watershed render and re-cuts it on every watershed_detail or
filter_speckle change: ~25 ms per re-cut vs ~40 ms rebuild on a 1400x775
photo, with the one-shot Frontend::segment path unchanged (build + cut),
so Session output still equals the one-shot pipeline exactly.

Tests: flatten-based stack invariants (solid bottom layer, full coverage,
final regions topmost, exact region counts), hierarchy re-cut == one-shot,
Session re-cut == one-shot across detail changes, and cutout keeping two
regions one gradient step apart that the color path's merge would rejoin.
Watershed goldens re-blessed for the new stack structure.
This commit is contained in:
Chris Tsang
2026-07-27 15:12:12 +01:00
parent 46a1b90ccd
commit c3f56a6339
9 changed files with 571 additions and 286 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CHANGELOG.md
+2 -1
View File
@@ -11,7 +11,8 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
* Binary thresholding: a tunable fixed threshold and Bradley–Roth adaptive thresholding for uneven lighting — CLI `--threshold` / `--adaptive` (`--adaptive-window`, `--adaptive-t`), also on `Config`, Python, and Node.
* Cutout mode merges neighbouring mosaic regions whose colors are within one gradient step — the flattened tessellation no longer keeps the near-identical faces that stacked gradient layering splits a smooth area into.
* Watershed clustering (`--clustering watershed`): an alternative region-forming frontend — a hierarchical watershed by volume on the pixel graph (Cousty et al., TPAMI 2009; Najman, Cousty & Perret, ISMM 2013), cut at a single `--watershed-detail` dial (0..=255, each +25.5 roughly doubles the region count). Content-adaptive regions with no watershed-line pixels; the partition drops straight into both stacked and cutout modes.
* Watershed clustering (`--clustering watershed`): an alternative region-forming frontend — a hierarchical watershed by volume on the pixel graph (Cousty et al., TPAMI 2009; Najman, Cousty & Perret, ISMM 2013), cut at a single `--watershed-detail` dial (0..=255, each +25.5 roughly doubles the region count). Content-adaptive regions with no watershed-line pixels. With `cutout` the partition reaches the mosaic natively (no gradient-step re-merge); with `stacked` the merge tree itself is the stack — coarse ancestors below, refined regions on top, the same principle as color clustering — so sub-pixel gaps show ancestor colors and overdraw stays seam-free.
* `WatershedHierarchy` is public and split into `build` (expensive, depends only on the image) and `cut` (near-instant): `Session` builds it once and re-cuts on every `watershed_detail`/`filter_speckle` change, making the detail slider fully interactive (~25 ms re-cut vs ~40 ms rebuild on a 1400×775 photo).
### Changed
crates/vtracer-cli/src/main.rs
+3 -5
View File
@@ -113,9 +113,7 @@ struct Args {
}
fn parse_segment_length(s: &str) -> Result<f64, String> {
let v: f64 = s
.parse()
.map_err(|_| format!("`{s}` is not a number"))?;
let v: f64 = s.parse().map_err(|_| format!("`{s}` is not a number"))?;
if !(3.5..=10.0).contains(&v) {
return Err(format!("segment length {v} is out of range [3.5, 10]"));
}
@@ -211,8 +209,8 @@ fn build_config(args: &Args) -> Result<Config, String> {
if let Some(text) = &args.palette {
config.palette = parse_palette(text)?;
} else if let Some(path) = &args.palette_file {
let text = std::fs::read_to_string(path)
.map_err(|e| format!("cannot read palette file: {e}"))?;
let text =
std::fs::read_to_string(path).map_err(|e| format!("cannot read palette file: {e}"))?;
config.palette = parse_palette(&text)?;
}
crates/vtracer/src/config.rs
+9 -4
View File
@@ -196,8 +196,8 @@ impl Config {
}
}
/// Speckle filter area (px), applied in the `finish` phase.
fn speckle_area(&self) -> usize {
/// Speckle filter area (px), fed to the frontend.
pub(crate) fn speckle_area(&self) -> usize {
self.filter_speckle * self.filter_speckle
}
@@ -295,8 +295,13 @@ impl Config {
fitter: self.segment_fitter(),
// Rejoin flattened neighbours the gradient layering split:
// clustering itself considers colors within one gradient step
// to be the same region (`deepen_diff`).
merge_diff: self.layer_difference,
// to be the same region (`deepen_diff`). The watershed
// hierarchy already decided every merge, so its partition
// passes to the mosaic untouched.
merge_diff: match self.clustering {
Clustering::Watershed => 0,
_ => self.layer_difference,
},
},
};
crates/vtracer/src/frontend/mod.rs
+1 -1
View File
@@ -16,7 +16,7 @@ mod watershed;
pub use binary::{BinaryFrontend, Threshold};
pub use color_cluster::ColorClusterFrontend;
pub use watershed::WatershedFrontend;
pub use watershed::{WatershedFrontend, WatershedHierarchy};
use visioncortex::ColorImage;
crates/vtracer/src/frontend/watershed.rs
+360 -221
View File
@@ -2,21 +2,34 @@
//!
//! The image is treated as a 4-adjacency edge-weighted graph (edge weight =
//! color difference between the two pixels; no gradient image is built). On it
//! we compute the watershed hierarchy by **volume extinction** and cut it at a
//! detail level, following:
//! we compute the watershed hierarchy by **volume extinction**, following:
//!
//! * Cousty, Bertrand, Najman, Couprie, *Watershed Cuts: Minimum Spanning
//! Forests and the Drop of Water Principle*, IEEE TPAMI 31(8), 2009.
//! * Najman, Cousty, Perret, *Playing with Kruskal: Algorithms for
//! Morphological Trees in Edge-Weighted Graphs*, ISMM 2013.
//!
//! Pipeline: Kruskal over counting-sorted edges builds the binary partition
//! tree (a flat `parents` array, leaves `0..n`, internal nodes created in
//! altitude order). A leaves-to-root pass computes each subtree's area and
//! volume; each internal node's *persistence* (the volume of the smaller of
//! the two merged regions) becomes the saliency of its MST edge. Cutting the
//! hierarchy at level λ is then single-linkage over MST edges with
//! persistence ≤ λ — every pixel gets a label, no watershed-line pixels.
//! The work is split in two so the expensive part can be cached (see
//! [`crate::Session`]):
//!
//! * [`WatershedHierarchy::build`] — Kruskal over counting-sorted edges builds
//! the binary partition tree (a flat `parents` array, leaves `0..n`,
//! internal nodes created in altitude order); a leaves-to-root pass computes
//! each subtree's area and volume; each internal node's *persistence* (the
//! volume of the smaller of the two merged basins) becomes the saliency of
//! its MST edge. This depends only on the image — no tuning parameters.
//! * [`WatershedHierarchy::cut`] — cutting at level λ is single-linkage over
//! MST edges with persistence ≤ λ (every pixel gets a label, no
//! watershed-line pixels), small basins are absorbed, and the surviving
//! merge tree above λ becomes the output layer stack.
//!
//! The cut emits a **stacked hierarchy**, the same principle as the color
//! clustering frontend: the root (whole canvas, mean color) is painted first,
//! then progressively finer ancestor regions, then the final regions on top.
//! Sub-pixel gaps between abutting regions therefore show their common
//! ancestor's color instead of an unrelated backdrop, and stacked mode stays
//! seam-free by overdraw. Flattening top-down (what cutout does) recovers the
//! exact partition, because the final regions are painted last.
//!
//! Everything is integer and allocation-flat: counting sort over 256 weight
//! buckets, path-halving union-find, `u32` node ids. Deterministic across
@@ -29,12 +42,17 @@ use crate::ir::{Layer, Paint, RegionMask, Segmentation};
use super::Frontend;
/// Cap on the total painted area of ancestor layers, as a multiple of the
/// canvas: keeps a pathological hierarchy (long chains of near-equal
/// persistence) from ballooning the stacked output. The root and the final
/// regions are always emitted, so coverage never depends on this.
const ANCESTOR_AREA_BUDGET: usize = 3;
/// Watershed frontend: hierarchical watershed by volume, cut at `detail`.
#[derive(Debug, Clone)]
pub struct WatershedFrontend {
/// Detail level (0..=255): where to cut the hierarchy. 255 keeps every
/// basin that survives a zero-persistence merge (finest useful partition);
/// 0 merges everything into a single region.
/// Detail level (0..=255): where to cut the hierarchy. Each +25.5 roughly
/// doubles the region count; 0 collapses the image to a single region.
pub detail: u8,
/// Absorb regions smaller than this many pixels into their most
/// color-similar neighbour after the cut (0 = keep all).
@@ -82,20 +100,47 @@ fn edge_weight(a: Color, b: Color) -> u8 {
dr.max(dg).max(db)
}
impl WatershedFrontend {
fn label_map(&self, img: &ColorImage) -> Vec<u32> {
/// The image's watershed hierarchy: the minimum spanning tree of the pixel
/// graph with a persistence (volume extinction) per edge. Building it is the
/// expensive step and depends only on the image; [`cut`](Self::cut) derives a
/// [`Segmentation`] for any detail level in near-linear time, so interactive
/// re-tuning never repays the build (see [`crate::Session`]).
pub struct WatershedHierarchy {
width: usize,
height: usize,
/// MST edges as pixel pairs, in Kruskal creation order.
mst: Vec<(u32, u32)>,
/// Persistence (volume of the smaller merged basin) per MST edge.
pers: Vec<u64>,
/// MST edge indices by ascending (persistence, index) — the cut order.
order: Vec<u32>,
}
impl WatershedHierarchy {
/// Build the hierarchy: counting-sorted Kruskal → binary partition tree →
/// volume persistence per MST edge. O(n α(n)).
pub fn build(img: &ColorImage) -> Result<Self, Error> {
let w = img.width;
let h = img.height;
if w == 0 || h == 0 {
return Err(Error::EmptyImage);
}
let n = w * h;
if n == 1 {
return Ok(Self {
width: w,
height: h,
mst: Vec::new(),
pers: Vec::new(),
order: Vec::new(),
});
}
// --- 4-adjacency edges, counting-sorted by weight -------------------
// Edge id encodes (pixel, direction): 2*p = right, 2*p+1 = down.
// The per-bucket fill preserves edge-id order, so the sort is stable
// and the whole construction is deterministic.
let px = |i: usize| {
let c = img.get_pixel(i % w, i / w);
c
};
let px = |i: usize| img.get_pixel(i % w, i / w);
let mut counts = [0u32; 256];
let mut weight_of = vec![0u8; 2 * n];
for i in 0..n {
@@ -118,19 +163,19 @@ impl WatershedFrontend {
start[b] = acc;
acc += counts[b] as usize;
}
let mut order = vec![0u32; n_edges];
let mut sorted = vec![0u32; n_edges];
let mut fill = start;
for i in 0..n {
if i % w + 1 < w {
let e = 2 * i;
let b = weight_of[e] as usize;
order[fill[b]] = e as u32;
sorted[fill[b]] = e as u32;
fill[b] += 1;
}
if i / w + 1 < h {
let e = 2 * i + 1;
let b = weight_of[e] as usize;
order[fill[b]] = e as u32;
sorted[fill[b]] = e as u32;
fill[b] += 1;
}
}
@@ -144,12 +189,12 @@ impl WatershedFrontend {
let mut parent = vec![u32::MAX; n_nodes];
let mut alt = vec![0u8; n_nodes]; // altitude; leaves at 0
let mut child = vec![[0u32; 2]; n - 1]; // children of internal node k
let mut mst_edge = vec![(0u32, 0u32); n - 1]; // pixel pair of edge k
let mut mst = vec![(0u32, 0u32); n - 1]; // pixel pair of edge k
let mut uf = Uf::new(n);
// Current tree node representing each union-find root's component.
let mut comp_node: Vec<u32> = (0..n as u32).collect();
let mut next = n as u32;
for &e in &order {
for &e in &sorted {
let p = (e / 2) as usize;
let q = if e % 2 == 0 { p + 1 } else { p + w };
let (rp, rq) = (uf.find(p as u32), uf.find(q as u32));
@@ -159,7 +204,7 @@ impl WatershedFrontend {
let k = (next - n as u32) as usize;
alt[next as usize] = weight_of[e as usize];
child[k] = [comp_node[rp as usize], comp_node[rq as usize]];
mst_edge[k] = (p as u32, q as u32);
mst[k] = (p as u32, q as u32);
parent[comp_node[rp as usize] as usize] = next;
parent[comp_node[rq as usize] as usize] = next;
uf.link(rp, rq);
@@ -207,234 +252,328 @@ impl WatershedFrontend {
pers[k] = corrected[c0 as usize].min(corrected[c1 as usize]);
}
// --- Cut level from the detail slider --------------------------------
// Merging every MST edge with persistence ≤ λ leaves exactly
// 1 + #{edges above λ} regions, so choosing λ as the k-th largest
// persistence targets k regions directly (ties merge a little more).
// The persistence distribution is extremely skewed — most merges are
// trivia with persistence ≈ 0 — so the slider maps to a region
// *count*, exponentially: every +25.5 of detail doubles the target,
// from 1 region at 0 up to 1024 at 255.
let target = (2f64).powf(self.detail as f64 / 25.5).round() as usize;
let target = target.clamp(1, pers.len());
let lambda = {
let mut sorted = pers.clone();
sorted.sort_unstable_by(|a, b| b.cmp(a));
sorted[target - 1]
};
let mut order: Vec<u32> = (0..(n - 1) as u32).collect();
order.sort_by_key(|&k| (pers[k as usize], k));
// --- Single-linkage cut over MST edges -------------------------------
let mut cut = Uf::new(n);
for k in 0..n - 1 {
if pers[k] <= lambda {
let (p, q) = mst_edge[k];
let (rp, rq) = (cut.find(p), cut.find(q));
if rp != rq {
cut.link(rp, rq);
}
}
}
let mut labels = vec![0u32; n];
for (i, l) in labels.iter_mut().enumerate() {
*l = cut.find(i as u32);
}
labels
Ok(Self {
width: w,
height: h,
mst,
pers,
order,
})
}
/// Absorb regions smaller than `min_area` into their most color-similar
/// 4-neighbour. Works on root-labels in place; areas and color sums are
/// maintained through the merges so chains stay well-behaved.
fn absorb_small(&self, img: &ColorImage, labels: &mut [u32]) {
if self.min_area <= 1 {
return;
}
let w = img.width;
let n = labels.len();
/// Cut the hierarchy at `detail` and emit the stacked [`Segmentation`].
/// Near-linear; safe to call repeatedly with different parameters.
pub fn cut(&self, img: &ColorImage, detail: u8, min_area: usize) -> Segmentation {
let (w, h) = (self.width, self.height);
let n = w * h;
let m = self.mst.len();
// --- Region formation: merge every MST edge with persistence ≤ λ ----
// Merging leaves exactly 1 + #{edges above λ} regions, so choosing λ
// as the k-th largest persistence targets k regions directly (ties
// merge a little more). The persistence distribution is extremely
// skewed — most merges are trivia at ≈ 0 — so the dial maps to a
// region *count*, exponentially: every +25.5 of detail doubles the
// target, from 1 region at 0 up to 1024 at 255.
let mut uf = Uf::new(n);
// Rebuild region stats keyed by current label (a pixel index).
let mut area = vec![0u64; n];
let mut sum = vec![[0u64; 3]; n];
let mut split_from = 0usize;
if m > 0 {
let target = (2f64).powf(detail as f64 / 25.5).round() as usize;
let target = target.clamp(1, m);
let lambda = self.pers[self.order[m - target] as usize];
for (i, &k) in self.order.iter().enumerate() {
if self.pers[k as usize] > lambda {
break;
}
let (p, q) = self.mst[k as usize];
let (rp, rq) = (uf.find(p), uf.find(q));
if rp != rq {
uf.link(rp, rq);
}
split_from = i + 1;
}
}
// --- Compact to region ids, region stats, boundary adjacency --------
// One find per pixel; everything after this works on the (small)
// region graph so re-cuts stay cheap.
let mut pre_of_root = vec![u32::MAX; n];
let mut pre = vec![0u32; n];
let mut kp = 0usize;
for i in 0..n {
let l = labels[i] as usize;
let r = uf.find(i as u32) as usize;
if pre_of_root[r] == u32::MAX {
pre_of_root[r] = kp as u32;
kp += 1;
}
pre[i] = pre_of_root[r];
}
let mut area = vec![0u64; kp];
let mut sum = vec![[0u64; 3]; kp];
let mut pairs: Vec<(u32, u32)> = Vec::new();
for i in 0..n {
let a = pre[i];
let c = img.get_pixel(i % w, i / w);
area[l] += 1;
sum[l][0] += c.r as u64;
sum[l][1] += c.g as u64;
sum[l][2] += c.b as u64;
}
let mean_diff = |sa: &[u64; 3], aa: u64, sb: &[u64; 3], ab: u64| -> u64 {
let mut d = 0i64;
for ch in 0..3 {
d += ((sa[ch] / aa) as i64 - (sb[ch] / ab) as i64).abs();
area[a as usize] += 1;
sum[a as usize][0] += c.r as u64;
sum[a as usize][1] += c.g as u64;
sum[a as usize][2] += c.b as u64;
if i % w + 1 < w && pre[i + 1] != a {
pairs.push((a, pre[i + 1]));
}
d as u64
};
// Sweep until no undersized region can be absorbed. Each sweep scans
// the boundary edges once and merges each small region into its best
// neighbour seen so far; region count strictly decreases, so this
// terminates quickly in practice.
loop {
// best[l] = (diff, neighbour_root) for undersized root l
let mut best: Vec<(u64, u32)> = vec![(u64::MAX, u32::MAX); n];
let mut any_small = false;
for i in 0..n {
let a = uf.find(labels[i]);
for j in [
if i % w + 1 < w { i + 1 } else { i },
if i / w + 1 < labels.len() / w { i + w } else { i },
] {
if j == i {
continue;
}
let b = uf.find(labels[j]);
if a == b {
continue;
}
for (s, t) in [(a, b), (b, a)] {
let (su, tu) = (s as usize, t as usize);
if area[su] < self.min_area as u64 {
any_small = true;
let d = mean_diff(&sum[su], area[su], &sum[tu], area[tu]);
if d < best[su].0 || (d == best[su].0 && t < best[su].1) {
best[su] = (d, t);
}
}
}
}
}
if !any_small {
break;
}
let mut merged = false;
for l in 0..n {
let (_, tgt) = best[l];
if tgt == u32::MAX {
continue;
}
let rl = uf.find(l as u32);
if rl as usize != l {
continue; // already absorbed this sweep
}
let rt = uf.find(tgt);
if rt == rl {
continue;
}
uf.link(rt, rl);
area[rt as usize] += area[l];
for ch in 0..3 {
sum[rt as usize][ch] += sum[l][ch];
}
merged = true;
}
if !merged {
break; // isolated undersized region (e.g. whole-canvas)
if i / w + 1 < h && pre[i + w] != a {
pairs.push((a, pre[i + w]));
}
}
for l in labels.iter_mut() {
*l = uf.find(*l);
}
}
/// Turn a root-label map into the layered [`Segmentation`]: one layer per
/// region with its mean color, the largest region first as a solid
/// full-canvas background so stacked mode stays seam-free by overdraw.
fn segmentation(img: &ColorImage, labels: &[u32]) -> Segmentation {
let w = img.width;
let h = img.height;
let n = labels.len();
// --- Small-basin absorption on the region graph ----------------------
let mut uf_r = Uf::new(kp);
absorb_small(min_area, &pairs, &mut uf_r, &mut area, &mut sum);
// Compact labels in raster order of first appearance (deterministic).
let mut compact = vec![u32::MAX; n];
let mut regions: Vec<u32> = Vec::new(); // compact id -> root label
// --- Final leaf ids in raster order of first appearance --------------
let mut leaf_of = vec![u32::MAX; kp];
let mut leaf_root: Vec<u32> = Vec::new(); // leaf id -> absorb root
let mut ids = vec![0u32; n];
for i in 0..n {
let l = labels[i] as usize;
if compact[l] == u32::MAX {
compact[l] = regions.len() as u32;
regions.push(labels[i]);
let r = uf_r.find(pre[i]) as usize;
if leaf_of[r] == u32::MAX {
leaf_of[r] = leaf_root.len() as u32;
leaf_root.push(r as u32);
}
ids[i] = compact[l];
ids[i] = leaf_of[r];
}
let m = regions.len();
let k = leaf_root.len();
let mut area = vec![0u64; m];
let mut sum = vec![[0u64; 3]; m];
let mut bbox = vec![(i32::MAX, i32::MAX, i32::MIN, i32::MIN); m];
for i in 0..n {
let id = ids[i] as usize;
let (x, y) = ((i % w) as i32, (i / w) as i32);
let c = img.get_pixel(i % w, i / w);
area[id] += 1;
sum[id][0] += c.r as u64;
sum[id][1] += c.g as u64;
sum[id][2] += c.b as u64;
let b = &mut bbox[id];
b.0 = b.0.min(x);
b.1 = b.1.min(y);
b.2 = b.2.max(x);
b.3 = b.3.max(y);
}
let mean = |id: usize| {
Color::new(
(sum[id][0] / area[id]) as u8,
(sum[id][1] / area[id]) as u8,
(sum[id][2] / area[id]) as u8,
)
let mean = |s: &[u64; 3], a: u64| {
Color::new((s[0] / a) as u8, (s[1] / a) as u8, (s[2] / a) as u8)
};
let background = (0..m).max_by_key(|&id| area[id]).unwrap_or(0);
let mut seg = Segmentation::new(w as u32, h as u32);
// Background: solid full canvas, painted first; the regions stacked on
// top stamp out everything that isn't actually background, so the
// flattened partition is exact while stacked mode keeps overdraw.
let mut bg = BinaryImage::new_w_h(w, h);
for y in 0..h {
for x in 0..w {
bg.set_pixel(x, y, true);
}
if k == 1 {
// Single region: one solid full-canvas layer.
let r = leaf_root[0] as usize;
seg.layers.push(Layer {
paint: Paint::Solid(mean(&sum[r], area[r])),
mask: full_canvas(w, h),
});
return seg;
}
seg.layers.push(Layer {
paint: Paint::Solid(mean(background)),
mask: RegionMask::new(bg, PointI32 { x: 0, y: 0 }),
});
for id in 0..m {
if id == background {
continue;
// --- Merge tree above the cut ----------------------------------------
// Re-run the remaining merges (ascending persistence) over the final
// regions: each one that still joins two components is a kept split.
// Nodes 0..k are the final regions; internal nodes are created in
// ascending persistence order, so the reverse is a root-first order in
// which every ancestor precedes its descendants.
let n_tree = 2 * k - 1;
let mut tree_child: Vec<[u32; 2]> = Vec::with_capacity(k - 1);
let mut tree_area = vec![0u64; n_tree];
let mut tree_sum = vec![[0u64; 3]; n_tree];
for (t, &r) in leaf_root.iter().enumerate() {
tree_area[t] = area[r as usize];
tree_sum[t] = sum[r as usize];
}
let mut uf2 = Uf::new(k);
let mut node_rep: Vec<u32> = (0..k as u32).collect();
let mut next = k as u32;
for &e in &self.order[split_from..] {
let (p, q) = self.mst[e as usize];
let (lp, lq) = (ids[p as usize], ids[q as usize]);
let (a, b) = (uf2.find(lp), uf2.find(lq));
if a == b {
continue; // rejoined by absorption; not a split anymore
}
let (x0, y0, x1, y1) = bbox[id];
let (bw, bh) = ((x1 - x0 + 1) as usize, (y1 - y0 + 1) as usize);
let mut image = BinaryImage::new_w_h(bw, bh);
for y in 0..bh {
for x in 0..bw {
let i = (y0 as usize + y) * w + (x0 as usize + x);
if ids[i] as usize == id {
image.set_pixel(x, y, true);
}
let node = next as usize;
tree_child.push([node_rep[a as usize], node_rep[b as usize]]);
for ch in [node_rep[a as usize], node_rep[b as usize]] {
tree_area[node] += tree_area[ch as usize];
for c in 0..3 {
tree_sum[node][c] += tree_sum[ch as usize][c];
}
}
uf2.link(a, b);
node_rep[a as usize] = next;
next += 1;
}
debug_assert_eq!(next as usize, n_tree);
// Per-leaf pixel lists, for painting ancestor masks.
let mut leaf_len = vec![0u32; k];
for &id in &ids {
leaf_len[id as usize] += 1;
}
let mut leaf_start = vec![0usize; k + 1];
for t in 0..k {
leaf_start[t + 1] = leaf_start[t] + leaf_len[t] as usize;
}
let mut leaf_px = vec![0u32; n];
let mut fill = leaf_start.clone();
for (i, &id) in ids.iter().enumerate() {
leaf_px[fill[id as usize]] = i as u32;
fill[id as usize] += 1;
}
// --- Emit: root, ancestors (budgeted), then the final regions --------
let root = n_tree - 1;
seg.layers.push(Layer {
paint: Paint::Solid(mean(&tree_sum[root], tree_area[root])),
mask: full_canvas(w, h),
});
let mut budget = ANCESTOR_AREA_BUDGET * n;
for node in (k..root).rev() {
let node_area = tree_area[node] as usize;
if node_area > budget {
continue;
}
budget -= node_area;
seg.layers.push(Layer {
paint: Paint::Solid(mean(id)),
mask: RegionMask::new(image, PointI32 { x: x0, y: y0 }),
paint: Paint::Solid(mean(&tree_sum[node], tree_area[node])),
mask: node_mask(node, k, &tree_child, &leaf_start, &leaf_px, w),
});
}
for t in 0..k {
seg.layers.push(Layer {
paint: Paint::Solid(mean(&tree_sum[t], tree_area[t])),
mask: node_mask(t, k, &tree_child, &leaf_start, &leaf_px, w),
});
}
seg
}
}
impl Frontend for WatershedFrontend {
fn segment(&self, img: &ColorImage) -> Result<Segmentation, Error> {
if img.width == 0 || img.height == 0 {
return Err(Error::EmptyImage);
}
if img.width * img.height == 1 {
// Degenerate single pixel: no edges, one region.
let labels = [0u32];
return Ok(Self::segmentation(img, &labels));
fn full_canvas(w: usize, h: usize) -> RegionMask {
let mut image = BinaryImage::new_w_h(w, h);
for y in 0..h {
for x in 0..w {
image.set_pixel(x, y, true);
}
}
RegionMask::new(image, PointI32 { x: 0, y: 0 })
}
let mut labels = self.label_map(img);
self.absorb_small(img, &mut labels);
Ok(Self::segmentation(img, &labels))
/// Paint a tree node's region (the union of the final regions beneath it)
/// into a bbox-cropped mask.
fn node_mask(
node: usize,
k: usize,
tree_child: &[[u32; 2]],
leaf_start: &[usize],
leaf_px: &[u32],
w: usize,
) -> RegionMask {
// Collect the node's leaves.
let mut leaves: Vec<usize> = Vec::new();
let mut stack = vec![node];
while let Some(t) = stack.pop() {
if t < k {
leaves.push(t);
} else {
let [a, b] = tree_child[t - k];
stack.push(a as usize);
stack.push(b as usize);
}
}
// Bounding box over all member pixels.
let (mut x0, mut y0, mut x1, mut y1) = (i32::MAX, i32::MAX, i32::MIN, i32::MIN);
for &t in &leaves {
for &p in &leaf_px[leaf_start[t]..leaf_start[t + 1]] {
let (x, y) = ((p as usize % w) as i32, (p as usize / w) as i32);
x0 = x0.min(x);
y0 = y0.min(y);
x1 = x1.max(x);
y1 = y1.max(y);
}
}
let (bw, bh) = ((x1 - x0 + 1) as usize, (y1 - y0 + 1) as usize);
let mut image = BinaryImage::new_w_h(bw, bh);
for &t in &leaves {
for &p in &leaf_px[leaf_start[t]..leaf_start[t + 1]] {
let (x, y) = (p as usize % w, p as usize / w);
image.set_pixel(x - x0 as usize, y - y0 as usize, true);
}
}
RegionMask::new(image, PointI32 { x: x0, y: y0 })
}
/// Absorb regions smaller than `min_area` into their most color-similar
/// neighbour, working entirely on the region graph: `pairs` are the boundary
/// adjacencies (duplicates fine), `uf` is a region-level union-find, and the
/// stats are merged along so downstream consumers see the final regions.
/// Sweeps until nothing undersized remains (or an undersized region has no
/// neighbour at all).
fn absorb_small(
min_area: usize,
pairs: &[(u32, u32)],
uf: &mut Uf,
area: &mut [u64],
sum: &mut [[u64; 3]],
) {
if min_area <= 1 {
return;
}
let k = area.len();
let mean_diff = |sa: &[u64; 3], aa: u64, sb: &[u64; 3], ab: u64| -> u64 {
let mut d = 0i64;
for ch in 0..3 {
d += ((sa[ch] / aa) as i64 - (sb[ch] / ab) as i64).abs();
}
d as u64
};
loop {
// best[r] = (diff, neighbour_root) for undersized root r
let mut best: Vec<(u64, u32)> = vec![(u64::MAX, u32::MAX); k];
let mut any_small = false;
for &(p, q) in pairs {
let (a, b) = (uf.find(p), uf.find(q));
if a == b {
continue;
}
for (s, t) in [(a, b), (b, a)] {
let (su, tu) = (s as usize, t as usize);
if area[su] < min_area as u64 {
any_small = true;
let d = mean_diff(&sum[su], area[su], &sum[tu], area[tu]);
if d < best[su].0 || (d == best[su].0 && t < best[su].1) {
best[su] = (d, t);
}
}
}
}
if !any_small {
break;
}
let mut merged = false;
for r in 0..k {
let (_, tgt) = best[r];
if tgt == u32::MAX {
continue;
}
let rr = uf.find(r as u32);
if rr as usize != r {
continue; // already absorbed this sweep
}
let rt = uf.find(tgt);
if rt == rr {
continue;
}
uf.link(rt, rr);
area[rt as usize] += area[r];
for ch in 0..3 {
sum[rt as usize][ch] += sum[r][ch];
}
merged = true;
}
if !merged {
break; // isolated undersized region (e.g. whole-canvas)
}
}
}
impl Frontend for WatershedFrontend {
fn segment(&self, img: &ColorImage) -> Result<Segmentation, Error> {
Ok(WatershedHierarchy::build(img)?.cut(img, self.detail, self.min_area))
}
}
crates/vtracer/src/session.rs
+52 -11
View File
@@ -7,6 +7,11 @@
//! re-segments only if a clustering parameter changed — the caller never has to
//! know which parameters those are.
//!
//! For watershed clustering there is a second cache level: the
//! [`WatershedHierarchy`] depends only on the image, so it is built once and
//! every re-segmentation (a detail or speckle change) is a near-instant re-cut
//! of the cached hierarchy rather than a rebuild.
//!
//! ```no_run
//! use vtracer::{Config, Session, ColorImage};
//! # fn load() -> ColorImage { todo!() }
@@ -27,10 +32,12 @@
use visioncortex::ColorImage;
use crate::config::{Config, SegmentKey};
use crate::config::{Clustering, Config, SegmentKey};
use crate::error::Error;
use crate::frontend::WatershedHierarchy;
use crate::ir::{Segmentation, VectorDoc};
use crate::progress::{CancelToken, Progress};
use crate::pipeline::Pipeline;
use crate::progress::{CancelToken, Ctx, Phase, Progress};
/// A reusable converter for one image: clusters once, re-renders many times.
///
@@ -42,12 +49,19 @@ pub struct Session {
/// The segmentation and the key it was produced with (`None` until the
/// first render).
cache: Option<(SegmentKey, Segmentation)>,
/// The image's watershed hierarchy, built lazily on the first watershed
/// render. Parameter-free, so it never goes stale while the image lives.
hierarchy: Option<WatershedHierarchy>,
}
impl Session {
/// Start a session over `img`. Nothing is clustered until the first render.
pub fn new(img: ColorImage) -> Self {
Self { img, cache: None }
Self {
img,
cache: None,
hierarchy: None,
}
}
/// Whether the cached segmentation is missing or was clustered with
@@ -61,13 +75,29 @@ impl Session {
&self.cache.as_ref().expect("cache populated by caller").1
}
/// Produce a fresh segmentation for `cfg`. Watershed goes through the
/// hierarchy cache (build once, cut cheaply); everything else runs the
/// pipeline's frontend.
fn segment(&mut self, cfg: &Config, pipeline: &Pipeline) -> Result<Segmentation, Error> {
if cfg.clustering == Clustering::Watershed {
if self.hierarchy.is_none() {
self.hierarchy = Some(WatershedHierarchy::build(&self.img)?);
}
let hierarchy = self.hierarchy.as_ref().expect("just built");
Ok(hierarchy.cut(&self.img, cfg.watershed_detail, cfg.speckle_area()))
} else {
pipeline.segment(&self.img)
}
}
/// Render to the document IR, re-segmenting only if `cfg`'s clustering
/// parameters differ from the cached segmentation's.
pub fn render(&mut self, cfg: &Config) -> Result<VectorDoc, Error> {
let pipeline = cfg.build()?;
let key = cfg.segment_key();
if self.stale(&key) {
self.cache = Some((key, pipeline.segment(&self.img)?));
let seg = self.segment(cfg, &pipeline)?;
self.cache = Some((key, seg));
}
pipeline.finish(self.segmentation())
}
@@ -77,7 +107,8 @@ impl Session {
let pipeline = cfg.build()?;
let key = cfg.segment_key();
if self.stale(&key) {
self.cache = Some((key, pipeline.segment(&self.img)?));
let seg = self.segment(cfg, &pipeline)?;
self.cache = Some((key, seg));
}
Ok(pipeline.writer.write(&pipeline.finish(self.segmentation())?))
}
@@ -87,9 +118,8 @@ impl Session {
/// When a re-segmentation is needed, progress covers the [`Phase::Segment`]
/// stage first, then the finish stages; on a cache hit only the finish
/// stages report. Hand a clone of `cancel` to the UI to abort a long
/// clustering pass.
///
/// [`Phase::Segment`]: crate::Phase::Segment
/// clustering pass. (A watershed re-cut over a cached hierarchy is fast
/// enough that it reports coarsely.)
pub fn render_with_progress(
&mut self,
cfg: &Config,
@@ -99,16 +129,27 @@ impl Session {
let pipeline = cfg.build()?;
let key = cfg.segment_key();
if self.stale(&key) {
let seg = pipeline.segment_with_progress(&self.img, cancel, on_progress)?;
let seg = if cfg.clustering == Clustering::Watershed {
let mut ctx = Ctx::new(cancel, on_progress);
ctx.check()?;
let seg = self.segment(cfg, &pipeline)?;
ctx.check()?;
ctx.report(Phase::Segment, 1.0);
seg
} else {
pipeline.segment_with_progress(&self.img, cancel, on_progress)?
};
self.cache = Some((key, seg));
}
pipeline.finish_with_progress(self.segmentation(), cancel, on_progress)
}
/// Drop the cached segmentation, forcing the next render to re-cluster.
/// Use after replacing the source image out of band; normally unnecessary.
/// Drop the cached segmentation and hierarchy, forcing the next render to
/// re-cluster. Use after replacing the source image out of band; normally
/// unnecessary.
pub fn invalidate(&mut self) {
self.cache = None;
self.hierarchy = None;
}
/// The source image this session renders.
crates/vtracer/tests/goldens/disc_watershed_spline.svg
+2 -1
View File
@@ -1,6 +1,7 @@
<?xml version="1.0" encoding="UTF-8"?>
<!-- Generator: visioncortex VTracer 1.0.0-alpha.1 -->
<svg version="1.1" xmlns="http://www.w3.org/2000/svg" width="48" height="48">
<path d="M0,0C15.84,0,31.68,0,48,0c0,15.84,0,31.68,0,48c-15.84,0-31.68,0-48,0C0,32.16,0,16.32,0,0Z" fill="#F0F0F0"/>
<path d="M0,0C15.84,0,31.68,0,48,0c0,15.84,0,31.68,0,48c-15.84,0-31.68,0-48,0C0,32.16,0,16.32,0,0Z" fill="#E2B1B1"/>
<path d="M0,0C15.84,0,31.68,0,48,0c0,15.84,0,31.68,0,48c-15.84,0-31.68,0-48,0C0,32.16,0,16.32,0,0ZM12.06,13.69c-2.83,4.54-4.51,8.24-3.72,13.65C9.88,32.82,12.06,35.92,17,39c4.43,2.03,7.98,2.31,12.69,1c5.03-1.99,7.88-4.88,10.16-9.85c1.58-4.77,1.09-8.47-.74-13.1c-2.46-4.55-6.13-6.96-10.92-8.55c-6.36-1.27-11.46,.92-16.13,5.19Z" fill="#F0F0F0"/>
<path d="M35.31,12.06c3.79,4.14,5.85,8.38,5.6,14.09c-1,4.97-3.14,8.4-6.91,11.85c-4.31,2.33-8.33,3.47-13.19,2.5c-4.79-1.59-8.46-4-10.92-8.55c-1.83-4.63-2.32-8.33-.74-13.1c2.28-4.97,5.13-7.86,10.16-9.85c6.33-1.77,10.55-.34,16,3.06Z" fill="#C83C3C"/>
</svg>
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 544 B

After

Width:  |  Height:  |  Size: 887 B

crates/vtracer/tests/watershed.rs
+141 -41
View File
@@ -1,8 +1,8 @@
//! Watershed frontend: partition invariants, the detail dial, and small-basin
//! absorption.
//! Watershed frontend: partition invariants, the detail dial, small-basin
//! absorption, and the hierarchy stack / cached re-cut behavior.
use vtracer::frontend::{Frontend, WatershedFrontend};
use vtracer::ColorImage;
use vtracer::frontend::{Frontend, WatershedFrontend, WatershedHierarchy};
use vtracer::{ColorImage, Clustering, Config, Hierarchical, Segmentation, Session};
fn image(w: usize, h: usize, f: impl Fn(usize, usize) -> (u8, u8, u8)) -> ColorImage {
let mut pixels = Vec::with_capacity(w * h * 4);
@@ -19,34 +19,57 @@ fn image(w: usize, h: usize, f: impl Fn(usize, usize) -> (u8, u8, u8)) -> ColorI
}
}
/// The core partition invariant behind the seam-free mosaic: painting the
/// layers bottom-to-top covers every canvas pixel exactly once per region —
/// i.e. the non-background masks are pairwise disjoint, and together with the
/// full-canvas background they tile the image.
fn assert_partition(seg: &vtracer::Segmentation) {
/// Flatten the stacked layers top-down (later layers win), returning one layer
/// index per pixel — the partition both compositors ultimately consume.
fn flatten(seg: &Segmentation) -> Vec<usize> {
let (w, h) = (seg.width as usize, seg.height as usize);
// Background layer must be first and cover the full canvas.
let bg = &seg.layers[0].mask;
assert_eq!((bg.width(), bg.height()), (w, h), "background is full-canvas");
assert_eq!(bg.area(), w * h, "background mask is solid");
// Later layers are pairwise disjoint.
let mut covered = vec![false; w * h];
for layer in &seg.layers[1..] {
let mut labels = vec![usize::MAX; w * h];
for (li, layer) in seg.layers.iter().enumerate() {
let m = &layer.mask;
for y in 0..m.image.height {
for x in 0..m.image.width {
if !m.image.get_pixel(x, y) {
continue;
if m.image.get_pixel(x, y) {
let gx = (m.offset.x + x as i32) as usize;
let gy = (m.offset.y + y as i32) as usize;
labels[gy * w + gx] = li;
}
let gx = (m.offset.x + x as i32) as usize;
let gy = (m.offset.y + y as i32) as usize;
assert!(gx < w && gy < h, "mask pixel out of canvas");
assert!(!covered[gy * w + gx], "overlapping region masks");
covered[gy * w + gx] = true;
}
}
}
labels
}
/// The stacked-hierarchy invariants: the bottom layer is a solid full canvas
/// (so overdraw is seam-free), every pixel is covered, and the flattened
/// partition has exactly `regions` distinct labels.
fn assert_stack(seg: &Segmentation, regions: usize) {
let (w, h) = (seg.width as usize, seg.height as usize);
let bottom = &seg.layers[0].mask;
assert_eq!((bottom.width(), bottom.height()), (w, h), "bottom layer is full-canvas");
assert_eq!(bottom.area(), w * h, "bottom layer is solid");
let labels = flatten(seg);
assert!(labels.iter().all(|&l| l != usize::MAX), "every pixel covered");
let mut distinct: Vec<usize> = labels.clone();
distinct.sort_unstable();
distinct.dedup();
assert_eq!(distinct.len(), regions, "flattened region count");
// The final regions must be the topmost layers (painted after every
// ancestor), or the flatten would not recover the partition.
let first_final = seg.layers.len() - regions;
assert!(
distinct.iter().all(|&l| l >= first_final),
"final regions are the topmost layers"
);
}
/// Region count of a segmentation's flattened partition.
fn regions(seg: &Segmentation) -> usize {
let mut labels = flatten(seg);
labels.sort_unstable();
labels.dedup();
labels.len()
}
/// A flat single-color image is one region no matter the detail level.
@@ -61,12 +84,12 @@ fn flat_image_is_one_region() {
.segment(&img)
.unwrap();
assert_eq!(seg.layers.len(), 1, "detail={detail}");
assert_partition(&seg);
assert_stack(&seg, 1);
}
}
/// Two clearly separated halves form two regions, with the boundary exactly on
/// the color edge (no watershed-line pixels — the partition is gapless).
/// Two clearly separated halves form two regions plus their common ancestor:
/// the stack is [root, half, half] and the flatten recovers the exact split.
#[test]
fn two_tone_image_is_two_regions() {
let img = image(32, 20, |x, _| {
@@ -82,10 +105,11 @@ fn two_tone_image_is_two_regions() {
}
.segment(&img)
.unwrap();
assert_eq!(seg.layers.len(), 2);
assert_partition(&seg);
// The non-background region is exactly one half of the canvas.
assert_eq!(seg.layers.len(), 3, "root + two final regions");
assert_stack(&seg, 2);
// Each final region is exactly one half of the canvas.
assert_eq!(seg.layers[1].mask.area(), 16 * 20);
assert_eq!(seg.layers[2].mask.area(), 16 * 20);
}
/// Raising detail never decreases the region count (the hierarchy cut is
@@ -108,19 +132,16 @@ fn detail_is_monotone() {
}
.segment(&img)
.unwrap();
assert!(
seg.layers.len() >= prev,
"detail={detail}: {} < {prev}",
seg.layers.len()
);
assert_partition(&seg);
prev = seg.layers.len();
let k = regions(&seg);
assert!(k >= prev, "detail={detail}: {k} < {prev}");
assert_stack(&seg, k);
prev = k;
}
assert!(prev > 1, "highest detail should find several regions");
}
/// Small basins are absorbed into a neighbour rather than dropped: the region
/// disappears but its pixels stay covered (the partition invariant holds).
/// disappears but its pixels stay covered.
#[test]
fn min_area_absorbs_small_basins() {
// Background plus a 3x3 fleck and a 12x12 block, all far apart in color.
@@ -146,9 +167,9 @@ fn min_area_absorbs_small_basins() {
.segment(&img)
.unwrap();
assert!(keep.layers.len() > absorb.layers.len(), "fleck absorbed");
assert_eq!(absorb.layers.len(), 2, "background + block survive");
assert_partition(&absorb);
assert!(regions(&keep) > regions(&absorb), "fleck absorbed");
assert_eq!(regions(&absorb), 2, "background + block survive");
assert_stack(&absorb, 2);
}
/// Output is deterministic: two runs produce identical layer geometry.
@@ -170,3 +191,82 @@ fn deterministic() {
assert_eq!(la.mask.area(), lb.mask.area());
}
}
/// A cut of a prebuilt hierarchy equals the one-shot frontend — the contract
/// behind `Session`'s cached re-cut.
#[test]
fn hierarchy_recut_matches_one_shot() {
let img = image(48, 32, |x, y| {
(((x * 5 + y * 3) % 200) as u8, ((x / 8) * 30) as u8, ((y / 8) * 40) as u8)
});
let hierarchy = WatershedHierarchy::build(&img).unwrap();
for detail in [64u8, 128, 200] {
let recut = hierarchy.cut(&img, detail, 16);
let one_shot = WatershedFrontend {
detail,
min_area: 16,
}
.segment(&img)
.unwrap();
assert_eq!(recut.layers.len(), one_shot.layers.len(), "detail={detail}");
for (a, b) in recut.layers.iter().zip(&one_shot.layers) {
assert_eq!(a.paint, b.paint);
assert_eq!(a.mask.offset, b.mask.offset);
assert_eq!(a.mask.area(), b.mask.area());
}
}
}
/// End-to-end through `Session`: retuning watershed detail re-cuts the cached
/// hierarchy, and the output still equals the one-shot pipeline.
#[test]
fn session_recut_matches_one_shot() {
let img = image(48, 32, |x, y| {
(((x * 5 + y * 3) % 200) as u8, ((x / 8) * 30) as u8, ((y / 8) * 40) as u8)
});
let mut session = Session::new(img.clone());
let base = Config {
clustering: Clustering::Watershed,
..Config::default()
};
for detail in [128u8, 200, 64] {
let cfg = Config {
watershed_detail: detail,
..base.clone()
};
assert_eq!(
session.render_svg(&cfg).unwrap(),
cfg.build().unwrap().to_svg(&img).unwrap(),
"detail={detail}: session re-cut must match the one-shot pipeline"
);
}
}
/// Watershed + cutout is native: the partition reaches the mosaic untouched,
/// so two regions within one gradient step stay separate faces (the color
/// path's `merge_similar` would have rejoined them).
#[test]
fn cutout_keeps_watershed_partition() {
// Two halves 4 gray-levels apart: close enough that the flatten merge
// (threshold = layer_difference = 16 >= 3*4) would union them.
let img = image(32, 20, |x, _| {
if x < 16 {
(100, 100, 100)
} else {
(104, 104, 104)
}
});
let cfg = Config {
clustering: Clustering::Watershed,
hierarchical: Hierarchical::Cutout,
watershed_detail: 255,
filter_speckle: 0,
..Config::default()
};
let doc = cfg.build().unwrap().run(&img).unwrap();
assert_eq!(
doc.shapes.len(),
2,
"watershed partition must pass to the mosaic unmerged"
);
}
docs/design/architecture.md
Generated
+1 -1
View File
@@ -114,7 +114,7 @@ Driver flow:
- **Frontends** (selected by `Config::clustering`)
- `ColorClusterFrontend` — wraps `visioncortex::color_clusters::Runner`, including the transparency-keying logic that currently lives in `converter.rs` (find unused key color, key fully-transparent pixels, `KeyingAction`).
- `BinaryFrontend` — threshold → `BinaryImage::to_clusters`.
- `WatershedFrontend` — hierarchical watershed by volume on the 4-adjacency pixel graph (Cousty et al. TPAMI 2009; Najman, Cousty & Perret ISMM 2013), cut at `watershed_detail`. Emits a flat partition with a solid full-canvas background layer so stacked overdraw stays seam-free.
- `WatershedFrontend` — hierarchical watershed by volume on the 4-adjacency pixel graph (Cousty et al. TPAMI 2009; Najman, Cousty & Perret ISMM 2013), cut at `watershed_detail`. Split into `WatershedHierarchy::build` (expensive, image-only) and `cut` (near-instant), so `Session` re-cuts a cached hierarchy when the detail changes. Emits the merge tree as a stacked hierarchy (root first, refined regions on top — the color-cluster principle), so stacked mode stays seam-free and sub-pixel gaps show ancestor colors; in cutout the partition reaches the mosaic untouched (`merge_diff = 0`).
- Third parties implement `Frontend` to feed external label maps or ML segmentation.
- **ColorFitters**
- `Identity` (today's behavior: mean cluster color)