dtourolle
0da8271836
Local masking needs to know where an image's regions are. The watershed spike (S15 arm A) found the boundaries but had no idea what any of them enclosed; its coarse levels were geometric accidents. This adds the other half and the thing that joins them. `core/dr-segment` is where region reasoning now lives — the hierarchy moves out of `dr-gpu`, which keeps only the pixel passes that are genuinely shaders. The new crate is device-free and, without its default features, model-free too: 20 of its tests need neither an adapter nor 11 MB of weights. Arm B runs YOLO26n-seg through `ort`. D13 framed inference as a choice between `ort`'s C++ runtime and the pure-Rust dependency policy; that was a false choice. `ort`'s `alternative-backend` feature unlinks the C entirely and `ort-tract` supplies the API from tract, which is pure Rust. Measured before committing to it: zero unsupported operators, 420 ms for 640x640, and correct masks on bus.jpg. No NDK problem to solve, so D13's largest tolerated exception is not needed. Arm C is `prior.rs`, and it ships because the two arms fail in opposite directions. Instance membership re-weights the merge saddles, so region pairs the model believes share an object merge early and pairs straddling its edge merge late. No boundary moves — only the order in which they dissolve — which is how the result stays pixel-accurate at every level while its coarse levels become named things. Two things the spec assumed that turned out to be false, both recorded in models/LICENCE.md: there is no usable ADE20K-trained YOLO, so the shipped vocabulary is COCO's 80 subjects and *stuff* like sky and foliage must come from arm A; and tract cannot parse a dynamic-shape export, so the graph's input is fixed and tiling is the only route to more semantic resolution. Weights are AGPL-3.0, which GPLv3 §13 permits and which makes the combined work effectively AGPL. Deliberate, not accidental. They live in Git LFS, and a build script fails with an instruction rather than embedding a pointer file when the clone lacks them.
430 lines
15 KiB
Rust
430 lines
15 KiB
Rust
//! Region adjacency and the merge hierarchy — arm A's CPU half (S15).
|
||
//!
|
||
//! **Deliberately device-free.** Nothing here touches wgpu, and that is the
|
||
//! point: the hierarchy is where arm A's behaviour actually lives, so it has
|
||
//! to be assertable on hand-built inputs rather than only on whatever a GPU
|
||
//! happened to produce (ARCH §6.5a). Every test below runs on CI machines
|
||
//! with no adapter.
|
||
//!
|
||
//! It is also why this is a legitimate CPU stage rather than a violation of
|
||
//! ARCH §6.1. The watershed's basins are found on the GPU, over pixels; what
|
||
//! follows runs on the **region adjacency graph**, which for a 2 MP proxy is
|
||
//! a few thousand nodes. Pixels stay on the GPU, the graph is CPU-side —
|
||
//! the same split ARCH §3.4 already draws for the edit graph.
|
||
//!
|
||
//! # Why a hierarchy rather than one partition
|
||
//!
|
||
//! A single segmentation has exactly one granularity, and no threshold is
|
||
//! right for both "her eye" and "her face". So the watershed deliberately
|
||
//! *over*-segments, and the merge order recorded here is what lets one
|
||
//! interaction — a scroll, a drag — walk from a fragment up to the object
|
||
//! that contains it.
|
||
|
||
use std::collections::HashMap;
|
||
|
||
/// One boundary between two adjacent regions.
|
||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||
pub struct Edge {
|
||
pub a: u32,
|
||
pub b: u32,
|
||
/// The **saddle**: the lowest point on the ridge separating the two
|
||
/// basins, which is the height at which water joining them would first
|
||
/// spill over.
|
||
///
|
||
/// The minimum over the shared boundary, not the mean or the maximum. A
|
||
/// long boundary that is strong everywhere except one weak gap describes
|
||
/// two regions that should merge early — the gap is exactly where a
|
||
/// person would say the edge fails.
|
||
pub saddle: f32,
|
||
}
|
||
|
||
/// A partition of the image into labelled regions, plus how they adjoin.
|
||
///
|
||
/// The shared interface from docs/segmentation.md §2: arm A produces this
|
||
/// from a watershed, arm B would produce it from a class map, and the
|
||
/// consumers above cannot tell which.
|
||
#[derive(Debug, Clone, PartialEq)]
|
||
pub struct RegionField {
|
||
pub width: usize,
|
||
pub height: usize,
|
||
/// Region id per pixel, compacted to `0..region_count`.
|
||
pub labels: Vec<u32>,
|
||
pub region_count: usize,
|
||
/// Deduplicated, `a < b`, sorted for reproducibility.
|
||
pub adjacency: Vec<Edge>,
|
||
}
|
||
|
||
impl RegionField {
|
||
/// Build from the watershed's raw output.
|
||
///
|
||
/// `roots` holds, per pixel, the linear index of its basin root — what
|
||
/// the pointer-jumping pass converges to. Those indices are sparse and
|
||
/// arbitrary, so they are compacted here to `0..region_count` in order of
|
||
/// first appearance, which makes them stable to serialise and cheap to
|
||
/// index.
|
||
pub fn from_roots(roots: &[u32], gradient: &[f32], width: usize, height: usize) -> Self {
|
||
assert_eq!(roots.len(), width * height, "roots must cover every pixel");
|
||
assert_eq!(
|
||
gradient.len(),
|
||
width * height,
|
||
"gradient must cover every pixel"
|
||
);
|
||
|
||
// Compact in raster order rather than by hashing the root value, so
|
||
// the same image always yields the same numbering (M5).
|
||
let mut compact: HashMap<u32, u32> = HashMap::new();
|
||
let mut labels = vec![0u32; roots.len()];
|
||
for (i, &root) in roots.iter().enumerate() {
|
||
let next = compact.len() as u32;
|
||
labels[i] = *compact.entry(root).or_insert(next);
|
||
}
|
||
let region_count = compact.len();
|
||
|
||
// Saddles, over 4-neighbour adjacency. Crossing a boundary means
|
||
// climbing to the higher of the two pixels, so a pair's pass height
|
||
// is the max; the region pair's saddle is the min over all its pairs.
|
||
let mut saddles: HashMap<(u32, u32), f32> = HashMap::new();
|
||
let note = |p: usize, q: usize, saddles: &mut HashMap<(u32, u32), f32>| {
|
||
let (la, lb) = (labels[p], labels[q]);
|
||
if la == lb {
|
||
return;
|
||
}
|
||
let key = (la.min(lb), la.max(lb));
|
||
let pass = gradient[p].max(gradient[q]);
|
||
saddles
|
||
.entry(key)
|
||
.and_modify(|s| {
|
||
if pass < *s {
|
||
*s = pass;
|
||
}
|
||
})
|
||
.or_insert(pass);
|
||
};
|
||
|
||
for y in 0..height {
|
||
for x in 0..width {
|
||
let i = y * width + x;
|
||
if x + 1 < width {
|
||
note(i, i + 1, &mut saddles);
|
||
}
|
||
if y + 1 < height {
|
||
note(i, i + width, &mut saddles);
|
||
}
|
||
}
|
||
}
|
||
|
||
let mut adjacency: Vec<Edge> = saddles
|
||
.into_iter()
|
||
.map(|((a, b), saddle)| Edge { a, b, saddle })
|
||
.collect();
|
||
// Sorted by (saddle, a, b) with `total_cmp` rather than `partial_cmp`:
|
||
// a total order over floats, so the sequence cannot depend on how the
|
||
// hash map happened to iterate. The merge order below is derived from
|
||
// this, and an unstable merge order is an unstable label field.
|
||
adjacency.sort_by(|x, y| {
|
||
x.saddle
|
||
.total_cmp(&y.saddle)
|
||
.then(x.a.cmp(&y.a))
|
||
.then(x.b.cmp(&y.b))
|
||
});
|
||
|
||
Self {
|
||
width,
|
||
height,
|
||
labels,
|
||
region_count,
|
||
adjacency,
|
||
}
|
||
}
|
||
|
||
/// Apply a cut's grouping to every pixel, for display or masking.
|
||
pub fn apply(&self, grouping: &[u32]) -> Vec<u32> {
|
||
self.labels.iter().map(|&l| grouping[l as usize]).collect()
|
||
}
|
||
}
|
||
|
||
/// One merge in the hierarchy.
|
||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||
pub struct Merge {
|
||
pub a: u32,
|
||
pub b: u32,
|
||
pub saddle: f32,
|
||
}
|
||
|
||
/// The merge order over a [`RegionField`] — the whole hierarchy.
|
||
///
|
||
/// Kruskal over the adjacency edges: sort by saddle, union-find, and record
|
||
/// each union that actually joined two distinct sets. The recording *is* the
|
||
/// dendrogram, which is why the multiscale part costs one `Vec` rather than a
|
||
/// second algorithm.
|
||
#[derive(Debug, Clone, PartialEq)]
|
||
pub struct MergeTree {
|
||
/// In merge order, so saddles are non-decreasing.
|
||
pub merges: Vec<Merge>,
|
||
pub region_count: usize,
|
||
}
|
||
|
||
impl MergeTree {
|
||
pub fn build(field: &RegionField) -> Self {
|
||
let mut uf = UnionFind::new(field.region_count);
|
||
let mut merges = Vec::new();
|
||
|
||
// `field.adjacency` is already sorted by (saddle, a, b).
|
||
for e in &field.adjacency {
|
||
if uf.union(e.a as usize, e.b as usize) {
|
||
merges.push(Merge {
|
||
a: e.a,
|
||
b: e.b,
|
||
saddle: e.saddle,
|
||
});
|
||
}
|
||
}
|
||
|
||
Self {
|
||
merges,
|
||
region_count: field.region_count,
|
||
}
|
||
}
|
||
|
||
/// The grouping after applying every merge weaker than `threshold`.
|
||
///
|
||
/// Returns one group id per original region, compacted to `0..groups`.
|
||
pub fn cut_at(&self, threshold: f32) -> Vec<u32> {
|
||
self.cut(self.merges.iter().take_while(|m| m.saddle <= threshold))
|
||
}
|
||
|
||
/// The grouping with at most `target` groups.
|
||
///
|
||
/// The interface a scroll wheel wants: thresholds are in gradient units
|
||
/// and mean nothing to anyone, where "about two hundred regions" is a
|
||
/// granularity a person can ask for.
|
||
pub fn cut_to(&self, target: usize) -> Vec<u32> {
|
||
let target = target.max(1);
|
||
let take = self.region_count.saturating_sub(target);
|
||
self.cut(self.merges.iter().take(take))
|
||
}
|
||
|
||
fn cut<'a>(&self, merges: impl Iterator<Item = &'a Merge>) -> Vec<u32> {
|
||
let mut uf = UnionFind::new(self.region_count);
|
||
for m in merges {
|
||
uf.union(m.a as usize, m.b as usize);
|
||
}
|
||
|
||
// Compact roots in region order, so a cut's numbering is stable and
|
||
// independent of the merge sequence that produced it.
|
||
let mut seen: HashMap<usize, u32> = HashMap::new();
|
||
(0..self.region_count)
|
||
.map(|r| {
|
||
let root = uf.find(r);
|
||
let next = seen.len() as u32;
|
||
*seen.entry(root).or_insert(next)
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
/// How many groups `cut_to` / `cut_at` would leave, without building one.
|
||
pub fn groups_at(&self, threshold: f32) -> usize {
|
||
let merged = self.merges.iter().filter(|m| m.saddle <= threshold).count();
|
||
self.region_count - merged
|
||
}
|
||
}
|
||
|
||
struct UnionFind {
|
||
parent: Vec<usize>,
|
||
rank: Vec<u8>,
|
||
}
|
||
|
||
impl UnionFind {
|
||
fn new(n: usize) -> Self {
|
||
Self {
|
||
parent: (0..n).collect(),
|
||
rank: vec![0; n],
|
||
}
|
||
}
|
||
|
||
fn find(&mut self, mut x: usize) -> usize {
|
||
while self.parent[x] != x {
|
||
// Path halving — the compression that keeps this near-linear
|
||
// without the second pass full compression needs.
|
||
self.parent[x] = self.parent[self.parent[x]];
|
||
x = self.parent[x];
|
||
}
|
||
x
|
||
}
|
||
|
||
/// Returns whether this union joined two distinct sets.
|
||
fn union(&mut self, a: usize, b: usize) -> bool {
|
||
let (ra, rb) = (self.find(a), self.find(b));
|
||
if ra == rb {
|
||
return false;
|
||
}
|
||
// Union by rank, but with a deterministic tie-break: equal ranks
|
||
// attach the higher index under the lower. Without it the tree shape
|
||
// depends on argument order, and `find` would return a different
|
||
// representative for the same input on a different day.
|
||
let (lo, hi) = if self.rank[ra] > self.rank[rb] {
|
||
(ra, rb)
|
||
} else if self.rank[rb] > self.rank[ra] {
|
||
(rb, ra)
|
||
} else {
|
||
let (lo, hi) = (ra.min(rb), ra.max(rb));
|
||
self.rank[lo] += 1;
|
||
(lo, hi)
|
||
};
|
||
self.parent[hi] = lo;
|
||
true
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
/// A 4×2 field split down the middle, with a weak gap in the boundary.
|
||
///
|
||
/// labels 0 0 1 1 gradient 0 5 5 0
|
||
/// 0 0 1 1 0 2 2 0
|
||
///
|
||
/// The saddle is 2, not 5: the bottom row is where the ridge is lowest.
|
||
fn split_field() -> RegionField {
|
||
let roots = vec![0, 0, 2, 2, 0, 0, 2, 2];
|
||
let gradient = vec![0.0, 5.0, 5.0, 0.0, 0.0, 2.0, 2.0, 0.0];
|
||
RegionField::from_roots(&roots, &gradient, 4, 2)
|
||
}
|
||
|
||
#[test]
|
||
fn roots_compact_to_dense_ids() {
|
||
// Basin roots are pixel indices and so are sparse and arbitrary.
|
||
// Anything that indexes by region needs them dense.
|
||
let f = split_field();
|
||
assert_eq!(f.region_count, 2);
|
||
assert_eq!(f.labels, vec![0, 0, 1, 1, 0, 0, 1, 1]);
|
||
}
|
||
|
||
#[test]
|
||
fn the_saddle_is_the_lowest_pass_not_the_typical_one() {
|
||
// The property that makes the hierarchy match human judgement: two
|
||
// regions joined by one weak gap belong together, however strong the
|
||
// rest of the boundary is.
|
||
let f = split_field();
|
||
assert_eq!(f.adjacency.len(), 1);
|
||
assert_eq!(f.adjacency[0].saddle, 2.0);
|
||
}
|
||
|
||
#[test]
|
||
fn a_uniform_image_is_one_region() {
|
||
// No minima to separate, so nothing to merge — and the tree must not
|
||
// invent a merge it cannot justify.
|
||
let f = RegionField::from_roots(&[0; 16], &[0.0; 16], 4, 4);
|
||
assert_eq!(f.region_count, 1);
|
||
assert!(f.adjacency.is_empty());
|
||
assert!(MergeTree::build(&f).merges.is_empty());
|
||
}
|
||
|
||
/// Three regions in a row: 0 | 1 | 2, with the 1–2 boundary weaker.
|
||
fn chain_field() -> RegionField {
|
||
let roots = vec![0, 0, 2, 2, 4, 4];
|
||
let gradient = vec![0.0, 9.0, 0.0, 3.0, 0.0, 0.0];
|
||
RegionField::from_roots(&roots, &gradient, 6, 1)
|
||
}
|
||
|
||
#[test]
|
||
fn the_weakest_boundary_merges_first() {
|
||
// The whole basis of the granularity ladder: walking up the tree must
|
||
// dissolve the least convincing edge before a strong one.
|
||
let tree = MergeTree::build(&chain_field());
|
||
assert_eq!(tree.merges.len(), 2);
|
||
assert_eq!(tree.merges[0].saddle, 3.0);
|
||
assert_eq!(tree.merges[1].saddle, 9.0);
|
||
assert_eq!((tree.merges[0].a, tree.merges[0].b), (1, 2));
|
||
}
|
||
|
||
#[test]
|
||
fn merges_are_ordered_by_saddle() {
|
||
let tree = MergeTree::build(&chain_field());
|
||
assert!(
|
||
tree.merges.windows(2).all(|w| w[0].saddle <= w[1].saddle),
|
||
"a cut at a threshold is only meaningful if merges are ordered"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn a_cut_walks_from_every_region_to_one() {
|
||
let tree = MergeTree::build(&chain_field());
|
||
|
||
// The bottom of the ladder: nothing merged.
|
||
let fine = tree.cut_to(3);
|
||
assert_eq!(fine, vec![0, 1, 2]);
|
||
|
||
// The middle: the weak boundary is gone, the strong one survives.
|
||
let mid = tree.cut_to(2);
|
||
assert_eq!(mid[1], mid[2], "the weak boundary should have dissolved");
|
||
assert_ne!(mid[0], mid[1], "the strong boundary should survive");
|
||
|
||
// The top: one region.
|
||
let coarse = tree.cut_to(1);
|
||
assert_eq!(coarse, vec![0, 0, 0]);
|
||
}
|
||
|
||
#[test]
|
||
fn cut_to_is_clamped_rather_than_panicking() {
|
||
// A scroll wheel runs past both ends of the ladder, and neither end
|
||
// is an error.
|
||
let tree = MergeTree::build(&chain_field());
|
||
assert_eq!(tree.cut_to(0), tree.cut_to(1));
|
||
assert_eq!(tree.cut_to(99), vec![0, 1, 2]);
|
||
}
|
||
|
||
#[test]
|
||
fn cut_at_and_cut_to_agree() {
|
||
let tree = MergeTree::build(&chain_field());
|
||
// Above the weak saddle, below the strong one.
|
||
assert_eq!(tree.cut_at(5.0), tree.cut_to(2));
|
||
assert_eq!(tree.groups_at(5.0), 2);
|
||
}
|
||
|
||
#[test]
|
||
fn a_grouping_maps_back_onto_pixels() {
|
||
let f = chain_field();
|
||
let tree = MergeTree::build(&f);
|
||
let px = f.apply(&tree.cut_to(2));
|
||
assert_eq!(px, vec![0, 0, 1, 1, 1, 1]);
|
||
}
|
||
|
||
#[test]
|
||
fn the_same_input_gives_a_bit_identical_result() {
|
||
// M5 in miniature. The CPU half must contribute no nondeterminism of
|
||
// its own, or there is no point asking whether the GPU half does:
|
||
// hash map iteration order is the obvious way to fail this, which is
|
||
// why both the adjacency list and the cut numbering are sorted rather
|
||
// than taken as the map yields them.
|
||
let (a, b) = (chain_field(), chain_field());
|
||
assert_eq!(a, b);
|
||
|
||
let (ta, tb) = (MergeTree::build(&a), MergeTree::build(&b));
|
||
assert_eq!(ta, tb);
|
||
assert_eq!(ta.cut_to(2), tb.cut_to(2));
|
||
}
|
||
|
||
#[test]
|
||
fn a_ladder_over_many_regions_is_monotone() {
|
||
// Region counts must fall as you climb, with no level skipped —
|
||
// otherwise a scroll step could jump past the granularity someone
|
||
// wanted.
|
||
let n = 32;
|
||
let roots: Vec<u32> = (0..n).map(|i| i as u32).collect();
|
||
let gradient: Vec<f32> = (0..n).map(|i| (i % 7) as f32).collect();
|
||
let f = RegionField::from_roots(&roots, &gradient, n, 1);
|
||
let tree = MergeTree::build(&f);
|
||
|
||
let mut last = usize::MAX;
|
||
for target in (1..=n).rev() {
|
||
let groups = tree.cut_to(target).iter().max().map(|m| *m as usize + 1);
|
||
let groups = groups.unwrap_or(0);
|
||
assert_eq!(groups, target, "cut_to({target}) should leave {target}");
|
||
assert!(groups < last);
|
||
last = groups;
|
||
}
|
||
}
|
||
}
|