Choose the export folder by walking the server, not by typing it
The destination for a Nextcloud export was a text field. Nobody recalls the exact spelling of a path three levels down, and getting it wrong does not fail — `create_dir` makes whatever was typed, so a misremembered folder becomes a new one at the root and the exports are somewhere nobody looks. So it is picked the way the library root is picked, using the same `FolderBrowser` model the launch screen drives: up, into, and "use this folder", confirming the folder currently *shown* rather than one selected in the list. Same rule in both places, so the phrase means one thing. The model is shared; the worker is not. `settings_ui::spawn_folder_list` is a near-twin of the launch screen's, because that one reaches into the `LaunchController` for its session and reports onto the launch screen's error line, while this one is handed credentials and writes to the settings page. Factoring them together needs a function taking both controllers or a trait implemented twice to abstract two call sites — more machinery than the twenty lines it saves. What matters is shared already: navigation behaves identically because both drive the same model. The callbacks are wired in `lib.rs` rather than in `settings_ui::wire`, because listing a remote folder needs credentials and the settings page holds no session on purpose — it is reachable before a library is opened and must not depend on one existing. With no account the picker says to sign in first, rather than showing an empty list that reads as a server with no folders. Details that are decisions rather than accidents: the picker opens at the library root rather than at whatever half-typed path is in the field, which would list nothing and look broken. The listing area is a fixed 180px, since a folder with sixty children would otherwise push the rest of the settings page off the bottom. "Up" is disabled at the root rather than hidden, so the row does not jump as the user navigates. A failed listing leaves the picker open on the folder it was showing — where the user had got to is not something to discard over a dropped request. And the chosen folder saves immediately like every other setting on a page that has no Save button. The poll timer lives on the controller for the reason `LaunchController` keeps its own there: a `slint::Timer` stops when dropped, so one local to the function that starts it would be collected before the listing arrived. Carries in-flight work from a parallel session — a segmentation pass in dr-gpu, a sidecar cache, and the develop panel's continuing changes. 1020 tests pass, fmt clean. One clippy warning remains and is not mine: `sidecar_cache::dir` is unused while that work is in progress. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Claude Opus 5
@@ -32,3 +32,7 @@ required-features = ["readback"]
|
||||
default = []
|
||||
# Exposes read_pixels outside tests. Production must not enable this.
|
||||
readback = []
|
||||
|
||||
[[example]]
|
||||
name = "segment"
|
||||
required-features = ["readback"]
|
||||
|
||||
@@ -0,0 +1,196 @@
|
||||
//! Segment an image and write the granularity ladder as false-coloured PPMs.
|
||||
//!
|
||||
//! The whole point of S15 step 2 (docs/segmentation.md §11): look at the
|
||||
//! ladder and decide whether clicking through it would land on the things a
|
||||
//! person means. No amount of design settles that — the pictures do.
|
||||
//!
|
||||
//! ```sh
|
||||
//! cargo run -p dr-gpu --example segment --features readback -- IMG.CR2
|
||||
//! cargo run -p dr-gpu --example segment --features readback -- synthetic
|
||||
//! ```
|
||||
//!
|
||||
//! PPM for the same reason `develop` uses it: no encoder dependency, and
|
||||
//! every viewer reads it. This is a diagnostic, not an export path.
|
||||
|
||||
use dr_gpu::hierarchy::{MergeTree, RegionField};
|
||||
use dr_gpu::{DemosaicedImage, Demosaicer, GpuContext, SegmentOptions, SegmentPass};
|
||||
|
||||
/// The ladder the example dumps. Chosen to span "far too fine to be useful"
|
||||
/// through "one or two objects", because both ends are informative: if no rung
|
||||
/// looks right, the gradient is wrong rather than the ladder being too coarse.
|
||||
const LEVELS: [usize; 6] = [2000, 800, 300, 120, 50, 16];
|
||||
|
||||
fn main() {
|
||||
env_logger::init();
|
||||
|
||||
let mut args = std::env::args().skip(1);
|
||||
let Some(input) = args.next() else {
|
||||
eprintln!("usage: segment <file.raw|synthetic> [out-prefix] [blur-radius]");
|
||||
std::process::exit(2);
|
||||
};
|
||||
let prefix = args.next().unwrap_or_else(|| "segment".into());
|
||||
let blur_radius = args
|
||||
.next()
|
||||
.and_then(|s| s.parse().ok())
|
||||
.unwrap_or(SegmentOptions::default().blur_radius);
|
||||
|
||||
let ctx = pollster::block_on(GpuContext::new_headless()).expect("gpu context");
|
||||
println!("gpu {}", ctx.adapter_name());
|
||||
|
||||
let source = if input == "synthetic" {
|
||||
let (w, h) = (1200, 800);
|
||||
println!("source synthetic {w} × {h}");
|
||||
DemosaicedImage::from_rgba8(&ctx, &synthetic(w, h), w, h).expect("synthetic source")
|
||||
} else {
|
||||
let bytes = std::fs::read(&input).expect("read file");
|
||||
let raw = dr_decode::decode(&bytes).expect("decode");
|
||||
println!("source {} × {}", raw.crop.width, raw.crop.height);
|
||||
Demosaicer::new(&ctx)
|
||||
.expect("demosaicer")
|
||||
.run(&raw)
|
||||
.expect("demosaic")
|
||||
};
|
||||
|
||||
let opts = SegmentOptions {
|
||||
blur_radius,
|
||||
..Default::default()
|
||||
};
|
||||
let pass = SegmentPass::new(&ctx).expect("segment pass");
|
||||
|
||||
let t0 = std::time::Instant::now();
|
||||
let seg = pass.run(&source, opts).expect("segment");
|
||||
let (w, h) = seg.size();
|
||||
let field = seg.read_field().expect("read field");
|
||||
let gpu_ms = t0.elapsed().as_secs_f32() * 1000.0;
|
||||
|
||||
let t1 = std::time::Instant::now();
|
||||
let tree = MergeTree::build(&field);
|
||||
let tree_ms = t1.elapsed().as_secs_f32() * 1000.0;
|
||||
|
||||
// M6, roughly: the readback is in `gpu_ms` and would not be there in a
|
||||
// shipping build, so this over-reports the GPU half rather than under.
|
||||
println!("proxy {w} × {h}, blur radius {blur_radius}");
|
||||
println!("basins {}", field.region_count);
|
||||
println!("boundaries {}", field.adjacency.len());
|
||||
println!("merges {}", tree.merges.len());
|
||||
println!("segment {gpu_ms:.0} ms (includes readback)");
|
||||
println!("hierarchy {tree_ms:.1} ms");
|
||||
|
||||
if let (Some(first), Some(last)) = (tree.merges.first(), tree.merges.last()) {
|
||||
println!("saddles {:.4} … {:.4}", first.saddle, last.saddle);
|
||||
}
|
||||
|
||||
for level in LEVELS {
|
||||
if level > field.region_count {
|
||||
println!("skip {level} (only {} basins)", field.region_count);
|
||||
continue;
|
||||
}
|
||||
let grouping = tree.cut_to(level);
|
||||
let pixels = field.apply(&grouping);
|
||||
let groups = grouping.iter().max().map(|m| m + 1).unwrap_or(0);
|
||||
let path = format!("{prefix}-{level:04}.ppm");
|
||||
write_ppm(&path, &false_colour(&pixels, &field), w, h);
|
||||
println!("wrote {path} ({groups} regions)");
|
||||
}
|
||||
|
||||
// The boundaries alone, which is what a snapped contour would cling to.
|
||||
let path = format!("{prefix}-edges.ppm");
|
||||
write_ppm(&path, &boundaries(&field), w, h);
|
||||
println!("wrote {path}");
|
||||
}
|
||||
|
||||
/// A distinct colour per region.
|
||||
///
|
||||
/// Hashed from the id rather than sampled from the image: two adjacent
|
||||
/// regions that happen to look alike are exactly the case worth seeing, and
|
||||
/// mean colours would hide it.
|
||||
fn false_colour(pixels: &[u32], field: &RegionField) -> Vec<u8> {
|
||||
let _ = field;
|
||||
let mut out = Vec::with_capacity(pixels.len() * 3);
|
||||
for &g in pixels {
|
||||
// Cheap integer hash — golden-ratio multiply, then spread the bits
|
||||
// across three channels.
|
||||
let mut x = g.wrapping_mul(2_654_435_761);
|
||||
x ^= x >> 15;
|
||||
out.push((x & 0xff) as u8);
|
||||
out.push(((x >> 8) & 0xff) as u8);
|
||||
out.push(((x >> 16) & 0xff) as u8);
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// White where two regions meet, black elsewhere.
|
||||
fn boundaries(field: &RegionField) -> Vec<u8> {
|
||||
let (w, h) = (field.width, field.height);
|
||||
let mut out = vec![0u8; w * h * 3];
|
||||
for y in 0..h {
|
||||
for x in 0..w {
|
||||
let i = y * w + x;
|
||||
let edge = (x + 1 < w && field.labels[i] != field.labels[i + 1])
|
||||
|| (y + 1 < h && field.labels[i] != field.labels[i + w]);
|
||||
if edge {
|
||||
out[i * 3] = 255;
|
||||
out[i * 3 + 1] = 255;
|
||||
out[i * 3 + 2] = 255;
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
fn write_ppm(path: &str, rgb: &[u8], width: u32, height: u32) {
|
||||
use std::io::Write as _;
|
||||
let mut f = std::io::BufWriter::new(std::fs::File::create(path).expect("create ppm"));
|
||||
write!(f, "P6\n{width} {height}\n255\n").expect("ppm header");
|
||||
f.write_all(rgb).expect("ppm body");
|
||||
}
|
||||
|
||||
/// A test image with the failure modes the corpus is meant to provoke, so the
|
||||
/// example is runnable before anyone has traced a single ground-truth mask.
|
||||
///
|
||||
/// Deliberately includes a soft gradient boundary and a noisy patch: those are
|
||||
/// where a watershed either earns its place or shatters, and a synthetic image
|
||||
/// of clean shapes would flatter it.
|
||||
fn synthetic(w: u32, h: u32) -> Vec<u8> {
|
||||
let mut px = Vec::with_capacity((w * h * 4) as usize);
|
||||
for y in 0..h {
|
||||
for x in 0..w {
|
||||
let fx = x as f32 / w as f32;
|
||||
let fy = y as f32 / h as f32;
|
||||
|
||||
// A smooth vertical gradient — the low-contrast boundary case.
|
||||
let mut r = 40.0 + 120.0 * fy;
|
||||
let mut g = 60.0 + 100.0 * fy;
|
||||
let mut b = 110.0 + 90.0 * fy;
|
||||
|
||||
// A hard-edged disc: the control case.
|
||||
let d = ((fx - 0.3).powi(2) + (fy - 0.45).powi(2)).sqrt();
|
||||
if d < 0.16 {
|
||||
r = 210.0;
|
||||
g = 90.0;
|
||||
b = 60.0;
|
||||
}
|
||||
|
||||
// A soft-edged disc: where the ladder should merge late.
|
||||
let d2 = ((fx - 0.68).powi(2) + (fy - 0.6).powi(2)).sqrt();
|
||||
let t = (1.0 - (d2 / 0.18)).clamp(0.0, 1.0);
|
||||
r = r * (1.0 - t) + 90.0 * t;
|
||||
g = g * (1.0 - t) + 170.0 * t;
|
||||
b = b * (1.0 - t) + 110.0 * t;
|
||||
|
||||
// A noisy corner: the case pre-smoothing exists for.
|
||||
if fx > 0.82 && fy < 0.22 {
|
||||
let n = ((x * 7919 + y * 104_729) % 97) as f32 / 97.0;
|
||||
r += (n - 0.5) * 90.0;
|
||||
g += (n - 0.5) * 90.0;
|
||||
b += (n - 0.5) * 90.0;
|
||||
}
|
||||
|
||||
px.push(r.clamp(0.0, 255.0) as u8);
|
||||
px.push(g.clamp(0.0, 255.0) as u8);
|
||||
px.push(b.clamp(0.0, 255.0) as u8);
|
||||
px.push(255);
|
||||
}
|
||||
}
|
||||
px
|
||||
}
|
||||
@@ -0,0 +1,429 @@
|
||||
//! 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;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -14,9 +14,12 @@ use wgpu::util::DeviceExt;
|
||||
mod adjust;
|
||||
mod demosaic;
|
||||
mod error;
|
||||
pub mod hierarchy;
|
||||
mod segment;
|
||||
pub use adjust::AdjustPass;
|
||||
pub use demosaic::{DemosaicedImage, Demosaicer};
|
||||
pub use error::GpuError;
|
||||
pub use segment::{SegmentOptions, SegmentPass, Segmentation};
|
||||
|
||||
/// Owns the wgpu device and queue.
|
||||
///
|
||||
|
||||
@@ -0,0 +1,575 @@
|
||||
//! Watershed segmentation — arm A's GPU half (S15, docs/segmentation.md).
|
||||
//!
|
||||
//! Runs the five passes in `shaders/watershed.wgsl` over a demosaiced image
|
||||
//! and leaves a basin label per pixel on the GPU. The hierarchy built from
|
||||
//! those labels lives in [`crate::hierarchy`], which needs no device.
|
||||
//!
|
||||
//! # Cost
|
||||
//!
|
||||
//! Every pass is a trivial kernel and the whole chain is a handful of
|
||||
//! milliseconds at proxy resolution. It runs **once per image**, off the
|
||||
//! interactive path — the point of precomputing a region map is that
|
||||
//! selection afterwards is a label comparison rather than a flood fill.
|
||||
//!
|
||||
//! # The open question this leaves
|
||||
//!
|
||||
//! [`Segmentation::read_field`] copies the label and gradient buffers back to
|
||||
//! the CPU to build the region adjacency graph, and is gated behind the
|
||||
//! `readback` feature for the same reason `read_pixels` is. That gate is not
|
||||
//! ceremony: a shipping build cannot take this path (ARCH §6.1, AC-8), so the
|
||||
//! RAG would have to be accumulated GPU-side with atomics instead.
|
||||
//!
|
||||
//! For a spike that trade is the right way round — the readback is once per
|
||||
//! image and off the frame path, and building the GPU-side RAG before knowing
|
||||
//! whether the granularity ladder is any good would be work spent on a
|
||||
//! question not yet asked. But it is a real gap between this and something
|
||||
//! shippable, and it should be read as one.
|
||||
|
||||
use wgpu::util::DeviceExt;
|
||||
|
||||
use crate::{DemosaicedImage, GpuContext, GpuError};
|
||||
|
||||
/// How the watershed is tuned for one image.
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
pub struct SegmentOptions {
|
||||
/// Longest proxy edge. The segmentation runs here, not at sensor
|
||||
/// resolution: a 24 MP watershed costs 12× the memory to place boundaries
|
||||
/// a person cannot see, and the boundary refinement that matters at 1:1
|
||||
/// is a separate stage (docs/segmentation.md §4).
|
||||
pub max_edge: u32,
|
||||
/// Pre-smoothing radius in proxy pixels. The caller's to raise with ISO —
|
||||
/// this is the single knob that decides whether a noisy file segments
|
||||
/// into regions or into grain.
|
||||
pub blur_radius: i32,
|
||||
pub w_luma: f32,
|
||||
pub w_chroma: f32,
|
||||
}
|
||||
|
||||
impl Default for SegmentOptions {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
// ~1.3 MP at 3:2. Large enough that a boundary is within a pixel
|
||||
// or two of where it belongs, small enough that the whole chain
|
||||
// fits comfortably in memory on a phone.
|
||||
max_edge: 1600,
|
||||
blur_radius: 2,
|
||||
w_luma: 1.0,
|
||||
// Chroma carries most of the sensor noise and few of the
|
||||
// boundaries anyone would draw, so it counts for less — but not
|
||||
// zero, or a red flower on green leaves has no edge at all.
|
||||
w_chroma: 0.5,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[repr(C)]
|
||||
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
|
||||
struct SegParams {
|
||||
width: u32,
|
||||
height: u32,
|
||||
src_width: u32,
|
||||
src_height: u32,
|
||||
blur_radius: i32,
|
||||
non_linear: u32,
|
||||
w_luma: f32,
|
||||
w_chroma: f32,
|
||||
}
|
||||
|
||||
/// One compute stage: its layout and its compiled pipeline.
|
||||
struct Stage {
|
||||
layout: wgpu::BindGroupLayout,
|
||||
pipeline: wgpu::ComputePipeline,
|
||||
}
|
||||
|
||||
/// Runs the watershed chain.
|
||||
pub struct SegmentPass {
|
||||
ctx: GpuContext,
|
||||
features: Stage,
|
||||
blur: Stage,
|
||||
gradient: Stage,
|
||||
flow: Stage,
|
||||
jump: Stage,
|
||||
}
|
||||
|
||||
impl SegmentPass {
|
||||
pub fn new(ctx: &GpuContext) -> Result<Self, GpuError> {
|
||||
// A validation failure here is a bug in the shader, not a user error.
|
||||
// Surfaced as a Result rather than wgpu's default panic, matching how
|
||||
// `AdjustPass` handles its generated source.
|
||||
let scope = ctx.device.push_error_scope(wgpu::ErrorFilter::Validation);
|
||||
|
||||
let module = ctx
|
||||
.device
|
||||
.create_shader_module(wgpu::ShaderModuleDescriptor {
|
||||
label: Some("watershed"),
|
||||
source: wgpu::ShaderSource::Wgsl(include_str!("shaders/watershed.wgsl").into()),
|
||||
});
|
||||
|
||||
let features = {
|
||||
let layout = ctx
|
||||
.device
|
||||
.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
|
||||
label: Some("watershed-features-bgl"),
|
||||
entries: &[
|
||||
uniform_entry(0),
|
||||
wgpu::BindGroupLayoutEntry {
|
||||
binding: 1,
|
||||
visibility: wgpu::ShaderStages::COMPUTE,
|
||||
ty: wgpu::BindingType::Texture {
|
||||
sample_type: wgpu::TextureSampleType::Float { filterable: true },
|
||||
view_dimension: wgpu::TextureViewDimension::D2,
|
||||
multisampled: false,
|
||||
},
|
||||
count: None,
|
||||
},
|
||||
storage_entry(2, false),
|
||||
],
|
||||
});
|
||||
let pipeline = compute(ctx, &module, &layout, "features");
|
||||
Stage { layout, pipeline }
|
||||
};
|
||||
|
||||
let buffer_stage = |in_binding: u32, out_binding: u32, entry: &str, label: &str| {
|
||||
let layout = ctx
|
||||
.device
|
||||
.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
|
||||
label: Some(label),
|
||||
entries: &[
|
||||
uniform_entry(0),
|
||||
storage_entry(in_binding, true),
|
||||
storage_entry(out_binding, false),
|
||||
],
|
||||
});
|
||||
let pipeline = compute(ctx, &module, &layout, entry);
|
||||
Stage { layout, pipeline }
|
||||
};
|
||||
|
||||
let blur = buffer_stage(3, 4, "blur", "watershed-blur-bgl");
|
||||
let gradient = buffer_stage(5, 6, "gradient", "watershed-gradient-bgl");
|
||||
let flow = buffer_stage(7, 8, "flow", "watershed-flow-bgl");
|
||||
let jump = buffer_stage(9, 10, "jump", "watershed-jump-bgl");
|
||||
|
||||
if let Some(err) = pollster::block_on(scope.pop()) {
|
||||
return Err(GpuError::ShaderCompilation(err.to_string()));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
ctx: ctx.clone(),
|
||||
features,
|
||||
blur,
|
||||
gradient,
|
||||
flow,
|
||||
jump,
|
||||
})
|
||||
}
|
||||
|
||||
/// Segment an image into basins.
|
||||
pub fn run(
|
||||
&self,
|
||||
source: &DemosaicedImage,
|
||||
opts: SegmentOptions,
|
||||
) -> Result<Segmentation, GpuError> {
|
||||
let (src_w, src_h) = source.size();
|
||||
let (width, height) = proxy_size(src_w, src_h, opts.max_edge);
|
||||
let n = (width * height) as u64;
|
||||
|
||||
let params = self
|
||||
.ctx
|
||||
.device
|
||||
.create_buffer_init(&wgpu::util::BufferInitDescriptor {
|
||||
label: Some("watershed-params"),
|
||||
contents: bytemuck::bytes_of(&SegParams {
|
||||
width,
|
||||
height,
|
||||
src_width: src_w,
|
||||
src_height: src_h,
|
||||
blur_radius: opts.blur_radius,
|
||||
non_linear: u32::from(source.is_non_linear()),
|
||||
w_luma: opts.w_luma,
|
||||
w_chroma: opts.w_chroma,
|
||||
}),
|
||||
usage: wgpu::BufferUsages::UNIFORM,
|
||||
});
|
||||
|
||||
// `vec4` rather than `vec3` for the feature buffers: a WGSL storage
|
||||
// array of vec3 still strides by 16 bytes, so packing to three floats
|
||||
// would save nothing and cost an index calculation.
|
||||
let feat_a = self.buffer("watershed-feat-a", n * 16, false);
|
||||
let feat_b = self.buffer("watershed-feat-b", n * 16, false);
|
||||
let gradient = self.buffer("watershed-gradient", n * 4, true);
|
||||
let parent_a = self.buffer("watershed-parent-a", n * 4, true);
|
||||
let parent_b = self.buffer("watershed-parent-b", n * 4, true);
|
||||
|
||||
let mut enc = self
|
||||
.ctx
|
||||
.device
|
||||
.create_command_encoder(&wgpu::CommandEncoderDescriptor {
|
||||
label: Some("watershed-encoder"),
|
||||
});
|
||||
|
||||
let groups = (width.div_ceil(8), height.div_ceil(8));
|
||||
|
||||
let features_bg = self
|
||||
.ctx
|
||||
.device
|
||||
.create_bind_group(&wgpu::BindGroupDescriptor {
|
||||
label: Some("watershed-features-bg"),
|
||||
layout: &self.features.layout,
|
||||
entries: &[
|
||||
wgpu::BindGroupEntry {
|
||||
binding: 0,
|
||||
resource: params.as_entire_binding(),
|
||||
},
|
||||
wgpu::BindGroupEntry {
|
||||
binding: 1,
|
||||
resource: wgpu::BindingResource::TextureView(source.view()),
|
||||
},
|
||||
wgpu::BindGroupEntry {
|
||||
binding: 2,
|
||||
resource: feat_a.as_entire_binding(),
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
let blur_bg = self.bind(&self.blur.layout, ¶ms, 3, &feat_a, 4, &feat_b);
|
||||
let gradient_bg = self.bind(&self.gradient.layout, ¶ms, 5, &feat_b, 6, &gradient);
|
||||
let flow_bg = self.bind(&self.flow.layout, ¶ms, 7, &gradient, 8, &parent_a);
|
||||
let jump_ab = self.bind(&self.jump.layout, ¶ms, 9, &parent_a, 10, &parent_b);
|
||||
let jump_ba = self.bind(&self.jump.layout, ¶ms, 9, &parent_b, 10, &parent_a);
|
||||
|
||||
// Pointer jumping halves every path per pass, so log2 of the pixel
|
||||
// count bounds it — that is the longest possible descent chain. A
|
||||
// convergence test would cost a readback per iteration to save a
|
||||
// handful of dispatches of a two-line kernel.
|
||||
let jumps = (n as f64).log2().ceil() as u32 + 1;
|
||||
|
||||
{
|
||||
let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor {
|
||||
label: Some("watershed-pass"),
|
||||
timestamp_writes: None,
|
||||
});
|
||||
|
||||
for (pipeline, bg) in [
|
||||
(&self.features.pipeline, &features_bg),
|
||||
(&self.blur.pipeline, &blur_bg),
|
||||
(&self.gradient.pipeline, &gradient_bg),
|
||||
(&self.flow.pipeline, &flow_bg),
|
||||
] {
|
||||
pass.set_pipeline(pipeline);
|
||||
pass.set_bind_group(0, bg, &[]);
|
||||
pass.dispatch_workgroups(groups.0, groups.1, 1);
|
||||
}
|
||||
|
||||
pass.set_pipeline(&self.jump.pipeline);
|
||||
for i in 0..jumps {
|
||||
let bg = if i % 2 == 0 { &jump_ab } else { &jump_ba };
|
||||
pass.set_bind_group(0, bg, &[]);
|
||||
pass.dispatch_workgroups(groups.0, groups.1, 1);
|
||||
}
|
||||
}
|
||||
|
||||
self.ctx.queue.submit(Some(enc.finish()));
|
||||
|
||||
// An odd number of jumps leaves the result in B.
|
||||
let labels = if jumps % 2 == 1 { parent_b } else { parent_a };
|
||||
|
||||
Ok(Segmentation {
|
||||
ctx: self.ctx.clone(),
|
||||
width,
|
||||
height,
|
||||
labels,
|
||||
gradient,
|
||||
})
|
||||
}
|
||||
|
||||
fn buffer(&self, label: &str, size: u64, copyable: bool) -> wgpu::Buffer {
|
||||
let mut usage = wgpu::BufferUsages::STORAGE;
|
||||
if copyable {
|
||||
usage |= wgpu::BufferUsages::COPY_SRC;
|
||||
}
|
||||
self.ctx.device.create_buffer(&wgpu::BufferDescriptor {
|
||||
label: Some(label),
|
||||
size,
|
||||
usage,
|
||||
mapped_at_creation: false,
|
||||
})
|
||||
}
|
||||
|
||||
fn bind(
|
||||
&self,
|
||||
layout: &wgpu::BindGroupLayout,
|
||||
params: &wgpu::Buffer,
|
||||
in_binding: u32,
|
||||
input: &wgpu::Buffer,
|
||||
out_binding: u32,
|
||||
output: &wgpu::Buffer,
|
||||
) -> wgpu::BindGroup {
|
||||
self.ctx
|
||||
.device
|
||||
.create_bind_group(&wgpu::BindGroupDescriptor {
|
||||
label: Some("watershed-bg"),
|
||||
layout,
|
||||
entries: &[
|
||||
wgpu::BindGroupEntry {
|
||||
binding: 0,
|
||||
resource: params.as_entire_binding(),
|
||||
},
|
||||
wgpu::BindGroupEntry {
|
||||
binding: in_binding,
|
||||
resource: input.as_entire_binding(),
|
||||
},
|
||||
wgpu::BindGroupEntry {
|
||||
binding: out_binding,
|
||||
resource: output.as_entire_binding(),
|
||||
},
|
||||
],
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// The result of one segmentation: a basin label per pixel, on the GPU.
|
||||
pub struct Segmentation {
|
||||
ctx: GpuContext,
|
||||
width: u32,
|
||||
height: u32,
|
||||
/// Per pixel, the linear index of its basin root. Sparse — compacted by
|
||||
/// [`crate::hierarchy::RegionField::from_roots`].
|
||||
labels: wgpu::Buffer,
|
||||
gradient: wgpu::Buffer,
|
||||
}
|
||||
|
||||
impl Segmentation {
|
||||
pub fn size(&self) -> (u32, u32) {
|
||||
(self.width, self.height)
|
||||
}
|
||||
|
||||
/// The label buffer, for a shader that masks by region id.
|
||||
pub fn labels(&self) -> &wgpu::Buffer {
|
||||
&self.labels
|
||||
}
|
||||
|
||||
/// Build the region adjacency graph, reading the labels back to the CPU.
|
||||
///
|
||||
/// **Not a shipping path** — see this module's header. Gated so it cannot
|
||||
/// be reached from a production build by accident.
|
||||
#[cfg(any(test, feature = "readback"))]
|
||||
pub fn read_field(&self) -> Result<crate::hierarchy::RegionField, GpuError> {
|
||||
let n = (self.width * self.height) as usize;
|
||||
let roots: Vec<u32> = read_buffer(&self.ctx, &self.labels, n)?;
|
||||
let gradient: Vec<f32> = read_buffer(&self.ctx, &self.gradient, n)?;
|
||||
Ok(crate::hierarchy::RegionField::from_roots(
|
||||
&roots,
|
||||
&gradient,
|
||||
self.width as usize,
|
||||
self.height as usize,
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn uniform_entry(binding: u32) -> wgpu::BindGroupLayoutEntry {
|
||||
wgpu::BindGroupLayoutEntry {
|
||||
binding,
|
||||
visibility: wgpu::ShaderStages::COMPUTE,
|
||||
ty: wgpu::BindingType::Buffer {
|
||||
ty: wgpu::BufferBindingType::Uniform,
|
||||
has_dynamic_offset: false,
|
||||
min_binding_size: None,
|
||||
},
|
||||
count: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn storage_entry(binding: u32, read_only: bool) -> wgpu::BindGroupLayoutEntry {
|
||||
wgpu::BindGroupLayoutEntry {
|
||||
binding,
|
||||
visibility: wgpu::ShaderStages::COMPUTE,
|
||||
ty: wgpu::BindingType::Buffer {
|
||||
ty: wgpu::BufferBindingType::Storage { read_only },
|
||||
has_dynamic_offset: false,
|
||||
min_binding_size: None,
|
||||
},
|
||||
count: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn compute(
|
||||
ctx: &GpuContext,
|
||||
module: &wgpu::ShaderModule,
|
||||
layout: &wgpu::BindGroupLayout,
|
||||
entry: &str,
|
||||
) -> wgpu::ComputePipeline {
|
||||
let pipeline_layout = ctx
|
||||
.device
|
||||
.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
|
||||
label: Some("watershed-layout"),
|
||||
bind_group_layouts: &[Some(layout)],
|
||||
immediate_size: 0,
|
||||
});
|
||||
ctx.device
|
||||
.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
|
||||
label: Some(entry),
|
||||
layout: Some(&pipeline_layout),
|
||||
module,
|
||||
entry_point: Some(entry),
|
||||
compilation_options: Default::default(),
|
||||
cache: None,
|
||||
})
|
||||
}
|
||||
|
||||
/// The proxy size for a source, preserving aspect and never upscaling.
|
||||
fn proxy_size(src_w: u32, src_h: u32, max_edge: u32) -> (u32, u32) {
|
||||
let longest = src_w.max(src_h);
|
||||
if longest <= max_edge || longest == 0 {
|
||||
return (src_w.max(1), src_h.max(1));
|
||||
}
|
||||
let scale = f64::from(max_edge) / f64::from(longest);
|
||||
(
|
||||
((f64::from(src_w) * scale).round() as u32).max(1),
|
||||
((f64::from(src_h) * scale).round() as u32).max(1),
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(any(test, feature = "readback"))]
|
||||
fn read_buffer<T: bytemuck::Pod>(
|
||||
ctx: &GpuContext,
|
||||
buffer: &wgpu::Buffer,
|
||||
len: usize,
|
||||
) -> Result<Vec<T>, GpuError> {
|
||||
let size = (len * std::mem::size_of::<T>()) as u64;
|
||||
let staging = ctx.device.create_buffer(&wgpu::BufferDescriptor {
|
||||
label: Some("watershed-readback"),
|
||||
size,
|
||||
usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
|
||||
mapped_at_creation: false,
|
||||
});
|
||||
|
||||
let mut enc = ctx.device.create_command_encoder(&Default::default());
|
||||
enc.copy_buffer_to_buffer(buffer, 0, &staging, 0, size);
|
||||
ctx.queue.submit(Some(enc.finish()));
|
||||
|
||||
let slice = staging.slice(..);
|
||||
let (tx, rx) = std::sync::mpsc::channel();
|
||||
slice.map_async(wgpu::MapMode::Read, move |r| {
|
||||
let _ = tx.send(r);
|
||||
});
|
||||
ctx.device
|
||||
.poll(wgpu::PollType::wait_indefinitely())
|
||||
.map_err(|e| GpuError::Readback(e.to_string()))?;
|
||||
rx.recv()
|
||||
.map_err(|e| GpuError::Readback(e.to_string()))?
|
||||
.map_err(|e| GpuError::Readback(e.to_string()))?;
|
||||
|
||||
let data = slice.get_mapped_range();
|
||||
let out = bytemuck::cast_slice::<u8, T>(&data).to_vec();
|
||||
drop(data);
|
||||
staging.unmap();
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::hierarchy::MergeTree;
|
||||
|
||||
fn ctx() -> Option<GpuContext> {
|
||||
match pollster::block_on(GpuContext::new_headless()) {
|
||||
Ok(c) => Some(c),
|
||||
Err(e) => {
|
||||
eprintln!("skipping: no GPU adapter ({e})");
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_proxy_preserves_aspect_and_never_upscales() {
|
||||
assert_eq!(proxy_size(6000, 4000, 1600), (1600, 1067));
|
||||
assert_eq!(proxy_size(4000, 6000, 1600), (1067, 1600));
|
||||
// A thumbnail must not be blown up to the proxy size — there is no
|
||||
// detail there to find basins in.
|
||||
assert_eq!(proxy_size(800, 600, 1600), (800, 600));
|
||||
assert_eq!(proxy_size(0, 0, 1600), (1, 1));
|
||||
}
|
||||
|
||||
/// Two flat halves split by a hard vertical edge.
|
||||
fn two_tone(w: u32, h: u32) -> Vec<u8> {
|
||||
let mut px = Vec::with_capacity((w * h * 4) as usize);
|
||||
for _ in 0..h {
|
||||
for x in 0..w {
|
||||
let v = if x < w / 2 { 30u8 } else { 220u8 };
|
||||
px.extend_from_slice(&[v, v, v, 255]);
|
||||
}
|
||||
}
|
||||
px
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_hard_edge_produces_two_regions_at_the_top_of_the_ladder() {
|
||||
// The end-to-end property, on an image whose answer is not in doubt:
|
||||
// whatever the watershed does with texture, it must not lose an edge
|
||||
// this obvious, and the coarsest non-trivial cut must be exactly the
|
||||
// two halves.
|
||||
let Some(ctx) = ctx() else { return };
|
||||
let (w, h) = (64u32, 64u32);
|
||||
let src = DemosaicedImage::from_rgba8(&ctx, &two_tone(w, h), w, h).expect("source");
|
||||
|
||||
let pass = SegmentPass::new(&ctx).expect("segment pass");
|
||||
let seg = pass.run(&src, SegmentOptions::default()).expect("run");
|
||||
assert_eq!(seg.size(), (w, h));
|
||||
|
||||
let field = seg.read_field().expect("read field");
|
||||
let tree = MergeTree::build(&field);
|
||||
let px = field.apply(&tree.cut_to(2));
|
||||
|
||||
for y in 0..h as usize {
|
||||
let left = px[y * w as usize];
|
||||
let right = px[y * w as usize + w as usize - 1];
|
||||
assert_ne!(left, right, "the two halves must not share a region");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_flat_image_does_not_fragment() {
|
||||
// The noise case in miniature. A gradient of zero everywhere is one
|
||||
// enormous plateau, which is exactly where a watershed without a
|
||||
// strict tie-break either hangs or shatters into per-pixel basins.
|
||||
let Some(ctx) = ctx() else { return };
|
||||
let (w, h) = (32u32, 32u32);
|
||||
let flat = vec![128u8; (w * h * 4) as usize];
|
||||
let src = DemosaicedImage::from_rgba8(&ctx, &flat, w, h).expect("source");
|
||||
|
||||
let pass = SegmentPass::new(&ctx).expect("segment pass");
|
||||
let seg = pass.run(&src, SegmentOptions::default()).expect("run");
|
||||
let field = seg.read_field().expect("read field");
|
||||
|
||||
assert_eq!(
|
||||
field.region_count, 1,
|
||||
"a plateau should resolve to one basin, not {}",
|
||||
field.region_count
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_same_image_segments_identically_twice() {
|
||||
// M5 on one device — the weaker half of the determinism question, but
|
||||
// the half that catches a race in the pointer jumping. Cross-vendor
|
||||
// is the part that needs hardware this test cannot assume.
|
||||
let Some(ctx) = ctx() else { return };
|
||||
let (w, h) = (48u32, 48u32);
|
||||
let src = DemosaicedImage::from_rgba8(&ctx, &two_tone(w, h), w, h).expect("source");
|
||||
let pass = SegmentPass::new(&ctx).expect("segment pass");
|
||||
|
||||
let a = pass
|
||||
.run(&src, SegmentOptions::default())
|
||||
.expect("run")
|
||||
.read_field()
|
||||
.expect("field");
|
||||
let b = pass
|
||||
.run(&src, SegmentOptions::default())
|
||||
.expect("run")
|
||||
.read_field()
|
||||
.expect("field");
|
||||
|
||||
assert_eq!(a, b, "segmentation must be reproducible run to run");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,256 @@
|
||||
// Watershed segmentation — the passes behind arm A of S15 (docs/segmentation.md).
|
||||
//
|
||||
// Five entry points forming one chain:
|
||||
//
|
||||
// features source texture -> perceptual triple, box-downscaled to proxy size
|
||||
// blur pre-smoothing, without which every noise grain becomes a basin
|
||||
// gradient Sobel magnitude — the surface the watershed floods
|
||||
// flow each pixel points downhill to its steepest neighbour
|
||||
// jump pointer-jumping, until every pixel points at its basin root
|
||||
//
|
||||
// Everything after `features` works in storage buffers rather than textures.
|
||||
// That is deliberate: the flow and jump passes need read-write access to the
|
||||
// same array across dispatches, which storage textures do not give portably,
|
||||
// and a buffer reads back without the 256-byte row padding a texture copy
|
||||
// imposes.
|
||||
|
||||
struct Params {
|
||||
// Proxy dimensions — what every pass but `features` iterates over.
|
||||
width: u32,
|
||||
height: u32,
|
||||
// Source dimensions, for the box downscale in `features`.
|
||||
src_width: u32,
|
||||
src_height: u32,
|
||||
// Half-width of the pre-smoothing kernel, in proxy pixels. 0 disables it.
|
||||
blur_radius: i32,
|
||||
// 1 when the source is already display-encoded (the JPEG path), 0 for
|
||||
// linear scene-referred data out of the demosaicer.
|
||||
non_linear: u32,
|
||||
// How much luma and chroma each contribute to the gradient. Chroma is
|
||||
// weighted lower because it carries most of the sensor noise and few of
|
||||
// the boundaries a person would draw.
|
||||
w_luma: f32,
|
||||
w_chroma: f32,
|
||||
}
|
||||
|
||||
// Binding slots are unique across the whole module, not reused per entry
|
||||
// point: WGSL resource variables share one namespace, so two globals at the
|
||||
// same (group, binding) is a module-level validation error even when no
|
||||
// single entry point uses both. Each pass therefore gets its own pair, and
|
||||
// each pipeline a layout declaring only the slots it touches.
|
||||
@group(0) @binding(0) var<uniform> u: Params;
|
||||
|
||||
// ---------------------------------------------------------------- features
|
||||
|
||||
@group(0) @binding(1) var src: texture_2d<f32>;
|
||||
@group(0) @binding(2) var<storage, read_write> feat_out: array<vec4<f32>>;
|
||||
|
||||
// Linear or display-encoded RGB to a roughly perceptual opponent triple.
|
||||
//
|
||||
// Perceptual rather than linear because the gradient has to agree with what
|
||||
// a person calls an edge. In linear light a highlight rolloff swamps the
|
||||
// boundary between two midtones, and the watershed would put its strongest
|
||||
// walls where nobody sees one.
|
||||
//
|
||||
// The two chroma axes are opponent differences rather than a real Lab
|
||||
// transform: they cost three subtractions instead of a matrix and a cube
|
||||
// root, and the watershed only needs the *magnitude* of colour change, not a
|
||||
// colorimetrically defensible value for it.
|
||||
fn perceptual(c_in: vec3<f32>) -> vec3<f32> {
|
||||
var c = max(c_in, vec3<f32>(0.0));
|
||||
if (u.non_linear == 0u) {
|
||||
c = pow(c, vec3<f32>(1.0 / 2.4));
|
||||
}
|
||||
let l = dot(c, vec3<f32>(0.2126, 0.7152, 0.0722));
|
||||
let a = c.r - c.g;
|
||||
let b = c.b - 0.5 * (c.r + c.g);
|
||||
return vec3<f32>(l, a, b);
|
||||
}
|
||||
|
||||
// Source -> proxy, averaging every source pixel that falls in the proxy
|
||||
// pixel's footprint.
|
||||
//
|
||||
// A box average rather than point sampling because the proxy is where the
|
||||
// segmentation happens: point sampling a 24 MP sensor down to 2 MP aliases
|
||||
// fine texture into false gradient, and the watershed would faithfully find
|
||||
// basins in the aliasing.
|
||||
@compute @workgroup_size(8, 8, 1)
|
||||
fn features(@builtin(global_invocation_id) gid: vec3<u32>) {
|
||||
if (gid.x >= u.width || gid.y >= u.height) {
|
||||
return;
|
||||
}
|
||||
|
||||
let sx0 = (gid.x * u.src_width) / u.width;
|
||||
let sy0 = (gid.y * u.src_height) / u.height;
|
||||
let sx1 = max(sx0 + 1u, ((gid.x + 1u) * u.src_width) / u.width);
|
||||
let sy1 = max(sy0 + 1u, ((gid.y + 1u) * u.src_height) / u.height);
|
||||
|
||||
var acc = vec3<f32>(0.0);
|
||||
var n = 0.0;
|
||||
for (var sy = sy0; sy < sy1; sy = sy + 1u) {
|
||||
for (var sx = sx0; sx < sx1; sx = sx + 1u) {
|
||||
let c = textureLoad(src, vec2<i32>(i32(sx), i32(sy)), 0).rgb;
|
||||
acc = acc + perceptual(c);
|
||||
n = n + 1.0;
|
||||
}
|
||||
}
|
||||
|
||||
feat_out[gid.y * u.width + gid.x] = vec4<f32>(acc / max(n, 1.0), 0.0);
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------- blur
|
||||
|
||||
@group(0) @binding(3) var<storage, read> blur_in: array<vec4<f32>>;
|
||||
@group(0) @binding(4) var<storage, read_write> blur_out: array<vec4<f32>>;
|
||||
|
||||
fn clamp_coord(v: i32, hi: u32) -> u32 {
|
||||
return u32(clamp(v, 0, i32(hi) - 1));
|
||||
}
|
||||
|
||||
// Pre-smoothing. Not a refinement — without it the watershed is unusable.
|
||||
//
|
||||
// A raw gradient over sensor data has a local minimum at every noise grain,
|
||||
// and one basin per local minimum means a 2 MP frame segments into hundreds
|
||||
// of thousands of regions that correspond to nothing. The radius is the
|
||||
// caller's to set from ISO.
|
||||
@compute @workgroup_size(8, 8, 1)
|
||||
fn blur(@builtin(global_invocation_id) gid: vec3<u32>) {
|
||||
if (gid.x >= u.width || gid.y >= u.height) {
|
||||
return;
|
||||
}
|
||||
let idx = gid.y * u.width + gid.x;
|
||||
|
||||
if (u.blur_radius <= 0) {
|
||||
blur_out[idx] = blur_in[idx];
|
||||
return;
|
||||
}
|
||||
|
||||
var acc = vec3<f32>(0.0);
|
||||
var wsum = 0.0;
|
||||
let r = u.blur_radius;
|
||||
for (var dy = -r; dy <= r; dy = dy + 1) {
|
||||
for (var dx = -r; dx <= r; dx = dx + 1) {
|
||||
let sx = clamp_coord(i32(gid.x) + dx, u.width);
|
||||
let sy = clamp_coord(i32(gid.y) + dy, u.height);
|
||||
let d2 = f32(dx * dx + dy * dy);
|
||||
let w = exp(-d2 / (2.0 * f32(r) * f32(r)));
|
||||
acc = acc + blur_in[sy * u.width + sx].rgb * w;
|
||||
wsum = wsum + w;
|
||||
}
|
||||
}
|
||||
blur_out[idx] = vec4<f32>(acc / wsum, 0.0);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------- gradient
|
||||
|
||||
@group(0) @binding(5) var<storage, read> grad_in: array<vec4<f32>>;
|
||||
@group(0) @binding(6) var<storage, read_write> grad_out: array<f32>;
|
||||
|
||||
fn feat_at(x: i32, y: i32) -> vec3<f32> {
|
||||
let sx = clamp_coord(x, u.width);
|
||||
let sy = clamp_coord(y, u.height);
|
||||
return grad_in[sy * u.width + sx].rgb;
|
||||
}
|
||||
|
||||
// Sobel magnitude over the weighted opponent triple.
|
||||
//
|
||||
// This is the surface the watershed floods, so its units matter for nothing
|
||||
// except ordering — only the *relative* height of one boundary against
|
||||
// another decides which regions merge first.
|
||||
@compute @workgroup_size(8, 8, 1)
|
||||
fn gradient(@builtin(global_invocation_id) gid: vec3<u32>) {
|
||||
if (gid.x >= u.width || gid.y >= u.height) {
|
||||
return;
|
||||
}
|
||||
let x = i32(gid.x);
|
||||
let y = i32(gid.y);
|
||||
|
||||
let tl = feat_at(x - 1, y - 1);
|
||||
let tc = feat_at(x, y - 1);
|
||||
let tr = feat_at(x + 1, y - 1);
|
||||
let ml = feat_at(x - 1, y);
|
||||
let mr = feat_at(x + 1, y);
|
||||
let bl = feat_at(x - 1, y + 1);
|
||||
let bc = feat_at(x, y + 1);
|
||||
let br = feat_at(x + 1, y + 1);
|
||||
|
||||
let gx = (tr + 2.0 * mr + br) - (tl + 2.0 * ml + bl);
|
||||
let gy = (bl + 2.0 * bc + br) - (tl + 2.0 * tc + tr);
|
||||
|
||||
let w = vec3<f32>(u.w_luma, u.w_chroma, u.w_chroma);
|
||||
let wx = gx * w;
|
||||
let wy = gy * w;
|
||||
|
||||
grad_out[gid.y * u.width + gid.x] = sqrt(dot(wx, wx) + dot(wy, wy));
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------- flow
|
||||
|
||||
@group(0) @binding(7) var<storage, read> flow_grad: array<f32>;
|
||||
@group(0) @binding(8) var<storage, read_write> flow_out: array<u32>;
|
||||
|
||||
// Each pixel points at the steepest-descent neighbour among its 8, or at
|
||||
// itself if it is a local minimum — a basin seed.
|
||||
//
|
||||
// **The tie-break is load-bearing, twice over.** Comparing on (value, index)
|
||||
// rather than value alone gives a strict total order, so the pointer graph
|
||||
// descends monotonically and cannot contain a cycle — plateaux, which are
|
||||
// everywhere in a smoothed image, would otherwise make two equal pixels point
|
||||
// at each other and hang the pointer-jumping below.
|
||||
//
|
||||
// It is also what makes the result reproducible. S15's M5 asks whether a
|
||||
// label field is stable enough across GPU vendors to be a cache key
|
||||
// (ARCH §6.13); an arbitrary tie-break would answer no before the question
|
||||
// was asked.
|
||||
@compute @workgroup_size(8, 8, 1)
|
||||
fn flow(@builtin(global_invocation_id) gid: vec3<u32>) {
|
||||
if (gid.x >= u.width || gid.y >= u.height) {
|
||||
return;
|
||||
}
|
||||
let idx = gid.y * u.width + gid.x;
|
||||
|
||||
var best_val = flow_grad[idx];
|
||||
var best_idx = idx;
|
||||
|
||||
for (var dy = -1; dy <= 1; dy = dy + 1) {
|
||||
for (var dx = -1; dx <= 1; dx = dx + 1) {
|
||||
if (dx == 0 && dy == 0) {
|
||||
continue;
|
||||
}
|
||||
let nx = i32(gid.x) + dx;
|
||||
let ny = i32(gid.y) + dy;
|
||||
if (nx < 0 || ny < 0 || nx >= i32(u.width) || ny >= i32(u.height)) {
|
||||
continue;
|
||||
}
|
||||
let ni = u32(ny) * u.width + u32(nx);
|
||||
let nv = flow_grad[ni];
|
||||
if (nv < best_val || (nv == best_val && ni < best_idx)) {
|
||||
best_val = nv;
|
||||
best_idx = ni;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
flow_out[idx] = best_idx;
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------- jump
|
||||
|
||||
@group(0) @binding(9) var<storage, read> jump_in: array<u32>;
|
||||
@group(0) @binding(10) var<storage, read_write> jump_out: array<u32>;
|
||||
|
||||
// Pointer jumping: parent = parent[parent].
|
||||
//
|
||||
// Halves every path length per dispatch, so ceil(log2(longest path)) passes
|
||||
// resolve every pixel to its basin root. The host runs a fixed count bounded
|
||||
// by log2(pixel count) rather than testing for convergence, because a
|
||||
// convergence test costs a readback per iteration and the bound is ~21
|
||||
// dispatches of a trivial kernel.
|
||||
@compute @workgroup_size(8, 8, 1)
|
||||
fn jump(@builtin(global_invocation_id) gid: vec3<u32>) {
|
||||
if (gid.x >= u.width || gid.y >= u.height) {
|
||||
return;
|
||||
}
|
||||
let idx = gid.y * u.width + gid.x;
|
||||
jump_out[idx] = jump_in[jump_in[idx]];
|
||||
}
|
||||
Reference in New Issue
Block a user