diff --git a/core/dr-segment/src/refine.rs b/core/dr-segment/src/refine.rs index 1009efe..3ff0478 100644 --- a/core/dr-segment/src/refine.rs +++ b/core/dr-segment/src/refine.rs @@ -51,54 +51,77 @@ //! alone would leave at even odds. The comparative test is what stops the //! absolute one from needing a per-category constant. //! -//! # Why the verdict is kept, and the decision is not +//! # 5. And then the pixels are asked where the edge is //! -//! Everything above is per-image and costs a distance transform, a k-means and -//! a pass over the pixels. What comes out of it is one number per pixel — how -//! much better the category explains this colour than the alternatives — and -//! [`Refinement`] stores exactly that, quantised. +//! Everything above judges a pixel on its colour alone, which decides *what* +//! and not *where*. A colour test cannot put a boundary on an edge: it has no +//! notion of one. //! -//! The *decision* is then a smoothstep on that number, which is arithmetic. So -//! [`Refinement::apply`] is cheap enough to run under a dragging slider, and -//! the model is never consulted again. +//! So the last step is a **marker-based watershed**. The mask is eroded to +//! give two markers — confidently inside, confidently outside — and the flood +//! runs in the ribbon left between them, meeting along the most expensive line +//! it can find. The cost is a sum of terms, as docs/segmentation.md §2 says it +//! should be: the photograph's own edges, and the colour model's disagreement. //! -//! This is [`crate::distance`]'s arrangement, deliberately: there, a signed -//! distance field is computed once and feather, grow and shrink become -//! arithmetic on it, "which is what makes those live controls rather than ones -//! that stall on every drag". Same shape, different field. +//! Markers are what make this the right shape rather than the watershed §15 +//! discarded. That path failed because the *merge ladder* collapsed — 45,808 +//! basins reduced to one region plus specks. There is no ladder here. Seeding +//! prevents over-segmentation instead of merging afterwards, so the one +//! component that broke is the one component this does not have. And the +//! markers are better than the usual ones: the textbook derives them by +//! thresholding the gradient, guessing where objects are, where these come +//! from a model that knows what sky is. +//! +//! # The bound, which is what makes growth safe +//! +//! [`RefineOptions::travel_cells`] fixes how wide the ribbon may be, and +//! everything beyond it is already a marker — so the flood never reaches it. +//! **The boundary cannot move further than the model's own uncertainty.** +//! +//! That single bound replaces a pile of rules. Growing and shrinking are one +//! operation rather than two mechanisms with two safeguards; there is no +//! connectivity test, no reachability radius, no separate additive path. A +//! blue car below the horizon cannot be gained not because a rule forbids it +//! but because it is outside the ribbon and the flood is never there. +//! +//! It also costs almost nothing. A ribbon of a few tens of pixels around one +//! contour is a small part of a proxy, and a seeded flood never visits a pixel +//! that already has a label — so running only where the answer is unresolved +//! is not an optimisation added on top, it is what the algorithm does. //! //! # What the strictness control means //! -//! [`Refinement::apply`] takes a strictness in nats, and it is a real quantity -//! rather than an arbitrary 0–100: it is how much more plausible than the -//! alternatives a colour must be before its weight is kept. +//! Strictness is in nats: how much more plausible than the alternatives a +//! colour must be before it is kept. It moves two things at once, and they +//! agree. //! -//! Because a mode's rarity is *added* to the evidence against it, the things a -//! photographer wants to remove come off in the order they want them to. A -//! flag holding two percent of the sky is charged several nats; a bank of -//! cloud holding a third of it is charged none. So the strictness that takes -//! out the flag leaves the cloud, and the cloud only starts to go some way -//! further up — which is the property that makes a single slider worth -//! offering. It is pinned by `the_flag_goes_before_the_cloud_does`. +//! It sets **the colour threshold**, and because a mode's rarity is added to +//! the evidence against it, things come off in the order a photographer wants. +//! A flag holding two percent of the sky is charged several nats; a bank of +//! cloud holding a third of it is charged none. The strictness that takes the +//! flag leaves the cloud, pinned by `the_flag_goes_before_the_cloud_does`. //! -//! At zero the refinement does nothing at all, which is what makes it safe as -//! a default-on feature: the control has an off position that is exactly the -//! old behaviour. +//! It also sets **the ribbon's width**, which is what makes the control +//! continuous. At zero the ribbon is empty, every pixel is a marker, the flood +//! has nothing to decide, and the answer is the unrefined mask *exactly* — +//! not nearly. Turning it up widens the band the pixels may redraw. //! -//! # Two properties that make it safe to apply blind +//! # What this gave up, deliberately //! -//! **It is subtractive.** The output is the input times a factor in `0..=1`, -//! never more. So the worst failure available to it is losing part of a real -//! sky; it can never *gain* a region, and a blue car below the horizon that -//! was never in the mask cannot be pulled into it. +//! An earlier form of this was strictly subtractive — the output was the input +//! times a factor in `0..=1` — which bounded its damage and kept `scene.rs`'s +//! partition true for free. //! -//! **The partition survives.** `scene.rs` rests entirely on the categories -//! summing to at most one at every pixel — that is what lets two adjacent -//! grades be feathered without painting both into the overlap. Multiplying by -//! a factor in `0..=1` cannot raise a sum, so refining every category -//! independently still leaves a partition. The weight taken off the flag lands -//! in the unlisted remainder, which is exactly where a flag belongs: ADE20K -//! has no class for one. +//! That is gone, and it had to be. A mask that may only shrink can sharpen a +//! horizon inward but never outward, so wherever the coarse contour sat inside +//! the true edge, the error survived every setting of the control. Inside the +//! ribbon the model has admitted it does not know, and refusing to move +//! outward there is not caution, it is a preference for one direction of being +//! wrong. +//! +//! What replaces it is the travel bound: not "can only remove" but "can only +//! move this far". Weaker, and still a real guarantee — where the old one was +//! also, in the outward direction, a guarantee of staying wrong. //! //! # And when it cannot tell //! @@ -108,6 +131,9 @@ //! a sky that silently vanished from a mask is the kind of failure nobody //! attributes to the right place. +use std::cmp::Reverse; +use std::collections::BinaryHeap; + use crate::distance::signed_distance; /// How the evidence for a refinement is gathered. @@ -181,6 +207,50 @@ pub struct RefineOptions { /// Samples below which a side is too small to fit anything to. pub min_samples: usize, + + /// How far the boundary may travel at full strictness, in logit cells. + /// + /// The bound on everything the flood can do, and the reason it is safe to + /// let a boundary *grow* as well as shrink. Beyond this distance every + /// pixel is already a marker, so the flood never reaches it — the mask + /// cannot run away, whatever the colour model believes. + /// + /// One and a half cells because that is the model's own uncertainty: one + /// cell is its resolution and `rasterise`'s bilinear spreads it across one + /// more. Inside that band it has no information and the pixels should + /// decide; outside it, it does, and they should not. + pub travel_cells: f32, + + /// Thinnest structure, in pixels, that still keeps a marker. + /// + /// Eroding by a cell and a half deletes anything thinner than three cells + /// — a flagpole, a mast, a bare branch, and equally a strip of sky between + /// two of them. The flood would then have nothing seeded there and would + /// fill it from whichever side surrounds it, which is exactly how a + /// flagpole comes back as sky. + /// + /// So erosion stops at the ridge of the distance transform: whatever would + /// otherwise vanish keeps a one-pixel seed down its centre. This is the + /// floor on how thin a thing may be and still get one — below it a + /// structure is treated as noise, because a hot pixel and a distant bird + /// have ridges too. + /// + /// Note where this puts the signal-versus-noise decision. It used to live + /// in colour space, as a share of a fitted distribution, which nobody can + /// picture. Here it is a width in pixels, which a photographer can both + /// reason about and see. + pub min_thickness: f32, + + /// How much the photograph's own edges count against the colour model in + /// the flood's cost, `0.0..=1.0`. + /// + /// docs/segmentation.md §2 specifies the cost as *a sum of terms* — image + /// gradient always available, semantic evidence added when a model is + /// present — and this is the mix. At one the boundary lands purely on the + /// strongest edge in the band; at zero purely where the colour verdict + /// changes sign. Neither extreme is right: an edge with no colour meaning + /// is a texture, and a colour change with no edge is a gradient. + pub edge_weight: f32, } impl Default for RefineOptions { @@ -194,6 +264,9 @@ impl Default for RefineOptions { soften: 1.5, smooth_px: 2.0, min_samples: 500, + travel_cells: 1.5, + min_thickness: 3.0, + edge_weight: 0.6, } } } @@ -264,6 +337,27 @@ const FEATURES: usize = 3; /// anything a photograph resolves and comfortably above zero. const VARIANCE_FLOOR: f32 = 1e-4; +/// How much wider the broad mode is than the sample it was fitted on. +/// +/// A correction for a bias that is built into the method rather than a fudge. +/// The seeds are *eroded* by a cell and a half before anything is fitted, so +/// the sample is drawn from the middle of the category and never from its +/// edge — and for anything with a gradient across it, the colours nearest the +/// boundary are exactly the ones left out. The fitted spread is therefore +/// narrower than the category's true one, every time, in a known direction. +/// +/// Four is two standard deviations of slack, which is about what the erosion +/// removes from a sky: measured on the frame +/// `the_flood_moves_the_boundary_onto_the_real_edge` builds, sky fifteen rows +/// past the sample scored a Mahalanobis 5.9 against the uncorrected mode and +/// 1.5 against this one — the difference between refusing the horizon and +/// reaching it. +/// +/// Only the broad mode is widened. The tight modes are what discriminate, and +/// loosening those would let an intruder back in; this one's job is coverage, +/// and it is the one whose sample is unrepresentative in a way that matters. +const SHOULDER: f32 = 4.0; + /// Cap on the samples fitted per side. /// /// k-means over a million pixels answers the same as k-means over twenty @@ -301,11 +395,31 @@ pub struct Refinement { /// pixel this would be a proxy-sized `f32` buffer per category, and the /// quantisation is far finer than the decision it feeds. verdict: Vec, + /// Edge strength from the photograph, `0..=255`. + /// + /// The topography the flood runs over. Kept per category rather than once + /// per image, which duplicates it across the handful a frame reports — + /// worth it to keep [`Refinement`] a self-contained thing a caller can + /// hold without also holding the frame it came from. + gradient: Vec, width: usize, height: usize, /// Carried from the fit so that a caller holding only this can apply it /// without also having to keep the options that produced it. soften: f32, + cell_pixels: f32, + travel_cells: f32, + min_thickness: f32, + edge_weight: f32, +} + +/// Which side of the boundary a pixel is on, while the flood is deciding. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Side { + /// Not yet decided — the ribbon the flood runs in. + Open, + In, + Out, } impl Refinement { @@ -380,11 +494,23 @@ impl Refinement { verdict = blur(&verdict, width, height, options.smooth_px.round() as usize); } + // The contour the ribbon surrounds is deliberately *not* cached here. + // It has to be measured from the mask as the colour gate leaves it, + // and the gate moves with the control — so a field computed now would + // describe the boundary before the flag was cut out, and would put the + // ribbon nowhere near the hole the flood most needs to redraw. + // `decide` recomputes it, which is one transform per change of the + // control rather than one per image. Ok(Self { verdict: verdict.iter().map(|&v| quantise(v)).collect(), + gradient: sobel(rgb, width, height, options.luma_weight), width, height, soften: options.soften, + cell_pixels: cell_pixels.max(1.0), + travel_cells: options.travel_cells, + min_thickness: options.min_thickness, + edge_weight: options.edge_weight.clamp(0.0, 1.0), }) } @@ -394,9 +520,9 @@ impl Refinement { /// The gate this refinement applies at `strictness`, per pixel. /// - /// Separate from [`Self::apply`] because the overlay wants to draw it and - /// the coverage path wants to multiply by it, and neither should have to - /// reimplement the smoothstep. + /// The colour model's opinion on its own, before the flood has weighed it + /// against the photograph's edges. The overlay draws this, and + /// [`Self::resolve`] uses it as one of the two terms in the flood's cost. pub fn gate(&self, strictness: f32) -> Vec { let strictness = strictness.max(0.0); self.verdict @@ -405,40 +531,322 @@ impl Refinement { .collect() } - /// Apply at a given strictness. The cheap half — a slider drags on this. + /// How far the boundary may travel at this strictness, in pixels. + /// + /// Scaled by the control rather than fixed, and that is what makes the + /// control continuous. At zero the ribbon is empty, every pixel is a + /// marker, and the flood has nothing to decide — which is *exactly* the + /// unrefined mask, not merely close to it. Turning the control up widens + /// the band the pixels are allowed to redraw. + fn ribbon(&self, strictness: f32) -> f32 { + let t = (strictness / STRICTNESS_MAX).clamp(0.0, 1.0); + t * self.travel_cells * self.cell_pixels + } + + /// Decide the boundary: markers from the mask, the flood between them. + /// + /// This is the marker-based watershed. The mask is eroded by the ribbon to + /// give two markers — confidently inside, confidently outside — and the + /// only pixels left over are the ones within the ribbon of the contour. + /// The flood runs *there and nowhere else*, which is not an optimisation + /// bolted on but simply what a seeded flood does: it never visits a pixel + /// that already has a label. + /// + /// Two properties fall out of that, and both are the reason this shape was + /// chosen over growing a mask outward by a rule: + /// + /// - **The boundary cannot travel further than the ribbon.** Everything + /// beyond is already a marker. Growing and shrinking become one + /// operation with one bound, and the bound is the model's own + /// uncertainty. + /// - **The cost is a fraction of the frame.** A ribbon of a few tens of + /// pixels around one contour is a small part of a proxy, so this is + /// cheap enough to sit under a control. + pub fn resolve(&self, weights: &[f32], strictness: f32) -> Vec { + self.decide(weights, strictness).2 + } + + /// The gate, the contours it leaves, and where the flood puts them. + /// + /// The two halves in the order they have to run, and the order is the + /// whole point. A colour gate removes an intruder *wherever it is*, + /// including deep in the interior where no boundary passes within a + /// ribbon of it — a flag forty pixels from the horizon is exactly that. + /// A flood cannot: it only refines contours that already exist, and there + /// is no contour around the flag precisely because the model never noticed + /// one. + /// + /// So the gate goes first and *creates* the contour, and the flood then + /// puts every contour — the horizon and the new hole alike — onto an edge + /// the photograph actually has. Neither does the other's job. + fn decide(&self, weights: &[f32], strictness: f32) -> (Vec, Vec, Vec) { + let gate = self.gate(strictness); + let gated: Vec = weights.iter().zip(&gate).map(|(&w, &g)| w * g).collect(); + + // The distance is measured from the *gated* mask, not the model's, so + // the ribbon surrounds the boundary as it stands after the colour test + // — including the one that has just appeared around the flag. + let coverage: Vec = gated + .iter() + .map(|&w| (w.clamp(0.0, 1.0) * 255.0).round() as u8) + .collect(); + let distance = signed_distance(&coverage, self.width, self.height, 128); + + let mut side = self.markers(&distance, self.ribbon(strictness)); + self.flood(&mut side, &gate); + ( + gated, + distance, + side.iter().map(|&s| s == Side::In).collect(), + ) + } + + /// Erode to the ribbon, then put back what erosion would have destroyed. + /// + /// The second half is what stops a flagpole reappearing. A structure + /// thinner than twice the ribbon has no pixel far enough from the contour + /// to become a marker, so it would be left open and the flood would fill + /// it from whichever side surrounds it — a mast in the sky becomes sky, a + /// strip of sky between branches becomes branch. Both directions, and the + /// mirror case is as real as the first. + /// + /// So a pixel also becomes a marker when it sits on the *ridge* of the + /// distance transform: no neighbour is further from the contour than it + /// is. That is the medial axis, and it is precisely what survives erosion + /// to nothing — every structure keeps a seed down its centre however thin + /// it is, provided it clears [`RefineOptions::min_thickness`]. + fn markers(&self, distance: &[f32], ribbon: f32) -> Vec { + let floor = 0.5 * self.min_thickness; + let (w, h) = (self.width, self.height); + let mut side = vec![Side::Open; distance.len()]; + + for y in 0..h { + for x in 0..w { + let p = y * w + x; + let d = distance[p]; + let magnitude = d.abs(); + + // Far enough from the contour to be trusted outright. + let decided = magnitude >= ribbon + // Or on the ridge of something too thin to survive the + // erosion, and thick enough not to be a speck. + || (magnitude >= floor && self.on_ridge(distance, x, y, magnitude)); + if decided { + side[p] = if d >= 0.0 { Side::In } else { Side::Out }; + } + } + } + side + } + + /// Whether no neighbour is further from the contour than this pixel. + /// + /// The ridge test, on the *magnitude* of the distance so that it finds the + /// centre line of a structure on either side of the boundary. Ties count + /// as ridge — a plateau down the middle of an even-width bar has no strict + /// maximum, and refusing it there would leave exactly the structures this + /// exists for unseeded. + fn on_ridge(&self, distance: &[f32], x: usize, y: usize, magnitude: f32) -> bool { + let (w, h) = (self.width, self.height); + for (dx, dy) in NEIGHBOURS { + let (nx, ny) = (x as isize + dx, y as isize + dy); + if nx < 0 || ny < 0 || nx >= w as isize || ny >= h as isize { + continue; + } + let n = distance[ny as usize * w + nx as usize].abs(); + if n > magnitude { + return false; + } + } + true + } + + /// Flood the open pixels from the markers, cheapest first. + /// + /// A priority flood: every open pixel adjacent to a marker enters a queue + /// keyed on the cost of crossing it, and the cheapest is always taken + /// next. A pixel is claimed by whichever side reaches it first, so the two + /// fronts meet along the most expensive line in the ribbon — which is the + /// watershed, and which is where the boundary belongs. + /// + /// The cost is a sum of terms, as docs/segmentation.md §2 says it should + /// be: the photograph's own edges, and the colour model's disagreement. + /// Neither alone is right. An edge with no colour meaning is a texture, + /// and a colour change with no edge is a gradient. + fn flood(&self, side: &mut [Side], gate: &[f32]) { + let (w, h) = (self.width, self.height); + let mut queue: BinaryHeap> = BinaryHeap::new(); + + let push = |queue: &mut BinaryHeap>, p: usize, from: Side| { + queue.push(Reverse((self.cost(p, gate), p, from as u8))); + }; + + for p in 0..side.len() { + if side[p] == Side::Open { + continue; + } + for n in neighbours(p, w, h) { + if side[n] == Side::Open { + push(&mut queue, n, side[p]); + } + } + } + + while let Some(Reverse((_, p, from))) = queue.pop() { + if side[p] != Side::Open { + continue; + } + side[p] = if from == Side::In as u8 { + Side::In + } else { + Side::Out + }; + for n in neighbours(p, w, h) { + if side[n] == Side::Open { + push(&mut queue, n, side[p]); + } + } + } + } + + /// What it costs the flood to cross one pixel. + /// + /// Quantised to an integer, and that is deliberate rather than sloppy: the + /// queue is ordered on it, floating-point ties would order by whatever the + /// comparison happened to do, and a mask that differs between machines + /// reaches the sidecar meaning one thing on the desktop and another on the + /// phone. Integer cost with the pixel index as the tiebreak makes the + /// whole flood a total order. + fn cost(&self, p: usize, gate: &[f32]) -> u32 { + let edge = self.gradient[p] as f32 / 255.0; + // Peaks where the colour model is undecided, so the boundary is drawn + // toward its own crossing as well as toward the photograph's edges. + let ambiguity = 1.0 - (2.0 * gate[p] - 1.0).abs(); + let cost = self.edge_weight * edge + (1.0 - self.edge_weight) * ambiguity; + (cost.clamp(0.0, 1.0) * u16::MAX as f32) as u32 + } + + /// Apply at a given strictness. /// /// At [`STRICTNESS_OFF`] this is the identity, exactly, and returns the - /// weights unchanged without touching the verdict at all. That is what - /// makes the control's zero position mean "as the model weighted it" - /// rather than "very nearly that". + /// weights unchanged without touching anything. That is what makes the + /// control's zero position mean "as the model weighted it" rather than + /// "very nearly that". + /// + /// Above zero the colour gate removes what does not belong, and the flood + /// then redraws every resulting contour within the ribbon. Outside the + /// ribbon the gated weight stands; inside it, the flood's answer does — + /// including where the flood puts the edge further out than the model did. pub fn apply(&self, weights: &[f32], strictness: f32) -> Vec { if strictness <= STRICTNESS_OFF || weights.len() != self.verdict.len() { return weights.to_vec(); } - weights - .iter() - .zip(self.gate(strictness)) - .map(|(&w, g)| w * g) + let ribbon = self.ribbon(strictness); + let (gated, distance, inside) = self.decide(weights, strictness); + (0..gated.len()) + .map(|p| { + if distance[p].abs() >= ribbon { + gated[p] + } else if inside[p] { + 1.0 + } else { + 0.0 + } + }) .collect() } /// [`Self::apply`] over a quantised coverage buffer. /// /// The form the UI actually holds a category in: the mask reaches the - /// rasteriser as one byte per pixel, and round-tripping it through floats - /// to reuse `apply` would allocate twice as much for no extra precision. + /// rasteriser as one byte per pixel. It goes through the float path rather + /// than repeating it, because the flood has to run over the same numbers + /// either way and a second implementation would be a second thing to keep + /// in step. pub fn apply_coverage(&self, coverage: &[u8], strictness: f32) -> Vec { if strictness <= STRICTNESS_OFF || coverage.len() != self.verdict.len() { return coverage.to_vec(); } - coverage + let weights: Vec = coverage.iter().map(|&c| c as f32 / 255.0).collect(); + self.apply(&weights, strictness) .iter() - .zip(self.gate(strictness)) - .map(|(&c, g)| (c as f32 * g).round().clamp(0.0, 255.0) as u8) + .map(|&w| (w * 255.0).round().clamp(0.0, 255.0) as u8) .collect() } } +/// The eight-neighbourhood, as offsets. +/// +/// Eight rather than four because a watershed on a four-neighbourhood produces +/// staircased diagonals, and a mask edge is looked at by someone at 100% zoom. +const NEIGHBOURS: [(isize, isize); 8] = [ + (-1, -1), + (0, -1), + (1, -1), + (-1, 0), + (1, 0), + (-1, 1), + (0, 1), + (1, 1), +]; + +fn neighbours(p: usize, w: usize, h: usize) -> impl Iterator { + let (x, y) = ((p % w) as isize, (p / w) as isize); + NEIGHBOURS.into_iter().filter_map(move |(dx, dy)| { + let (nx, ny) = (x + dx, y + dy); + (nx >= 0 && ny >= 0 && nx < w as isize && ny < h as isize) + .then(|| ny as usize * w + nx as usize) + }) +} + +/// Edge strength over the opponent features, as one byte per pixel. +/// +/// Sobel over the same three numbers the colour model is fitted on, rather +/// than over plain luma — docs/segmentation.md §3 is explicit that a +/// channel-weighted RGB gradient reads a saturated red edge as weaker than it +/// looks, and a flag against sky is exactly that edge. +/// +/// The magnitude is normalised against the frame's own strongest edge, so the +/// cost mix in [`RefineOptions::edge_weight`] means the same thing on a flat +/// photograph as on a contrasty one. +fn sobel(rgb: &[f32], width: usize, height: usize, luma_weight: f32) -> Vec { + let f: Vec<[f32; FEATURES]> = (0..width * height) + .map(|p| feature(rgb, p, luma_weight)) + .collect(); + let at = |x: usize, y: usize| f[y * width + x]; + + let mut raw = vec![0.0f32; width * height]; + let mut peak = 0.0f32; + for y in 0..height { + for x in 0..width { + // Clamped at the border rather than skipped: a zero edge around + // the frame would be the cheapest path in the image, and the flood + // would run along it in preference to any real boundary. + let (x0, x1) = (x.saturating_sub(1), (x + 1).min(width - 1)); + let (y0, y1) = (y.saturating_sub(1), (y + 1).min(height - 1)); + + let mut sum = 0.0f32; + for c in 0..FEATURES { + let gx = at(x1, y0)[c] + 2.0 * at(x1, y)[c] + at(x1, y1)[c] + - at(x0, y0)[c] + - 2.0 * at(x0, y)[c] + - at(x0, y1)[c]; + let gy = at(x0, y1)[c] + 2.0 * at(x, y1)[c] + at(x1, y1)[c] + - at(x0, y0)[c] + - 2.0 * at(x, y0)[c] + - at(x1, y0)[c]; + sum += gx * gx + gy * gy; + } + let g = sum.sqrt(); + raw[y * width + x] = g; + peak = peak.max(g); + } + } + + let norm = if peak > 0.0 { 255.0 / peak } else { 0.0 }; + raw.iter().map(|&g| (g * norm).round() as u8).collect() +} + /// TRACES: FR-DEV-3 /// Fit and apply in one call, for a caller with no control to offer. /// @@ -524,11 +932,22 @@ fn sample( #[derive(Debug, Clone, Copy)] struct Mode { mean: [f32; FEATURES], - /// Per-feature, not a full covariance. The opponent axes are close enough - /// to decorrelated by construction that the off-diagonal terms buy little, - /// and a diagonal fit stays well-conditioned on the few hundred samples a - /// small mode gets. - variance: [f32; FEATURES], + /// The **full** inverse covariance, not a per-feature one. + /// + /// A diagonal fit was tried first and is wrong here, for the reason this + /// whole module exists: a category with a gradient across it moves along + /// all three features *together*. Sky is the case — luminance rises, the + /// red-green axis drifts and the blue-yellow axis falls, in lockstep, from + /// zenith to horizon. + /// + /// Treating those as independent charges a colour two standard deviations + /// along that gradient three times over, once per axis, and a chi-square + /// on three degrees of freedom then rejects it. Measured on the frame + /// `the_flood_moves_the_boundary_onto_the_real_edge` builds: sky five rows + /// past the sampled range scored 11.6 against a threshold of 11.34, so the + /// model refused the very thing it was refining. With the correlation kept + /// the same colour scores about 4. + precision: [[f32; FEATURES]; FEATURES], /// `0.5 * ln|Σ|`, precomputed because it is constant per mode and the /// inner loop runs once per pixel. half_log_det: f32, @@ -602,13 +1021,15 @@ fn fit(samples: &[[f32; FEATURES]], k: usize) -> Vec { // Mean, spread and share from the final assignment. let mut sums = vec![[0.0f32; FEATURES]; k]; - let mut squares = vec![[0.0f32; FEATURES]; k]; + let mut cross = vec![[[0.0f32; FEATURES]; FEATURES]; k]; let mut counts = vec![0usize; k]; for (s, sample) in samples.iter().enumerate() { let c = assignment[s]; - for f in 0..FEATURES { - sums[c][f] += sample[f]; - squares[c][f] += sample[f] * sample[f]; + for i in 0..FEATURES { + sums[c][i] += sample[i]; + for j in 0..FEATURES { + cross[c][i][j] += sample[i] * sample[j]; + } } counts[c] += 1; } @@ -621,24 +1042,135 @@ fn fit(samples: &[[f32; FEATURES]], k: usize) -> Vec { continue; } let n = counts[c] as f32; - let mut mean = [0.0f32; FEATURES]; - let mut variance = [0.0f32; FEATURES]; - let mut half_log_det = 0.0f32; - for f in 0..FEATURES { - mean[f] = sums[c][f] / n; - variance[f] = (squares[c][f] / n - mean[f] * mean[f]).max(VARIANCE_FLOOR); - half_log_det += 0.5 * variance[f].ln(); + let rarity = -((n / total) * even).ln().min(0.0); + if let Some(mode) = gaussian(&sums[c], &cross[c], n, rarity, 1.0) { + modes.push(mode); } - modes.push(Mode { - mean, - variance, - half_log_det, - rarity: -((n / total) * even).ln().min(0.0), - }); } + + // And one mode over the whole side, on top of the k tight ones. + // + // Without it this method rejects the very thing it is refining, whenever a + // category has a smooth gradient across it. Sky is the case: the seeds are + // *eroded*, so the sample never contains the colours nearest the boundary, + // and k-means then splits the narrow range it did see into tight bands. A + // pixel a little beyond that range — which is exactly every pixel the + // boundary might move onto — sits well outside the nearest band and is + // refused. The erosion that makes the seeds trustworthy is the same + // erosion that makes them unrepresentative. + // + // A single mode over every sample carries the category's full spread + // *including the direction its gradient runs in*, so it covers that + // shoulder while the tight modes keep their discrimination: `nearest` + // takes whichever explains a colour best, so a tight mode wins wherever it + // applies and this one catches what falls off the end. + // + // It does not reopen the door to an intruder. Its covariance is wide along + // the gradient and narrow across it, so a flag — which differs in a + // direction the sky never travels — is as far outside this mode as it is + // outside the tight ones. + // + // Rarity zero: it explains the whole side by construction, so there is + // nothing to charge it for. + if !modes.is_empty() { + let mut sum = [0.0f32; FEATURES]; + let mut square = [[0.0f32; FEATURES]; FEATURES]; + for c in 0..k { + for i in 0..FEATURES { + sum[i] += sums[c][i]; + for j in 0..FEATURES { + square[i][j] += cross[c][i][j]; + } + } + } + if let Some(mode) = gaussian(&sum, &square, total, 0.0, SHOULDER) { + modes.push(mode); + } + } + modes } +/// One Gaussian from accumulated first and second moments. +/// +/// Regularised on the diagonal before inversion. A mode over a flat patch of +/// sky has a covariance of nearly nothing, and nothing in the denominator +/// makes every other colour infinitely improbable; a mode whose samples lie on +/// a line — which a gradient's samples very nearly do — is singular outright +/// and has no inverse at all. [`VARIANCE_FLOOR`] is a standard deviation of +/// one part in a hundred, comfortably below anything a photograph resolves and +/// comfortably above zero. +fn gaussian( + sum: &[f32; FEATURES], + cross: &[[f32; FEATURES]; FEATURES], + n: f32, + rarity: f32, + spread: f32, +) -> Option { + let mut mean = [0.0f32; FEATURES]; + for i in 0..FEATURES { + mean[i] = sum[i] / n; + } + + let mut covariance = [[0.0f32; FEATURES]; FEATURES]; + for i in 0..FEATURES { + for j in 0..FEATURES { + covariance[i][j] = (cross[i][j] / n - mean[i] * mean[j]) * spread; + } + covariance[i][i] += VARIANCE_FLOOR; + } + + let (precision, determinant) = invert(covariance)?; + Some(Mode { + mean, + precision, + half_log_det: 0.5 * determinant.ln(), + rarity, + }) +} + +/// Invert a 3x3 symmetric matrix, returning it with its determinant. +/// +/// Written out rather than pulled in: three features is the whole vocabulary +/// here, cofactors of a 3x3 are six lines, and a linear-algebra dependency +/// under the Android NDK is exactly what D13 says to avoid for this. +fn invert(m: [[f32; FEATURES]; FEATURES]) -> Option<([[f32; FEATURES]; FEATURES], f32)> { + let (a, b, c) = (m[0][0], m[0][1], m[0][2]); + let (d, e, f) = (m[1][0], m[1][1], m[1][2]); + let (g, h, i) = (m[2][0], m[2][1], m[2][2]); + + let determinant = a * (e * i - f * h) - b * (d * i - f * g) + c * (d * h - e * g); + // Regularisation above should make this unreachable; a mode that reaches + // it anyway is dropped rather than trusted, because a precision matrix + // built from a near-zero determinant is numerically meaningless and would + // score every colour as wildly improbable. + if !determinant.is_finite() || determinant <= 1e-20 { + return None; + } + + let scale = 1.0 / determinant; + Some(( + [ + [ + (e * i - f * h) * scale, + (c * h - b * i) * scale, + (b * f - c * e) * scale, + ], + [ + (f * g - d * i) * scale, + (a * i - c * g) * scale, + (c * d - a * f) * scale, + ], + [ + (d * h - e * g) * scale, + (b * g - a * h) * scale, + (a * e - b * d) * scale, + ], + ], + determinant, + )) +} + /// k-means++ seeding, with a fixed sequence. /// /// The usual seeding is random, and random would make a mask depend on which @@ -691,10 +1223,15 @@ fn seed_centres(samples: &[[f32; FEATURES]], k: usize) -> Vec<[f32; FEATURES]> { fn nearest(modes: &[Mode], x: &[f32; FEATURES]) -> (f32, f32) { let mut best = (f32::INFINITY, f32::INFINITY); for mode in modes { + let mut delta = [0.0f32; FEATURES]; + for f in 0..FEATURES { + delta[f] = x[f] - mode.mean[f]; + } let mut maha = 0.0f32; - for ((v, mean), variance) in x.iter().zip(&mode.mean).zip(&mode.variance) { - let d = v - mean; - maha += d * d / variance; + for i in 0..FEATURES { + for j in 0..FEATURES { + maha += delta[i] * mode.precision[i][j] * delta[j]; + } } let nll = 0.5 * maha + mode.half_log_det + mode.rarity; if nll < best.1 { @@ -953,57 +1490,165 @@ mod tests { assert_eq!(r.apply(&weights, STRICTNESS_OFF), weights); } - /// Raising the strictness may only ever remove more. + /// The one guarantee, and the one that replaced "subtractive". /// - /// A slider that gave weight back on the way up would be one whose - /// direction the user cannot predict, and it would break the ordering the - /// whole control rests on. + /// The flood may not move a boundary further than the ribbon. Everything + /// beyond it is already a marker, so no pixel out there can be reached at + /// any strictness — which is what bounds the damage now that an edge is + /// allowed to travel outward as well as in. + /// + /// Checked against the *gated* mask, because that is the boundary the + /// flood actually works on, and across the whole travel of the control, + /// because a bound that only held at the bottom would be no bound at all. #[test] - fn strictness_is_monotonic() { + fn the_flood_cannot_travel_further_than_the_ribbon() { + let (rgb, weights) = with(FLAG, RED); + let opts = RefineOptions::default(); + let r = Refinement::compute(&weights, &rgb, EDGE, EDGE, CELL, &opts) + .expect("both sides present"); + + for step in 1..=8 { + let strictness = step as f32 * STRICTNESS_MAX / 8.0; + let ribbon = r.ribbon(strictness); + let (gated, distance, inside) = r.decide(&weights, strictness); + let out = r.apply(&weights, strictness); + + for p in 0..weights.len() { + if distance[p].abs() < ribbon { + continue; + } + assert!( + (out[p] - gated[p]).abs() < 1e-6, + "pixel {p} at {} from the gated contour moved at {strictness}: {} -> {}", + distance[p], + gated[p], + out[p] + ); + assert_eq!( + inside[p], + distance[p] >= 0.0, + "pixel {p} beyond the ribbon was labelled against its own side" + ); + } + } + } + + /// The colour gate, on its own, may only ever remove. + /// + /// The flood is what gained the right to add, and only inside the ribbon. + /// The gate keeps the older and stronger property, and it is worth pinning + /// separately: it is what removes an intruder nowhere near a boundary, + /// where the flood cannot reach. + #[test] + fn the_colour_gate_only_removes() { let (rgb, weights) = with(FLAG, RED); let r = Refinement::compute(&weights, &rgb, EDGE, EDGE, CELL, &RefineOptions::default()) .expect("both sides present"); - let mut previous = r.apply(&weights, 0.0); - for step in 1..=12 { - let next = r.apply(&weights, step as f32 * STRICTNESS_MAX / 12.0); - for (p, (&before, &after)) in previous.iter().zip(&next).enumerate() { + for step in 0..=8 { + let strictness = step as f32 * STRICTNESS_MAX / 8.0; + let gated = r.decide(&weights, strictness).0; + for (p, (&before, &after)) in weights.iter().zip(&gated).enumerate() { assert!( after <= before + 1e-6, - "pixel {p} gained weight at step {step}: {before} -> {after}" + "the gate gave pixel {p} weight at {strictness}: {before} -> {after}" ); } - previous = next; } } - /// The refinement may only ever take weight away. + /// A structure thinner than the erosion keeps a marker. /// - /// This is what bounds its damage and what keeps `scene.rs`'s partition - /// true when several categories are refined independently, so it is - /// asserted rather than left as a property of the arithmetic. + /// The flagpole case, and the reason [`RefineOptions::min_thickness`] + /// exists. A mast a few pixels wide has no pixel far enough from the + /// contour to survive an erosion of a cell and a half, so without the + /// ridge rule it would be left open and the flood would fill it from the + /// sky surrounding it. The pole would come back — and come back + /// *confident*, which is worse than the coarse mask that started it. #[test] - fn refinement_is_subtractive() { - let (rgb, weights) = with(FLAG, RED); - let (out, _) = refine_category( - &weights, - &rgb, - EDGE, - EDGE, - CELL, - STRICTNESS_DEFAULT, - &RefineOptions::default(), - ); - for (p, (&before, &after)) in weights.iter().zip(&out).enumerate() { - assert!( - after <= before + 1e-6, - "pixel {p} gained weight: {before} -> {after}" - ); - assert!( - (0.0..=1.0).contains(&after), - "pixel {p} out of range: {after}" - ); + fn a_structure_thinner_than_the_erosion_keeps_its_marker() { + let (w, h) = (EDGE, EDGE); + let (mut rgb, mut weights) = landscape(w, h); + + // A five-pixel mast standing in the sky, correctly excluded by the + // model — so it is a thin sliver of *exterior* inside a large + // interior, which is exactly what erosion destroys. + for y in 20..80 { + for x in 94..99 { + let p = y * w + x; + rgb[p * 3..p * 3 + 3].copy_from_slice(&[0.15, 0.14, 0.16]); + weights[p] = 0.0; + } } + + let r = Refinement::compute(&weights, &rgb, w, h, CELL, &RefineOptions::default()) + .expect("both sides present"); + let distance = r.decide(&weights, STRICTNESS_DEFAULT).1; + let markers = r.markers(&distance, r.ribbon(STRICTNESS_DEFAULT)); + + let seeded = (20..80) + .flat_map(|y| (94..99).map(move |x| y * w + x)) + .any(|p| markers[p] == Side::Out); + assert!( + seeded, + "the mast must keep at least one marker down its centre" + ); + + // And it must still be excluded once the flood has run. + let out = r.apply(&weights, STRICTNESS_DEFAULT); + let core = mean(&out, w, (95, 40, 98, 60)); + assert!(core < 0.3, "the mast should not have been flooded: {core}"); + } + + /// The boundary lands on the photograph's edge, not on the model's guess. + /// + /// The whole point of the flood, and the one thing the colour gate alone + /// could never do: the mask's contour is put eight pixels *above* the real + /// horizon, so the correction is outward. A gate that may only remove + /// weight would leave this wrong at every setting. + /// + /// The travel bound is widened for this test rather than the displacement + /// shrunk, because the two are coupled through one control: strictness + /// sets both the ribbon's width and the colour threshold, so buying a + /// wider ribbon from the default options also means demanding a colour + /// agreement strict enough to cut the sky itself. Moving `travel_cells` + /// separates them, and it is the quantity actually under test. + #[test] + fn the_flood_moves_the_boundary_onto_the_real_edge() { + let (w, h) = (EDGE, EDGE); + let (rgb, _) = landscape(w, h); + + // The picture's own horizon is at h/2. Claim only as far as h/2 - 8. + let mut weights = vec![0.0f32; w * h]; + for y in 0..(h / 2 - 8) { + for x in 0..w { + weights[y * w + x] = 1.0; + } + } + + let options = RefineOptions { + travel_cells: 3.0, + ..RefineOptions::default() + }; + let r = + Refinement::compute(&weights, &rgb, w, h, CELL, &options).expect("both sides present"); + // Half scale: a 12px ribbon, comfortably past the 8px error, at a + // colour threshold the sky itself still passes. + let out = r.apply(&weights, STRICTNESS_DEFAULT); + + // The band the model gave away: sky in the photograph, not in the mask. + let reclaimed = mean(&out, w, (20, h / 2 - 6, w - 20, h / 2 - 2)); + assert!( + reclaimed > 0.7, + "the flood should have taken back sky the model missed: {reclaimed}" + ); + + // And it must not have run past the real horizon into the ground. + let ground = mean(&out, w, (20, h / 2 + 4, w - 20, h / 2 + 10)); + assert!( + ground < 0.3, + "the flood should stop at the edge, not cross it: {ground}" + ); } /// The coverage path must agree with the float one, since the application @@ -1027,6 +1672,40 @@ mod tests { } } + /// The flood must not depend on the order a heap happens to pop ties. + /// + /// Two pixels of equal cost are ordinary — a flat region has thousands — + /// and if the winner were decided by whatever the comparison did, the mask + /// would differ between runs. It reaches the sidecar as identity, so that + /// would mean one thing on the desktop and another on the phone. + #[test] + fn the_flood_is_deterministic() { + // A deliberately flat frame: no gradient anywhere, so *every* cost in + // the ribbon ties and the tiebreak is the only thing ordering it. + let (w, h) = (96usize, 96usize); + let mut rgb = vec![0.5f32; w * h * 3]; + for y in h / 2..h { + for x in 0..w { + let p = y * w + x; + rgb[p * 3 + 2] = 0.2; + } + } + let mut weights = vec![0.0f32; w * h]; + for y in 0..h / 2 { + for x in 0..w { + weights[y * w + x] = 1.0; + } + } + + let opts = RefineOptions::default(); + let r = Refinement::compute(&weights, &rgb, w, h, CELL, &opts).expect("both sides present"); + assert_eq!( + r.resolve(&weights, STRICTNESS_DEFAULT), + r.resolve(&weights, STRICTNESS_DEFAULT), + "the same refinement must flood the same way twice" + ); + } + /// No confident interior means no distribution, and no distribution must /// mean "leave it alone" rather than "reject everything". #[test] diff --git a/ui/dr-ui/ui/masks.slint b/ui/dr-ui/ui/masks.slint index 7f775d5..758b7cd 100644 --- a/ui/dr-ui/ui/masks.slint +++ b/ui/dr-ui/ui/masks.slint @@ -249,9 +249,10 @@ component MaskEntry inherits Rectangle { // sample of a drag. if root.data.refinable: SliderRow { label: "Refine"; - hint: "Drop pixels whose colour does not match the rest of " - + "the category — a flag in the sky, a chimney, a bare " - + "branch. Zero is the model's own outline."; + hint: "Drop pixels whose colour does not belong — a flag in " + + "the sky, a chimney, a bare branch — and pull the " + + "outline onto the edge the photograph actually has. " + + "Zero is the model's own outline."; value: root.data.refine; default-value: 4; minimum: 0;