Files
DarkRoom/ui/dr-ui/src/segmentation.rs
T
dtourolleandClaude Opus 5 b55812812a Name segmented people from the faces already recognised in them
The segmenter knows it found a person; the face index knows which person.
Joining them turns "person" in the mask list into "Anna", which is the
difference between a vocabulary of eighty COCO classes and one that
includes the user's family. Selecting a subject in a group photograph
stops being a guessing game between three identical rows.

Containment, not IoU. A face is a small part of the person it belongs to,
so a correct pairing has an IoU near zero and anything IoU-based would
reject every true match.

Confirmed names only. A suggestion is the system's guess, and printing a
guessed name onto a mask region would launder it into a fact.

Writing the tests corrected the design once: a tight head-and-shoulders
portrait, where the face fills most of the person box, is the case where
naming is most certain, not least. An earlier guard rejected exactly that
and has been removed, with the reasoning left as a test because it is
easy to get backwards a second time.

The names hang on the develop session, set when the image opens because
that is the one moment the catalog and the image id are both in reach.
Every segmentation run afterwards picks them up for free, and a library
with no face indexing behaves exactly as it did before.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-08-26 23:10:46 +02:00

511 lines
19 KiB
Rust

//! Finding what can be selected in the photograph on screen.
//!
//! One model pass, and what it recognised. That is the whole subsystem now.
//!
//! # What used to be here
//!
//! A watershed over-segmented the image, a merge tree turned that into a
//! granularity ladder, and a click walked up it (docs/segmentation.md, arms A
//! and C). It is gone from this path, and the reason is measured rather than
//! aesthetic: on a real photograph the saddles are near zero almost
//! everywhere, so the merge order joins everything meaningful before it joins
//! anything spurious. Cutting a 45,808-basin field of a 22 MP frame to 2,000
//! regions left **one** region covering nearly the whole picture plus specks
//! (§15). A ladder that collapses is not a ladder.
//!
//! The passes and the hierarchy still exist in `dr-gpu` and `dr-segment`,
//! tested and documented, because it is the *merge criterion* that fails and
//! that is one function. What does not exist any more is a product path
//! through them, or a control offering a choice that does nothing.
//!
//! # This is a precompute, and it is slow
//!
//! ~700 ms on a 22 MP frame: a proxy render, a readback, and the model. It
//! runs **once per image, when the user asks**, and never on the frame path.
//! Every interaction it enables — click a subject, grow a mask, change a
//! falloff — reads its cached output.
//!
//! All of it is on a worker. [`compute`] takes an owned buffer and a
//! `GpuContext`, which is what makes that possible — nothing here touches the
//! develop session, and [`crate::develop::SegmentationJob`] is the piece that
//! carries the proxy render across with it.
use std::sync::Arc;
use dr_gpu::GpuContext;
use dr_pipeline::mask::segmentation_signature;
/// One recognised object.
#[derive(Debug, Clone)]
pub struct InstanceSummary {
pub class_name: Arc<str>,
pub score: f32,
/// Coverage at proxy resolution, quantised to a byte.
///
/// A byte rather than the `f32` the model produces: 256 levels is far
/// finer than an edge anyone can see, and at four bytes a pixel a handful
/// of objects would be most of a hundred megabytes for one photograph.
///
/// This is the *source* a distance field is built from, not the mask
/// itself — `dr_segment::Shaped` turns it into one.
pub mask: Vec<u8>,
/// `(x0, y0, x1, y1)` in [`Segmentation::proxy_size`] pixels.
///
/// Carried through from `dr_segment::Instance` rather than re-derived
/// from the mask, so a refine pass knows what region to crop without
/// scanning a proxy-sized buffer for its own extent.
pub bbox: (f32, f32, f32, f32),
}
/// One image's recognised objects, ready to mask.
pub struct Segmentation {
instances: Vec<InstanceSummary>,
/// Identifies this run, so a stored layer can tell whether the index it
/// holds still means what it meant.
signature: u64,
/// The space instance masks are defined in, in **source** proxy pixels.
///
/// Everything a mask is built from lives here, which is what lets the
/// render sample it *after* the framing map rather than before — so a mask
/// stays on the photograph through a zoom, a pan and a crop.
proxy: (usize, usize),
}
impl Segmentation {
pub fn signature(&self) -> u64 {
self.signature
}
pub fn proxy_size(&self) -> (usize, usize) {
self.proxy
}
pub fn instances(&self) -> &[InstanceSummary] {
&self.instances
}
/// One instance's coverage, at [`Self::proxy_size`].
pub fn instance_mask(&self, index: usize) -> Option<&[u8]> {
self.instances.get(index).map(|i| i.mask.as_slice())
}
/// Replace one instance in place, keeping every other index and the
/// signature unchanged.
///
/// What a refine pass calls once it has a sharper mask for the subject at
/// `index`: the layers pointing at this run by index still mean what they
/// meant, they just read better pixels now.
pub fn replace_instance(&mut self, index: usize, instance: InstanceSummary) {
if let Some(slot) = self.instances.get_mut(index) {
*slot = instance;
}
}
/// The strongest instance covering a point in normalised image
/// coordinates.
///
/// Strongest rather than smallest: detections are score-ordered and
/// overlapping ones are usually the same object found twice, so the more
/// confident is the better guess. A person in front of a bus wins over the
/// bus, because the person's mask is the one under the cursor at all.
pub fn instance_at(&self, x: f32, y: f32) -> Option<usize> {
if !(0.0..1.0).contains(&x) || !(0.0..1.0).contains(&y) {
return None;
}
let (w, h) = self.proxy;
if w == 0 || h == 0 {
return None;
}
let px = ((x * w as f32) as usize).min(w - 1);
let py = ((y * h as f32) as usize).min(h - 1);
let p = py * w + px;
self.instances
.iter()
.enumerate()
.filter(|(_, i)| i.mask.get(p).is_some_and(|&v| v >= 128))
.max_by(|(_, a), (_, b)| a.score.total_cmp(&b.score))
.map(|(i, _)| i)
}
/// A false-coloured picture of what a click can select, in source space.
///
/// **Transparent where nothing is selectable.** The region map this
/// replaced covered every pixel and so hid the photograph it was drawn
/// over; the question an overlay exists to answer is whether an outline
/// follows the subject, and that can only be answered by seeing both.
pub fn overlay_rgba(&self) -> (Vec<u8>, u32, u32) {
let (w, h) = self.proxy;
let mut out = vec![0u8; w * h * 4];
// Weakest first, so where two detections overlap the more confident
// one is the colour on top — matching which a click would select.
let mut order: Vec<usize> = (0..self.instances.len()).collect();
order.sort_by(|&a, &b| self.instances[a].score.total_cmp(&self.instances[b].score));
for &i in &order {
let [r, g, b] = instance_colour(i as u32);
for (p, &cov) in self.instances[i].mask.iter().enumerate() {
if cov < 128 || p * 4 + 3 >= out.len() {
continue;
}
out[p * 4] = r;
out[p * 4 + 1] = g;
out[p * 4 + 2] = b;
out[p * 4 + 3] = 255;
}
}
// The outline drawn opaque white over the fill. It is the part being
// judged — a fill can look right while its edge sits several pixels
// off the subject — and it survives the low opacity the fill is
// composited at.
let solid = |p: usize| out.get(p * 4 + 3).is_some_and(|&a| a > 0);
let mut edges = Vec::new();
for y in 0..h {
for x in 0..w {
let p = y * w + x;
if !solid(p) {
continue;
}
let boundary = (x + 1 == w || !solid(p + 1))
|| (x == 0 || !solid(p - 1))
|| (y + 1 == h || !solid(p + w))
|| (y == 0 || !solid(p - w));
if boundary {
edges.push(p);
}
}
}
for p in edges {
out[p * 4..p * 4 + 4].copy_from_slice(&[255, 255, 255, 255]);
}
(out, w as u32, h as u32)
}
}
/// A distinct colour per instance.
///
/// Golden-angle hue stepping, deterministic rather than random: the same
/// object is the same colour every time the overlay is drawn, so the eye can
/// track it while a mask is being shaped.
fn instance_colour(index: u32) -> [u8; 3] {
let h = (index as f32 * 137.508) % 360.0;
let c = 230.0;
let x = c * (1.0 - ((h / 60.0) % 2.0 - 1.0).abs());
let (r, g, b) = match (h / 60.0) as u32 {
0 => (c, x, 0.0),
1 => (x, c, 0.0),
2 => (0.0, c, x),
3 => (0.0, x, c),
4 => (x, 0.0, c),
_ => (c, 0.0, x),
};
[r as u8 + 25, g as u8 + 25, b as u8 + 25]
}
/// What to run.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Options {
/// Detections below this are dropped.
///
/// Deliberately low. A weak detection costs a spurious entry in a list the
/// user is choosing from, and its score is shown beside it; a missed one
/// costs a subject that cannot be selected at all, which is the worse
/// failure for a selection tool.
pub confidence: f32,
/// TRACES: FR-DEV-3
/// Run the model over overlapping tiles instead of the whole frame once.
///
/// Off by default and deliberately so. The graph's input is fixed at
/// 640x640 (docs/segmentation.md F6), so every image is letterboxed into
/// it and a subject 200px across in a 1600px proxy reaches the model at
/// 80px — which is where a coarse outline comes from. Tiling is the only
/// route to more resolution with a fixed window, and it costs one
/// inference per tile: about 2.8s for a 3x2 grid against 470ms whole-frame.
///
/// Six times the wait is the wrong default for the common case, where the
/// subject is large in frame and whole-frame inference is already the best
/// answer. It is the right answer for a bird against sky, so it is offered
/// per-image rather than chosen once for all of them.
pub fine: bool,
}
impl Default for Options {
fn default() -> Self {
Self {
confidence: 0.30,
fine: false,
}
}
}
/// Find what can be selected in this photograph.
///
/// `rgb` is the proxy the model reads — tightly packed RGB floats at
/// `(width, height)`. Passed in rather than derived here because the caller
/// already has the rendered proxy, and re-deriving it would mean a second
/// readback of something the CPU is holding.
pub fn compute(
_ctx: &GpuContext,
rgb: &[f32],
width: usize,
height: usize,
options: &Options,
) -> Result<Segmentation, String> {
let found = detect(rgb, width, height, options.fine)?;
let instances: Vec<InstanceSummary> = found
.iter()
.filter(|i| i.score >= options.confidence)
.map(|i| InstanceSummary {
class_name: i.class_name.clone(),
score: i.score,
mask: quantise(&i.mask),
bbox: i.bbox,
})
.collect();
let signature = segmentation_signature(
width as u32,
height as u32,
instances.len() as u32,
// The tiling choice belongs in the signature as much as the
// confidence does. A mask stores the signature of the segmentation its
// region ids index into (`MaskSource::Regions`), and a tiled run finds
// different instances in a different order — so if the two runs shared
// a signature, a layer built against the coarse pass would be silently
// reinterpreted against the fine one. That is a *wrong* mask, which is
// far worse than a stale one, because nothing announces it.
options.confidence.to_bits() as u64
^ if options.fine {
0x9E37_79B9_7F4A_7C15
} else {
0
},
);
Ok(Segmentation {
instances,
signature,
proxy: (width, height),
})
}
/// Classes a recognised face may put a name on.
///
/// Only these. A face inside a `tv` or a `laptop` is a photograph of someone on
/// a screen, and renaming the television to "Anna" would be worse than leaving
/// it as the model found it.
fn is_person_class(name: &str) -> bool {
name == "person"
}
impl Segmentation {
/// Relabel `person` instances with the name of the face inside them.
///
/// The segmenter knows it found *a person*; the face index knows *which*
/// person. Joining them turns "person" in the mask list into "Anna", which
/// is the difference between COCO's eighty classes and a vocabulary that
/// includes the user's family — and it is the same click either way, so the
/// gain is entirely in being able to tell two people apart before clicking.
///
/// `faces` are in the same proxy pixels as [`InstanceSummary::bbox`]. The
/// geometry lives in `dr_face::naming`, where it is testable without a
/// model or a catalog.
///
/// Returns how many instances gained a name. An unrecognised person keeps
/// the model's own label, which is the right default: "person" is merely
/// unhelpful, where a wrong name is wrong and the user cannot tell which
/// they are looking at.
pub fn apply_names(&mut self, faces: &[dr_face::NamedFace<'_>]) -> usize {
dr_face::name_instances(
&mut self.instances,
faces,
|i| i.bbox,
|i| is_person_class(&i.class_name),
|i, name| i.class_name = name.into(),
)
}
}
/// Load the model and run it.
///
/// Loading is ~24 ms against the ~470 ms of inference that follows, and this
/// runs once per image — so caching the session would keep 11 MB of weights
/// resident for the life of the app to save five percent of a background task.
fn detect(
rgb: &[f32],
width: usize,
height: usize,
fine: bool,
) -> Result<Vec<dr_segment::Instance>, String> {
let mut model = dr_segment::SemanticModel::embedded().map_err(|e| e.to_string())?;
let options = dr_segment::SemanticOptions {
// A quarter shared with each neighbour. It has to exceed zero at all,
// or a subject sitting on a seam is cut in half by both tiles and
// recognised by neither; a quarter is enough to carry a whole subject
// inside one tile at the sizes tiling is reached for.
tiling: if fine {
dr_segment::Tiling::Grid { overlap: 0.25 }
} else {
dr_segment::Tiling::Whole
},
..dr_segment::SemanticOptions::default()
};
model
.detect(rgb, width, height, &options)
.map_err(|e| e.to_string())
}
/// The model's soft coverage, to a byte per pixel.
///
/// Rounded rather than truncated, so a coverage of exactly 0.5 lands on the
/// threshold the selection tests against instead of one below it.
fn quantise(mask: &[f32]) -> Vec<u8> {
mask.iter()
.map(|&v| (v.clamp(0.0, 1.0) * 255.0).round() as u8)
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
/// Two objects: a big weak one on the left, a small strong one that
/// overlaps it.
fn overlapping() -> Segmentation {
let (w, h) = (8usize, 4usize);
let mut big = vec![0u8; w * h];
let mut small = vec![0u8; w * h];
for y in 0..h {
for x in 0..6 {
big[y * w + x] = 255;
}
for x in 4..8 {
small[y * w + x] = 255;
}
}
Segmentation {
instances: vec![
InstanceSummary {
class_name: "bus".into(),
score: 0.5,
mask: big,
bbox: (0.0, 0.0, 6.0, h as f32),
},
InstanceSummary {
class_name: "person".into(),
score: 0.9,
mask: small,
bbox: (4.0, 0.0, 8.0, h as f32),
},
],
signature: 1,
proxy: (w, h),
}
}
#[test]
fn a_click_finds_the_object_under_it() {
let seg = overlapping();
assert_eq!(seg.instance_at(0.1, 0.5), Some(0), "only the bus here");
assert_eq!(seg.instance_at(0.95, 0.5), Some(1), "only the person here");
}
/// The overlap rule, and the one that decides what a click means where two
/// detections cover the same pixel.
#[test]
fn overlapping_objects_resolve_to_the_more_confident() {
let seg = overlapping();
assert_eq!(
seg.instance_at(0.6, 0.5),
Some(1),
"the person at 0.9 beats the bus at 0.5"
);
}
#[test]
fn a_click_outside_the_frame_selects_nothing() {
let seg = overlapping();
assert_eq!(seg.instance_at(-0.1, 0.5), None);
assert_eq!(seg.instance_at(1.5, 0.5), None);
assert_eq!(seg.instance_at(0.5, 1.0), None, "the far edge is exclusive");
}
#[test]
fn a_click_on_nothing_selects_nothing() {
let seg = Segmentation {
instances: Vec::new(),
signature: 1,
proxy: (4, 4),
};
assert_eq!(seg.instance_at(0.5, 0.5), None);
}
#[test]
fn colours_are_stable_and_distinct() {
assert_eq!(instance_colour(7), instance_colour(7));
assert_ne!(instance_colour(0), instance_colour(1));
assert_ne!(instance_colour(1), instance_colour(2));
}
#[test]
fn every_colour_is_visible_against_a_photograph() {
for i in 0..64u32 {
let [r, g, b] = instance_colour(i);
assert!(r.max(g).max(b) >= 200, "instance {i} is too dark");
}
}
/// The overlay must not cover the picture: that is the difference between
/// this and the region map it replaced.
#[test]
fn the_overlay_is_transparent_where_nothing_was_found() {
let seg = Segmentation {
instances: Vec::new(),
signature: 1,
proxy: (4, 4),
};
let (px, w, h) = seg.overlay_rgba();
assert_eq!((w, h), (4, 4));
assert!(px.chunks_exact(4).all(|p| p[3] == 0), "nothing to draw");
}
#[test]
fn the_overlay_outlines_what_it_fills() {
let seg = overlapping();
let (px, w, _) = seg.overlay_rgba();
let at = |x: usize, y: usize| {
let p = (y * w as usize + x) * 4;
[px[p], px[p + 1], px[p + 2], px[p + 3]]
};
// The rightmost column of the person is against the frame edge, so it
// is an outline pixel.
assert_eq!(at(7, 1), [255, 255, 255, 255]);
// And an interior pixel keeps its fill.
assert_ne!(at(2, 1)[3], 0, "the bus is filled");
assert_ne!(at(2, 1), [255, 255, 255, 255], "and not all outline");
}
#[test]
fn quantising_rounds_rather_than_truncates() {
// Exactly half must reach the threshold the selection tests against.
assert_eq!(quantise(&[0.5]), vec![128]);
assert_eq!(quantise(&[0.0, 1.0]), vec![0, 255]);
// And values outside the range cannot wrap.
assert_eq!(quantise(&[-1.0, 2.0]), vec![0, 255]);
}
#[test]
fn confidence_changes_the_signature() {
// A different threshold is a different instance list, so the indices a
// stored layer holds mean something else.
let a = segmentation_signature(100, 100, 3, Options::default().confidence.to_bits() as u64);
let b = segmentation_signature(100, 100, 3, 0.9f32.to_bits() as u64);
assert_ne!(a, b);
}
}