Files
DarkRoom/core/dr-segment/src/refine.rs
T
dtourolle 6b1aac477d Put the developer docs under docs/dev and index the folder for users first
docs/ had 26 developer documents flat beside the manual, and the two
audiences are very differently sized: most readers want the manual and
the gesture reference, a few want the register, the designs and the
measurements. The manual and gestures.md stay at the top; everything for
someone changing the code moves to docs/dev/, and the two documents that
name their own successors — the v0.1 milestone and the UI-refinement plan
— go to docs/dev/archive/ rather than being deleted, since both are still
cited. docs/README.md is the index, users first.

Every reference follows: code comments, Cargo manifests, the workflows,
the pre-commit hook, the bench and traceability tools (which locate the
repo root by docs/dev/requirements.md now), packaging, the Docker READMEs,
CLAUDE.md, CONTRIBUTING.md and the README. The matrix links one level
deeper and is regenerated. Links out of the moved documents into the tree
gain a level; a link checker over every Markdown file finds none broken.
2026-09-20 16:20:15 +02:00

1960 lines
81 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//! Sharpening a coarse category mask against the photograph's own colours.
//!
//! # The problem, as a number
//!
//! The `scene` module keeps the model's native `[1, 150, 80, 80]` logit grid,
//! so one cell is eight of the graph's input pixels across. At the 1600px
//! proxy the application segments at, the letterbox scale is `0.4` and **one
//! cell is 20 proxy pixels** — which the bilinear in `Scene::rasterise` then
//! spreads across one more either side.
//!
//! So a flag in the sky, or a chimney, or a bare branch, sits inside a handful
//! of cells whose softmax is dominated by the sky around it, and comes out
//! weighted as sky. No feather setting recovers it, because the information
//! was never in the grid.
//!
//! # What this does about it
//!
//! The photograph is at full proxy resolution even though the weights are not,
//! and it knows exactly where the flag is. The move is to let the *model*
//! decide what the category is and the *pixels* decide which of them belong to
//! it — the same division of labour [`crate::prior`] already draws between the
//! instance model and the watershed, applied to categories.
//!
//! 1. **Seeds.** Threshold the weights high, then erode by
//! [`RefineOptions::margin_cells`] logit cells using
//! [`crate::distance::signed_distance`]. What survives is confidently
//! inside; the mirror of it is confidently outside. The erosion radius is
//! derived from the grid rather than picked, because one cell *is* the
//! model's resolution and anything within a cell of the boundary is
//! precisely what there is no evidence about.
//!
//! 2. **A distribution per side, with several modes.** A single Gaussian over
//! sky is wrong: sky is blue at the zenith, white where the cloud is, and
//! pale at the horizon, and one blob over all three rejects two of them. So
//! k-means, which is GrabCut's mixture without the EM.
//!
//! 3. **A mode is charged for its rarity.** A *small* flag deep in the sky has
//! both a high weight and a large distance from the mask's boundary, so it
//! lands in the interior sample and fits itself a mode. It cannot be
//! excluded geometrically. What separates it from a real appearance of the
//! category is that it explains almost none of the category — so a colour
//! whose only explanation is a thin mode is charged `−ln(share)` nats for
//! it, and a flag ends up far less plausible than the sky around it without
//! ever being ruled out by fiat.
//!
//! 4. **Two tests, and a pixel must pass both.** An absolute one — is this
//! colour plausible under the category at all, as a chi-square on the
//! Mahalanobis distance — and a comparative one, is it likelier under the
//! interior than the exterior. The absolute test is what catches the flag,
//! whose colour is far from *both* sides and which the comparative test
//! alone would leave at even odds. The comparative test is what stops the
//! absolute one from needing a per-category constant.
//!
//! # 5. And then the pixels are asked where the edge is
//!
//! 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.
//!
//! 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/dev/segmentation.md §2 says it
//! should be: the photograph's own edges, and the colour model's disagreement.
//!
//! 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
//!
//! 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.
//!
//! 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`.
//!
//! 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.
//!
//! # What this gave up, deliberately
//!
//! 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.
//!
//! 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
//!
//! Every path that lacks the evidence to judge refuses to build a
//! [`Refinement`] at all and says why, rather than returning a plausible one.
//! An empty seed set fitted to a distribution would reject *every* pixel, and
//! 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.
///
/// Everything here feeds the *fit*, which happens once per image. The one
/// control that moves afterwards is the strictness passed to
/// [`Refinement::apply`], and it is deliberately not in this struct: mixing
/// the two would invite a caller to change a fit parameter under a slider and
/// wonder why the drag stalled.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct RefineOptions {
/// Weight at or above which a pixel may seed the interior distribution.
///
/// High on purpose. This is not the mask's threshold — it is the standard
/// for "the model was sure here", and everything downstream is a
/// consequence of what these pixels look like.
pub seed_weight: f32,
/// How far inside the boundary a seed must lie, in logit cells.
///
/// One cell is the model's own resolution and `rasterise`'s bilinear
/// spreads a cell's influence across one more, so the first cell and a
/// half either side of the boundary is smear rather than evidence.
pub margin_cells: f32,
/// Modes fitted per side.
///
/// Rarity is measured against an even split ([`Mode::rarity`]), so this
/// number does not shift the verdicts on its own — raising it buys a finer
/// description of the category without making every colour look rarer.
pub clusters: usize,
/// How much luminance counts against chrominance in judging a colour.
///
/// Well under one, and that is the difference between this working and not
/// on sky. Sky's variance is *dominated* by luminance — the zenith is
/// several stops off the horizon and a cloud is brighter than either — so
/// at equal weight the distribution is a long bright streak that a
/// mid-grey flag sits comfortably inside. A flag is separated by
/// chrominance; a cloud is separated by luminance alone. Down-weighting
/// keeps the cloud and catches the flag.
///
/// Not zero, or a dark bird against a bright sky is kept.
///
/// The same reasoning `dr_gpu::SegmentOptions` applies to the watershed's
/// gradient, for the same reason.
pub luma_weight: f32,
/// Squared Mahalanobis distance at which a colour stops being plausible
/// under the category.
///
/// Chi-square with three degrees of freedom: `11.34` is the 99th
/// percentile, so under the fitted model one confident pixel in a hundred
/// is expected to fail this on its own. It sets where the verdict scale's
/// zero falls; the strictness then moves the decision along that scale.
pub outlier: f32,
/// Width of the band, in nats, over which the gate falls from keep to cut.
///
/// A hard cut would put an aliased edge through the photograph at exactly
/// the place a person is looking.
pub soften: f32,
/// Radius, in pixels, the *verdict* is blurred by.
///
/// A per-pixel colour test on a real sky speckles: sensor noise puts
/// individual pixels over the line in both directions. Blurring the
/// evidence rather than the decision means it is paid for once, in
/// [`Refinement::compute`], instead of on every frame of a drag.
pub smooth_px: f32,
/// 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/dev/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 {
fn default() -> Self {
Self {
seed_weight: 0.9,
margin_cells: 1.5,
clusters: 4,
luma_weight: 0.25,
outlier: 11.34,
soften: 1.5,
smooth_px: 2.0,
min_samples: 500,
travel_cells: 1.5,
min_thickness: 3.0,
edge_weight: 0.6,
}
}
}
/// Strictness at which a category is left exactly as the model weighted it.
pub const STRICTNESS_OFF: f32 = 0.0;
/// The top of the useful strictness range, in nats.
///
/// Past this a category's own dominant colours have gone, so a control
/// offering more would only offer ways to delete the mask. Fixed here rather
/// than in the UI so that the slider and the tests agree on what the end of
/// the scale means.
pub const STRICTNESS_MAX: f32 = 8.0;
/// Where the control sits until someone moves it.
///
/// Half scale, and that position is the measurement rather than a round
/// number. On the synthetic frame `the_flag_goes_before_the_cloud_does`
/// builds, a flag holding 1.6% of the sky is more than half gone by **2.95**
/// and a cloud bank holding a third of it survives to **5.75** — so the whole
/// of the interval between them is a place the control can sit, and this is
/// near its middle.
///
/// That separation is what rarity buys. Without it both land at the same
/// verdict, there is no interval, and the control would appear dead until it
/// suddenly ate the sky.
///
/// Treat the numbers as a calibration and not a law: the synthetic sky is
/// smooth enough to sit on [`VARIANCE_FLOOR`], where a real one has noise and
/// therefore a real spread, which moves every crossing down together. The
/// ordering survives that; the exact placement is what the slider is for.
///
/// **Not where a new layer starts.** That warning turned out to be an
/// understatement — on a real photograph this position removes most of every
/// category that is not sky, and the measurements are in
/// [`Refinement::gentle`], which is what a layer starts at instead. What this
/// constant still names is the midpoint of the control's travel, and the
/// synthetic frame the tests hold it against.
pub const STRICTNESS_DEFAULT: f32 = 4.0;
/// Why a refinement could not be built.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SkipReason {
/// Nothing survived the erosion — the category is present but everywhere
/// thinner than the model's own resolution, so there is no pixel it is
/// sure about.
NoInterior,
/// The frame is essentially all this category, leaving nothing to
/// contrast it against.
NoExterior,
/// The buffers handed in do not describe one image.
Mismatched,
}
/// What a one-shot [`refine_category`] did.
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Refined {
/// Applied. `removed` is the fraction of the category's original total
/// weight the gate took away.
Applied {
removed: f32,
},
Skipped(SkipReason),
}
/// Number of features a pixel is judged on: one luminance, two chrominance.
const FEATURES: usize = 3;
/// Floor on a fitted variance.
///
/// A mode over a flat patch of sky has a variance of nearly nothing, and
/// nothing in the denominator makes every other pixel infinitely improbable.
/// This is a standard deviation of one part in a hundred, comfortably below
/// 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
/// thousand of them, and costs fifty times as much. The subsample is a fixed
/// stride rather than a random draw so the result is reproducible.
const MAX_SAMPLES: usize = 20_000;
/// Lloyd iterations.
///
/// Fixed rather than run to convergence: a stopping rule that depends on
/// floating-point comparison is a stopping rule that can differ between
/// machines, and a mask that differs between machines reaches the sidecar as
/// indices meaning one thing on the desktop and another on the phone
/// (docs/dev/segmentation.md §6).
const ITERATIONS: usize = 12;
/// Half-width of the verdict scale, in nats.
///
/// The verdict is stored as a byte, so it needs an end. Sixteen is comfortably
/// past [`STRICTNESS_MAX`] plus a soften band, and it puts the quantisation
/// step at an eighth of a nat — an order finer than the narrowest transition
/// the gate can be asked for.
const VERDICT_RANGE: f32 = 16.0;
/// The evidence for refining one category, ready for a control to act on.
///
/// Built once per image by [`Refinement::compute`]; read by
/// [`Refinement::apply`] as often as a slider moves.
#[derive(Debug, Clone, PartialEq)]
pub struct Refinement {
/// Per pixel, how much better the category explains this colour than the
/// alternatives, in nats, quantised over ±[`VERDICT_RANGE`].
///
/// A byte for the same reason a coverage buffer is one: at four bytes a
/// pixel this would be a proxy-sized `f32` buffer per category, and the
/// quantisation is far finer than the decision it feeds.
verdict: Vec<u8>,
/// 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<u8>,
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 {
/// Fit the colour models and score every pixel. The expensive half.
///
/// `weights` is one category's coverage at `width * height`, as
/// `Scene::rasterise` produces it. `rgb` is the same picture, tightly
/// packed `f32` RGB — the proxy the model itself read, so the two describe
/// one frame.
///
/// `cell_pixels` is how many pixels of *this* buffer one logit cell spans;
/// see `Scene::cell_pixels`, which is where the caller should get it
/// rather than re-deriving a letterbox inverse.
pub fn compute(
weights: &[f32],
rgb: &[f32],
width: usize,
height: usize,
cell_pixels: f32,
options: &RefineOptions,
) -> Result<Self, SkipReason> {
let pixels = width * height;
if weights.len() != pixels || rgb.len() != pixels * 3 || pixels == 0 {
return Err(SkipReason::Mismatched);
}
// The distance field is over the *thresholded* weights, so the
// boundary it measures from is where the model stopped being sure —
// not where the mask will eventually be cut, which is what the
// strictness decides later.
let coverage: Vec<u8> = weights
.iter()
.map(|&w| (w.clamp(0.0, 1.0) * 255.0).round() as u8)
.collect();
let threshold = (options.seed_weight.clamp(0.0, 1.0) * 255.0).round() as u8;
let distance = signed_distance(&coverage, width, height, threshold);
let margin = options.margin_cells * cell_pixels.max(1.0);
let interior = sample(&distance, rgb, options.luma_weight, |d| d >= margin);
if interior.len() < options.min_samples {
return Err(SkipReason::NoInterior);
}
let exterior = sample(&distance, rgb, options.luma_weight, |d| d <= -margin);
if exterior.len() < options.min_samples {
return Err(SkipReason::NoExterior);
}
let inside = fit(&interior, options.clusters);
let outside = fit(&exterior, options.clusters);
let mut verdict = vec![0.0f32; pixels];
for (p, cell) in verdict.iter_mut().enumerate() {
let x = feature(rgb, p, options.luma_weight);
// Squared Mahalanobis for the absolute test, because that is what
// is chi-square distributed; the density — the same distance with
// the mode's own spread folded in — for the comparative one, so
// that a tight mode and a loose one are compared fairly. Both
// carry the mode's rarity.
let (maha_in, nll_in) = nearest(&inside, &x);
let (_, nll_out) = nearest(&outside, &x);
let plausible = 0.5 * (options.outlier - maha_in);
let likelier = nll_out - nll_in;
*cell = plausible.min(likelier);
}
// The evidence is blurred, not the decision — see
// `RefineOptions::smooth_px`.
if options.smooth_px >= 1.0 {
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),
})
}
pub fn size(&self) -> (usize, usize) {
(self.width, self.height)
}
/// The gate this refinement applies at `strictness`, per pixel.
///
/// 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<f32> {
let strictness = strictness.max(0.0);
self.verdict
.iter()
.map(|&v| smoothstep(-self.soften, self.soften, dequantise(v) - strictness))
.collect()
}
/// 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<bool> {
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<f32>, Vec<f32>, Vec<bool>) {
let gate = self.gate(strictness);
let gated: Vec<f32> = 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<u8> = 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<Side> {
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/dev/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<Reverse<(u32, usize, u8)>> = BinaryHeap::new();
let push = |queue: &mut BinaryHeap<Reverse<(u32, usize, u8)>>, 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 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<f32> {
if strictness <= STRICTNESS_OFF || weights.len() != self.verdict.len() {
return weights.to_vec();
}
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. 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<u8> {
if strictness <= STRICTNESS_OFF || coverage.len() != self.verdict.len() {
return coverage.to_vec();
}
let weights: Vec<f32> = coverage.iter().map(|&c| c as f32 / 255.0).collect();
self.apply(&weights, strictness)
.iter()
.map(|&w| (w * 255.0).round().clamp(0.0, 255.0) as u8)
.collect()
}
/// The strongest strictness this photograph can be started at without the
/// category disappearing.
///
/// # Why a constant could not do this job
///
/// [`STRICTNESS_DEFAULT`] was measured on the synthetic frame the tests
/// build, and its own note warns that a real photograph's noise "moves
/// every crossing down together". It moves them a great deal further than
/// that reads. Measured on seven ordinary frames, `4.0` — half scale, the
/// position a control would naturally start at — removes:
///
/// | category | weight removed at 4.0 |
/// |----------|----------------------|
/// | sky | 0% – 6% |
/// | vegetation | 18% – 91% |
/// | ground | 36% – 93% |
/// | architecture | 76% – **99.5%** |
///
/// So a layer created at the constant is an *empty mask* on most
/// photographs that are not mostly sky, and empty is indistinguishable
/// from broken: the adjustment moves and no pixel changes. That is the
/// whole of the fault, and it cannot be fixed by choosing a smaller
/// constant — the useful position is 5.0 on one frame and below 1.0 on the
/// next, because a nat of evidence means different things over a smooth
/// sky and over a stone facade.
///
/// # What this does instead
///
/// It asks the photograph. Walking down from half scale, the first rung
/// whose gate takes no more than [`GENTLE_CUT`] of the category's weight
/// is the answer, and [`STRICTNESS_OFF`] is the answer when none of them
/// does — the model's own outline, which is never wrong about *where the
/// category is*, only about where it stops.
///
/// Downwards rather than upwards because the friendly case is the common
/// one and it exits on the first rung: a sky costs one `apply`, and only a
/// frame the refinement disagrees with pays for all four.
///
/// This is a starting position and not a limit. The slider still offers
/// the whole range, and it is the control's job to let a photographer go
/// past what this considered safe.
pub fn gentle(&self, coverage: &[u8]) -> f32 {
let total: f64 = coverage.iter().map(|&c| c as f64).sum();
if total <= 0.0 {
return STRICTNESS_OFF;
}
for &strictness in GENTLE_LADDER {
let kept: f64 = self
.apply_coverage(coverage, strictness)
.iter()
.map(|&c| c as f64)
.sum();
if (total - kept) / total <= GENTLE_CUT as f64 {
return strictness;
}
}
STRICTNESS_OFF
}
}
/// How much of a category's weight a *starting* strictness may take.
///
/// A sixth, and the number is doing one job: separating "the gate tidied the
/// edge" from "the gate ate the category". A twenty-proxy-pixel boundary
/// around a subject covering a fifth of the frame is a few percent of its
/// weight, so a refinement doing what it is for lands well under this; the
/// failures measured on [`Refinement::gentle`]'s table are all at 76% and
/// above. Nothing sits near the line, which is what makes it safe to state as
/// a constant rather than fit.
const GENTLE_CUT: f32 = 1.0 / 6.0;
/// The rungs [`Refinement::gentle`] tries, strongest first.
///
/// Whole nats, because that is the unit the evidence is denominated in and the
/// spacing the sweep in `examples/scene.rs` is read at. Stopping at half scale
/// rather than at [`STRICTNESS_MAX`]: past there a category's own dominant
/// colours have gone, and a photograph on which 8.0 removed under a sixth of
/// the weight would be one where the gate is finding nothing to cut and the
/// starting position may as well be gentler.
const GENTLE_LADDER: &[f32] = &[4.0, 3.0, 2.0, 1.0];
/// 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<Item = usize> {
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/dev/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<u8> {
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.
///
/// The example and the tests use this. The application does not: it keeps the
/// [`Refinement`] so that its slider costs a smoothstep rather than a k-means.
pub fn refine_category(
weights: &[f32],
rgb: &[f32],
width: usize,
height: usize,
cell_pixels: f32,
strictness: f32,
options: &RefineOptions,
) -> (Vec<f32>, Refined) {
match Refinement::compute(weights, rgb, width, height, cell_pixels, options) {
Ok(refinement) => {
let refined = refinement.apply(weights, strictness);
let before: f32 = weights.iter().sum();
let after: f32 = refined.iter().sum();
let removed = if before > 0.0 {
((before - after) / before).clamp(0.0, 1.0)
} else {
0.0
};
(refined, Refined::Applied { removed })
}
Err(why) => (weights.to_vec(), Refined::Skipped(why)),
}
}
fn quantise(v: f32) -> u8 {
let t = (v.clamp(-VERDICT_RANGE, VERDICT_RANGE) + VERDICT_RANGE) / (2.0 * VERDICT_RANGE);
(t * 255.0).round() as u8
}
fn dequantise(b: u8) -> f32 {
(b as f32 / 255.0) * 2.0 * VERDICT_RANGE - VERDICT_RANGE
}
/// One pixel's colour, as the three numbers the distributions are fitted over.
///
/// Opponent axes rather than raw RGB, and no division anywhere: normalising
/// chrominance by luminance would be exposure-invariant and would also blow up
/// in the shadows, which on a landscape is the whole foreground.
fn feature(rgb: &[f32], p: usize, luma_weight: f32) -> [f32; FEATURES] {
let (r, g, b) = (rgb[p * 3], rgb[p * 3 + 1], rgb[p * 3 + 2]);
let luma = 0.2126 * r + 0.7152 * g + 0.0722 * b;
[luma_weight * luma, r - g, b - 0.5 * (r + g)]
}
/// Every pixel whose distance passes `want`, subsampled to [`MAX_SAMPLES`].
///
/// Two passes — count, then take a fixed stride — rather than one pass
/// collecting everything and thinning afterwards, so a sky filling a 22 MP
/// proxy never allocates twenty million features to throw most of them away.
fn sample(
distance: &[f32],
rgb: &[f32],
luma_weight: f32,
want: impl Fn(f32) -> bool,
) -> Vec<[f32; FEATURES]> {
let total = distance.iter().filter(|&&d| want(d)).count();
if total == 0 {
return Vec::new();
}
let stride = total.div_ceil(MAX_SAMPLES).max(1);
let mut out = Vec::with_capacity(total.div_ceil(stride));
let mut seen = 0usize;
for (p, &d) in distance.iter().enumerate() {
if !want(d) {
continue;
}
if seen.is_multiple_of(stride) {
out.push(feature(rgb, p, luma_weight));
}
seen += 1;
}
out
}
/// One mode of a side's colour distribution.
#[derive(Debug, Clone, Copy)]
struct Mode {
mean: [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,
/// What it costs, in nats, to explain a colour only by this mode.
///
/// `−ln(share × k)`, floored at zero. Three things follow from that shape:
///
/// - It is measured against an **even split**, so raising
/// [`RefineOptions::clusters`] describes the category more finely
/// without making every colour in it look rarer.
/// - It is floored, so a dominant mode earns no *bonus*. A discount there
/// would let the category's commonest colour outvote a genuinely bad
/// chi-square, which is the one direction this must not bend.
/// - It is what replaces pruning small modes outright. A flag holding two
/// percent of the sky is charged a few nats rather than deleted, so
/// where it falls relative to the cloud is a *distance* the strictness
/// control can travel across, instead of a threshold it jumps over.
rarity: f32,
}
/// Fit `k` modes to one side.
fn fit(samples: &[[f32; FEATURES]], k: usize) -> Vec<Mode> {
let k = k.max(1).min(samples.len());
let mut centres = seed_centres(samples, k);
let k = centres.len();
let mut assignment = vec![0usize; samples.len()];
for _ in 0..ITERATIONS {
let mut moved = false;
for (s, sample) in samples.iter().enumerate() {
let mut best = 0usize;
let mut best_d = f32::INFINITY;
for (c, centre) in centres.iter().enumerate() {
let d = squared(sample, centre);
if d < best_d {
best_d = d;
best = c;
}
}
if assignment[s] != best {
assignment[s] = best;
moved = true;
}
}
if !moved {
break;
}
let mut sums = vec![[0.0f32; 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];
}
counts[c] += 1;
}
for c in 0..k {
// An emptied centre is left where it was rather than re-seeded.
// Re-seeding is the usual advice, but it needs a random draw
// inside the loop and the determinism is worth more here than the
// empty cluster costs.
if counts[c] == 0 {
continue;
}
for f in 0..FEATURES {
centres[c][f] = sums[c][f] / counts[c] as f32;
}
}
}
// Mean, spread and share from the final assignment.
let mut sums = 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 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;
}
let total = samples.len() as f32;
let even = k as f32;
let mut modes = Vec::new();
for c in 0..k {
if counts[c] == 0 {
continue;
}
let n = counts[c] as f32;
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);
}
}
// 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<Mode> {
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
/// run produced it. The draw is a fixed-seed LCG instead: still spread out,
/// still far better than taking the first `k` samples, and the same every time
/// on every machine.
fn seed_centres(samples: &[[f32; FEATURES]], k: usize) -> Vec<[f32; FEATURES]> {
let mut rng = Lcg(0x2545_F491_4F6C_DD1D);
let mut centres = vec![samples[0]];
let mut nearest = vec![f32::INFINITY; samples.len()];
while centres.len() < k {
let last = *centres.last().expect("seeded with one");
let mut total = 0.0f64;
for (s, sample) in samples.iter().enumerate() {
nearest[s] = nearest[s].min(squared(sample, &last));
total += nearest[s] as f64;
}
if total <= 0.0 {
// Every sample coincides with a centre — a perfectly flat side.
// More modes would all be the same mode.
break;
}
let mut target = rng.unit() as f64 * total;
let mut pick = samples.len() - 1;
for (s, &d) in nearest.iter().enumerate() {
target -= d as f64;
if target <= 0.0 {
pick = s;
break;
}
}
centres.push(samples[pick]);
}
centres
}
/// The closest mode, as `(squared Mahalanobis, negative log density)`, both
/// charged for that mode's rarity.
///
/// The rarity enters the Mahalanobis term doubled because that term is a
/// squared distance and the other is a log density: a chi-square of `2r` is
/// the same amount of evidence as `r` nats, so the two tests are then denominated
/// in the same currency and the strictness means one thing across both.
///
/// Nearest rather than a proper mixture sum: the largest term dominates a
/// well-separated mixture, and taking the max is what makes the two numbers
/// this returns describe *the same* mode.
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 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 {
best = (maha + 2.0 * mode.rarity, nll);
}
}
best
}
fn squared(a: &[f32; FEATURES], b: &[f32; FEATURES]) -> f32 {
(0..FEATURES).map(|f| (a[f] - b[f]).powi(2)).sum()
}
fn smoothstep(edge0: f32, edge1: f32, x: f32) -> f32 {
if edge1 <= edge0 {
return if x >= edge1 { 1.0 } else { 0.0 };
}
let t = ((x - edge0) / (edge1 - edge0)).clamp(0.0, 1.0);
t * t * (3.0 - 2.0 * t)
}
/// Separable box blur, clamped at the edges.
///
/// A box rather than a gaussian because what is wanted is the removal of
/// single-pixel speckle, and the shape of the kernel that does it does not
/// matter.
fn blur(src: &[f32], width: usize, height: usize, radius: usize) -> Vec<f32> {
if radius == 0 || width == 0 || height == 0 {
return src.to_vec();
}
let span = (2 * radius + 1) as f32;
let mut mid = vec![0.0f32; src.len()];
for y in 0..height {
for x in 0..width {
let mut sum = 0.0f32;
for k in 0..=2 * radius {
let sx = (x + k).saturating_sub(radius).min(width - 1);
sum += src[y * width + sx];
}
mid[y * width + x] = sum / span;
}
}
let mut out = vec![0.0f32; src.len()];
for y in 0..height {
for x in 0..width {
let mut sum = 0.0f32;
for k in 0..=2 * radius {
let sy = (y + k).saturating_sub(radius).min(height - 1);
sum += mid[sy * width + x];
}
out[y * width + x] = sum / span;
}
}
out
}
/// A fixed-sequence generator, for the one place a draw is needed.
struct Lcg(u64);
impl Lcg {
fn unit(&mut self) -> f32 {
self.0 = self
.0
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
// The high bits are the well-mixed ones in an LCG; the low bits cycle
// with a short period and would make the draw far from uniform.
((self.0 >> 40) as f32) / ((1u64 << 24) as f32)
}
}
#[cfg(test)]
mod tests {
use super::*;
/// One logit cell, in the pixels of the synthetic frames below.
const CELL: f32 = 8.0;
/// Edge of every synthetic frame.
///
/// Large enough that an intruder can be a realistic *fraction* of the
/// category rather than a realistic number of pixels — which is what
/// rarity is measured in, and getting that proportion wrong is the
/// difference between this method working and not.
const EDGE: usize = 192;
/// A blue upper half over a green lower half, and a category mask that
/// claims the whole upper half — which is what the coarse grid produces.
fn landscape(w: usize, h: usize) -> (Vec<f32>, Vec<f32>) {
let mut rgb = vec![0.0f32; w * h * 3];
let mut weights = vec![0.0f32; w * h];
for y in 0..h {
for x in 0..w {
let p = y * w + x;
// A vertical gradient in the sky, because a real one has one
// and a single Gaussian over luma is exactly what it breaks.
let t = y as f32 / h as f32;
if y < h / 2 {
rgb[p * 3] = 0.25 + 0.4 * t;
rgb[p * 3 + 1] = 0.45 + 0.35 * t;
rgb[p * 3 + 2] = 0.85;
} else {
rgb[p * 3] = 0.20;
rgb[p * 3 + 1] = 0.45;
rgb[p * 3 + 2] = 0.15;
}
weights[p] = if y < h / 2 { 1.0 } else { 0.0 };
}
}
(rgb, weights)
}
/// Paint a rectangle into the frame, and let the coarse mask claim it —
/// the flag in the sky.
fn intrude(
rgb: &mut [f32],
weights: &mut [f32],
w: usize,
rect: (usize, usize, usize, usize),
colour: [f32; 3],
) {
let (x0, y0, x1, y1) = rect;
for y in y0..y1 {
for x in x0..x1 {
let p = y * w + x;
rgb[p * 3..p * 3 + 3].copy_from_slice(&colour);
weights[p] = 1.0;
}
}
}
fn mean(v: &[f32], w: usize, rect: (usize, usize, usize, usize)) -> f32 {
let (x0, y0, x1, y1) = rect;
let mut sum = 0.0;
for y in y0..y1 {
for x in x0..x1 {
sum += v[y * w + x];
}
}
sum / ((x1 - x0) * (y1 - y0)) as f32
}
const FLAG: (usize, usize, usize, usize) = (88, 40, 104, 56);
const FLAG_CORE: (usize, usize, usize, usize) = (92, 44, 100, 52);
const CLOUD: (usize, usize, usize, usize) = (30, 12, 150, 60);
const CLOUD_CORE: (usize, usize, usize, usize) = (50, 24, 130, 48);
const RED: [f32; 3] = [0.75, 0.10, 0.12];
const WHITE: [f32; 3] = [0.92, 0.94, 0.96];
fn with(rect: (usize, usize, usize, usize), colour: [f32; 3]) -> (Vec<f32>, Vec<f32>) {
let (mut rgb, mut weights) = landscape(EDGE, EDGE);
intrude(&mut rgb, &mut weights, EDGE, rect, colour);
(rgb, weights)
}
/// The case the module exists for: a red flag inside the sky, claimed by
/// the coarse mask, must come back out at the default strictness — while
/// the sky around it stays.
#[test]
fn a_flag_in_the_sky_is_removed() {
let (rgb, weights) = with(FLAG, RED);
let (out, what) = refine_category(
&weights,
&rgb,
EDGE,
EDGE,
CELL,
STRICTNESS_DEFAULT,
&RefineOptions::default(),
);
assert!(
matches!(what, Refined::Applied { .. }),
"should have had the evidence to judge: {what:?}"
);
let inside = mean(&out, EDGE, FLAG_CORE);
assert!(inside < 0.2, "the flag should be cut out, got {inside}");
let sky = mean(&out, EDGE, (8, 8, 60, 60));
assert!(sky > 0.8, "the sky around it should survive, got {sky}");
}
/// A cloud is a legitimate part of the sky and is separated from the blue
/// by luminance alone, which is what `luma_weight` is for — and it holds a
/// large share, which is what rarity is for.
#[test]
fn a_cloud_is_not_mistaken_for_an_intruder() {
let (rgb, weights) = with(CLOUD, WHITE);
let (out, _) = refine_category(
&weights,
&rgb,
EDGE,
EDGE,
CELL,
STRICTNESS_DEFAULT,
&RefineOptions::default(),
);
let cloud = mean(&out, EDGE, CLOUD_CORE);
assert!(cloud > 0.7, "the cloud should stay sky, got {cloud}");
}
/// The property that makes a single slider worth offering.
///
/// Rarity puts the flag and the cloud at *different places on one scale*
/// rather than on two sides of a threshold, so there is a strictness that
/// has taken the flag and not the cloud — and the flag always goes first.
/// Without the rarity term both sit at the same verdict and the control
/// would appear dead until it suddenly ate the sky.
#[test]
fn the_flag_goes_before_the_cloud_does() {
let opts = RefineOptions::default();
let (flag_rgb, flag_w) = with(FLAG, RED);
let flag = Refinement::compute(&flag_w, &flag_rgb, EDGE, EDGE, CELL, &opts)
.expect("the flag frame has both sides");
let (cloud_rgb, cloud_w) = with(CLOUD, WHITE);
let cloud = Refinement::compute(&cloud_w, &cloud_rgb, EDGE, EDGE, CELL, &opts)
.expect("the cloud frame has both sides");
// The first strictness on a fine sweep at which each is more than half
// gone. `None` would mean it never goes at all.
let crossing = |r: &Refinement, weights: &[f32], core| {
(0..=120)
.map(|i| i as f32 * STRICTNESS_MAX / 120.0)
.find(|&s| mean(&r.apply(weights, s), EDGE, core) < 0.5)
};
let flag_at = crossing(&flag, &flag_w, FLAG_CORE).expect("the flag must go somewhere");
let cloud_at = crossing(&cloud, &cloud_w, CLOUD_CORE);
// `None` is the better outcome, not a missing case: it means the cloud
// survived the whole range, so the flag went first by a margin wider
// than the control can travel.
if let Some(cloud_at) = cloud_at {
assert!(
flag_at < cloud_at,
"the flag must go first: flag at {flag_at}, cloud at {cloud_at}"
);
}
assert!(
flag_at <= STRICTNESS_DEFAULT,
"the default must already clear a flag, but it only goes at {flag_at}"
);
}
/// The coverage buffer as the application holds it: one byte a pixel.
fn quantised(weights: &[f32]) -> Vec<u8> {
weights
.iter()
.map(|&w| (w.clamp(0.0, 1.0) * 255.0).round() as u8)
.collect()
}
/// What fraction of a category's weight a strictness takes away.
fn removed(r: &Refinement, coverage: &[u8], strictness: f32) -> f32 {
let total: f64 = coverage.iter().map(|&c| c as f64).sum();
let kept: f64 = r
.apply_coverage(coverage, strictness)
.iter()
.map(|&c| c as f64)
.sum();
((total - kept) / total) as f32
}
/// The promise a starting position has to keep: whatever it chooses, the
/// category is still there afterwards.
///
/// This is the fault it exists to fix, stated as an assertion. A fixed
/// strictness removed 76% to 99.5% of `architecture` on real photographs
/// — a mask that is empty on arrival, and indistinguishable from a broken
/// one, because the adjustment moves and no pixel changes.
#[test]
fn a_starting_strictness_never_empties_the_category() {
let opts = RefineOptions::default();
for (rect, colour, what) in [(FLAG, RED, "flag"), (CLOUD, WHITE, "cloud")] {
let (rgb, weights) = with(rect, colour);
let refinement = Refinement::compute(&weights, &rgb, EDGE, EDGE, CELL, &opts)
.expect("the frame has both sides");
let coverage = quantised(&weights);
let start = refinement.gentle(&coverage);
let cut = removed(&refinement, &coverage, start);
assert!(
cut <= GENTLE_CUT + 1e-3,
"the {what} frame starts at {start}, which takes {:.1}% of the category",
cut * 100.0
);
}
}
/// And it must not answer with zero out of caution.
///
/// A starting position that is always "off" would be a safe way of not
/// having the feature. On the frame the module was built for — a red flag
/// inside the sky — there *is* a strictness that takes the flag and leaves
/// the sky, and this has to find it.
#[test]
fn a_gentle_start_still_takes_the_flag_out_of_the_sky() {
let (rgb, weights) = with(FLAG, RED);
let refinement =
Refinement::compute(&weights, &rgb, EDGE, EDGE, CELL, &RefineOptions::default())
.expect("the flag frame has both sides");
let coverage = quantised(&weights);
let start = refinement.gentle(&coverage);
assert!(start > STRICTNESS_OFF, "gave up rather than choosing");
let refined = refinement.apply(&weights, start);
let inside = mean(&refined, EDGE, FLAG_CORE);
assert!(inside < 0.2, "the flag should be cut out, got {inside}");
let sky = mean(&refined, EDGE, (8, 8, 60, 60));
assert!(sky > 0.8, "the sky around it should survive, got {sky}");
}
/// A category nothing has claimed has nothing to judge, and asking must
/// not divide by its zero total.
#[test]
fn an_empty_category_starts_at_off() {
let (rgb, weights) = with(FLAG, RED);
let refinement =
Refinement::compute(&weights, &rgb, EDGE, EDGE, CELL, &RefineOptions::default())
.expect("the flag frame has both sides");
assert_eq!(
refinement.gentle(&vec![0u8; EDGE * EDGE]),
STRICTNESS_OFF,
"nothing to cut back"
);
}
/// Zero strictness is exactly the model's own weighting.
///
/// The control's off position has to be the old behaviour bit for bit, or
/// "turn it off and compare" does not answer the question it is asked.
#[test]
fn strictness_zero_changes_nothing() {
let (rgb, weights) = with(FLAG, RED);
let r = Refinement::compute(&weights, &rgb, EDGE, EDGE, CELL, &RefineOptions::default())
.expect("both sides present");
assert_eq!(r.apply(&weights, STRICTNESS_OFF), weights);
}
/// The one guarantee, and the one that replaced "subtractive".
///
/// 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 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");
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,
"the gate gave pixel {p} weight at {strictness}: {before} -> {after}"
);
}
}
}
/// A structure thinner than the erosion keeps a marker.
///
/// 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 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
/// uses the first and every test here uses the second.
#[test]
fn the_coverage_path_agrees_with_the_float_one() {
let (rgb, weights) = with(FLAG, RED);
let r = Refinement::compute(&weights, &rgb, EDGE, EDGE, CELL, &RefineOptions::default())
.expect("both sides present");
let coverage: Vec<u8> = weights.iter().map(|&w| (w * 255.0).round() as u8).collect();
let bytes = r.apply_coverage(&coverage, STRICTNESS_DEFAULT);
let floats = r.apply(&weights, STRICTNESS_DEFAULT);
for (p, (&b, &f)) in bytes.iter().zip(&floats).enumerate() {
let expected = (f * 255.0).round() as u8;
assert!(
b.abs_diff(expected) <= 1,
"pixel {p}: coverage {b}, float {expected}"
);
}
}
/// 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]
fn a_mask_thinner_than_the_grid_is_left_alone() {
let (rgb, _) = landscape(EDGE, EDGE);
// A three-pixel stripe: narrower than one cell, so the erosion at
// 1.5 cells empties it.
let mut weights = vec![0.0f32; EDGE * EDGE];
for y in 0..EDGE {
for x in 60..63 {
weights[y * EDGE + x] = 1.0;
}
}
assert_eq!(
Refinement::compute(&weights, &rgb, EDGE, EDGE, CELL, &RefineOptions::default()),
Err(SkipReason::NoInterior)
);
let (out, what) = refine_category(
&weights,
&rgb,
EDGE,
EDGE,
CELL,
STRICTNESS_DEFAULT,
&RefineOptions::default(),
);
assert_eq!(what, Refined::Skipped(SkipReason::NoInterior));
assert_eq!(out, weights, "a skip must return the input untouched");
}
/// A frame that is entirely one category has nothing to contrast against.
#[test]
fn a_frame_of_nothing_but_sky_is_left_alone() {
let rgb = vec![0.6f32; EDGE * EDGE * 3];
let weights = vec![1.0f32; EDGE * EDGE];
let (out, what) = refine_category(
&weights,
&rgb,
EDGE,
EDGE,
CELL,
STRICTNESS_DEFAULT,
&RefineOptions::default(),
);
assert_eq!(what, Refined::Skipped(SkipReason::NoExterior));
assert_eq!(out, weights);
}
/// Masks reach the sidecar as indices, so the same input must give the
/// same mask on every run and every machine — which is why the k-means
/// seeding uses a fixed sequence and the iteration count is fixed.
#[test]
fn the_same_input_gives_the_same_answer() {
let (rgb, weights) = with(FLAG, RED);
let opts = RefineOptions::default();
let a = Refinement::compute(&weights, &rgb, EDGE, EDGE, CELL, &opts).unwrap();
let b = Refinement::compute(&weights, &rgb, EDGE, EDGE, CELL, &opts).unwrap();
assert_eq!(
a.apply(&weights, STRICTNESS_DEFAULT),
b.apply(&weights, STRICTNESS_DEFAULT)
);
}
#[test]
fn buffers_that_disagree_are_refused() {
let (out, what) = refine_category(
&[1.0; 4],
&[0.5; 6],
2,
2,
CELL,
STRICTNESS_DEFAULT,
&RefineOptions::default(),
);
assert_eq!(what, Refined::Skipped(SkipReason::Mismatched));
assert_eq!(out, vec![1.0; 4]);
}
}