dtourolle/DarkRoom
1
0
Fork 0
Code Issues 39 Pull Requests Actions Packages Projects Releases 20 Wiki Activity

Merge branch 'worktree-watershed-plateaux'

Browse Source
This commit is contained in:
dtourolle
2026-08-17 12:25:41 +02:00
parent b8e9793908 0b20436445
commit 76bb6b2847
3 changed files with 511 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/dr-gpu/src/shaders/watershed.wgsl
+186 -16
View File
@@ -1,12 +1,14 @@
// Watershed segmentation — the passes behind arm A of S15 (docs/segmentation.md).
//
// Five entry points forming one chain:
// Seven 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
// features source texture -> perceptual triple, downscaled to proxy size
// blur pre-smoothing, without which every grain becomes a basin
// gradient Sobel magnitude — the surface the watershed floods
// plateau_init seed the distance field at every real descent
// plateau_step carry it inward, so flat ground drains toward its exit
// 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
@@ -184,19 +186,162 @@ fn gradient(@builtin(global_invocation_id) gid: vec3<u32>) {
grad_out[gid.y * u.width + gid.x] = sqrt(dot(wx, wx) + dot(wy, wy));
}
// ---------------------------------------------------------- lower-complete
//
// A watershed needs every non-minimum pixel to have a lower neighbour. A real
// gradient does not oblige: a flat wall, a clipped sky or the inside of a
// uniform object is a **plateau**, where every neighbour is exactly equal and
// there is no downhill direction to follow.
//
// Left alone, the tie-break in `flow` sends every plateau pixel to its
// lowest-indexed neighbour, which is up and to the left. Each pixel therefore
// walks diagonally until it falls off the plateau, and one flat region becomes
// a fan of diagonal chains rather than one basin — visible as hatching across
// what should be a single area (docs/segmentation.md §12, F1).
//
// The fix is the standard lower-completion: give each plateau pixel its
// geodesic distance to the nearest pixel that *does* have a lower neighbour,
// then let `flow` order on (gradient, distance). Water on a plateau now runs
// toward the plateau's exit, which is what it would physically do.
//
// A plateau with no exit at all is a genuine regional minimum — the inside of
// a uniform disc, say. Those pixels keep `PLATEAU_UNRESOLVED`, tie with each
// other, and fall through to the index tie-break, which collapses the whole
// connected plateau onto its lowest-indexed pixel. One basin, which is the
// right answer for a regional minimum.
const PLATEAU_UNRESOLVED: u32 = 0xffffffffu;
// How close two gradients must be to count as the same level.
//
// **Exact equality does not work here, and that is not a rounding nicety.**
// The gradient is a float computed from 8-bit samples, so a region the eye
// and the algorithm both consider flat still has neighbours differing in the
// sixth decimal. With `==`, the breadth-first step never advances past its
// seeds and the whole pass is a no-op; with `<`, nearly every pixel finds
// some marginally lower neighbour and is seeded at zero, which is the same
// no-op wearing a different hat. Both were measured before this constant
// existed.
//
// Sized against the gradient's own scale: features are normalised to 0..1, so
// a Sobel magnitude runs to a few units, and 1e-4 is far below any step a
// real edge produces while sitting comfortably above f32 noise from a blur.
const LEVEL_EPS: f32 = 1e-4;
// Whether `b` lies below `a` by more than the level tolerance.
fn strictly_below(b: f32, a: f32) -> bool {
return b < a - LEVEL_EPS;
}
// Whether two gradients belong to the same plateau.
fn same_level(a: f32, b: f32) -> bool {
return abs(a - b) <= LEVEL_EPS;
}
@group(0) @binding(7) var<storage, read> pinit_grad: array<f32>;
@group(0) @binding(8) var<storage, read_write> pinit_out: array<u32>;
// Seed the distance field: zero where a real descent exists, unresolved on a
// plateau.
@compute @workgroup_size(8, 8, 1)
fn plateau_init(@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;
let here = pinit_grad[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;
}
if (strictly_below(pinit_grad[u32(ny) * u.width + u32(nx)], here)) {
pinit_out[idx] = 0u;
return;
}
}
}
pinit_out[idx] = PLATEAU_UNRESOLVED;
}
@group(0) @binding(9) var<storage, read> pstep_grad: array<f32>;
@group(0) @binding(10) var<storage, read> pstep_in: array<u32>;
@group(0) @binding(11) var<storage, read_write> pstep_out: array<u32>;
// One breadth-first step inward from the plateau's rim.
//
// Iterated by the host a fixed number of times rather than to convergence: a
// convergence test costs a readback per pass, and the count only has to cover
// the widest plateau in the frame. Pixels still unresolved when the budget
// runs out keep `PLATEAU_UNRESOLVED` and behave exactly as they did before
// this pass existed — the degradation is graceful, not a wrong answer.
@compute @workgroup_size(8, 8, 1)
fn plateau_step(@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;
let current = pstep_in[idx];
if (current != PLATEAU_UNRESOLVED) {
pstep_out[idx] = current;
return;
}
let here = pstep_grad[idx];
var best = PLATEAU_UNRESOLVED;
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);
// Only within the same plateau: a neighbour at a different height
// is across a boundary, and its distance says nothing about the
// way out of this one.
if (!same_level(pstep_grad[ni], here)) {
continue;
}
let nd = pstep_in[ni];
if (nd != PLATEAU_UNRESOLVED && nd < best) {
best = nd;
}
}
}
if (best == PLATEAU_UNRESOLVED) {
pstep_out[idx] = PLATEAU_UNRESOLVED;
} else {
pstep_out[idx] = best + 1u;
}
}
// -------------------------------------------------------------------- flow
@group(0) @binding(7) var<storage, read> flow_grad: array<f32>;
@group(0) @binding(8) var<storage, read_write> flow_out: array<u32>;
@group(0) @binding(12) var<storage, read> flow_grad: array<f32>;
@group(0) @binding(13) var<storage, read> flow_dist: array<u32>;
@group(0) @binding(14) 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.
// **The ordering is load-bearing, twice over.** Comparing on (gradient,
// plateau distance, index) rather than gradient 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
@@ -210,6 +355,7 @@ fn flow(@builtin(global_invocation_id) gid: vec3<u32>) {
let idx = gid.y * u.width + gid.x;
var best_val = flow_grad[idx];
var best_dist = flow_dist[idx];
var best_idx = idx;
for (var dy = -1; dy <= 1; dy = dy + 1) {
@@ -224,8 +370,32 @@ fn flow(@builtin(global_invocation_id) gid: vec3<u32>) {
}
let ni = u32(ny) * u.width + u32(nx);
let nv = flow_grad[ni];
if (nv < best_val || (nv == best_val && ni < best_idx)) {
let nd = flow_dist[ni];
// Lexicographic on (gradient, plateau distance, index) rather
// than one fused scalar. Folding the distance into the gradient
// as a small epsilon would need a scale factor that is small
// enough never to cross a real gradient step and large enough to
// survive f32 — a tuning problem with a silent failure mode,
// where three explicit keys have neither.
// The same tolerance the plateau passes use, and for the same
// reason: with exact equality this tie never fires on real data,
// so the distance carried inward above would be computed and then
// never consulted — the pass measurably did nothing.
var better = false;
if (strictly_below(nv, best_val)) {
better = true;
} else if (same_level(nv, best_val)) {
if (nd < best_dist) {
better = true;
} else if (nd == best_dist && ni < best_idx) {
better = true;
}
}
if (better) {
best_val = nv;
best_dist = nd;
best_idx = ni;
}
}
@@ -236,8 +406,8 @@ fn flow(@builtin(global_invocation_id) gid: vec3<u32>) {
// -------------------------------------------------------------------- jump
@group(0) @binding(9) var<storage, read> jump_in: array<u32>;
@group(0) @binding(10) var<storage, read_write> jump_out: array<u32>;
@group(0) @binding(15) var<storage, read> jump_in: array<u32>;
@group(0) @binding(16) var<storage, read_write> jump_out: array<u32>;
// Pointer jumping: parent = parent[parent].
//