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

Mask the subject the model found, not the regions underneath it

Browse Source 
The watershed hierarchy does not survive a photograph, so local masking
stops depending on it. A layer can now be one recognised object, and the
object's own coverage is the mask.

`Options::watershed` defaults off. It costs ~80 ms plus a full-resolution
readback to produce a ladder that collapses, and paying that on every
photograph buys a control that misleads. Kept switchable rather than
deleted: the passes and the hierarchy are correct in themselves and it is
the merge criterion that fails, which is a change to one function.

Masks now rasterise in **source** space at proxy resolution and are sampled
by the composed shader after the framing map. That fixes a real bug: they
were rasterised in output space, so zooming slid the photograph underneath
a mask that stayed pinned to the viewport, and cropping moved every
adjustment to a different part of the picture. Doing it this way also
leaves the framing map in exactly one place — a second copy in the mask
shader would have been a second thing to keep in step, failing only when
straightened.

A subject is stored as identity, not pixels: the mask is megabytes and is
reproducible by running the same model over the same image, so the sidecar
carries the index, the class and the score, and the session carries the
pixels. The class is there to be checked — if instance 3 comes back a "car"
where it was a "dog", something changed and the layer is stale rather than
silently masking the wrong thing.

The overlay now draws instances and is transparent everywhere else. The
region version covered every pixel and so hid the photograph it was drawn
over; the question it exists to answer is whether an outline follows the
subject, which you can only answer by seeing both.

`examples/local.rs` is the worked example: subject in colour with the rest
monochrome, and the subject lifted out of its background. Run on a 5472x3648
CR2 it finds two people and two cars, and the colour-pop keeps her hat and
hair while the wall and grass behind go grey.
This commit is contained in:
dtourolle
2026-08-22 08:39:17 +02:00
parent b1433ad4a9
commit ee10097435
13 changed files with 963 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ui/dr-ui/src/segmentation.rs
+268 -66
View File
@@ -35,9 +35,17 @@ pub const DEFAULT_LEVEL: u32 = 300;
pub struct Segmentation {
/// The watershed's partition, with the semantic prior already applied to
/// its edge weights.
field: RegionField,
///
/// **Optional, and off by default.** The hierarchy built on it does not
/// survive a photograph: saddles are near zero almost everywhere, so the
/// merge order joins everything meaningful before it joins anything
/// spurious and a global cut returns one region plus noise. Measured on
/// both a compressed JPEG and a 22 MB CR2 — see docs/segmentation.md §15.
/// Until the merge criterion is replaced, local masking runs on the
/// model's instances and this stays off.
field: Option<RegionField>,
/// The merge order over `field`, which is the granularity ladder.
tree: MergeTree,
tree: Option<MergeTree>,
/// What the model found, for "select the subject" and for the layer names
/// a user actually recognises.
instances: Vec<InstanceSummary>,
@@ -55,6 +63,13 @@ pub struct Segmentation {
/// The grouping at the current level: region id -> group id.
grouping: Vec<u32>,
level: u32,
/// Size of the space instance masks and the label field are defined in.
///
/// Everything a mask is built from lives here, in **source** proxy pixels,
/// 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),
}
/// A detected object, reduced to what the interface needs.
@@ -68,7 +83,18 @@ pub struct InstanceSummary {
pub class_name: Arc<str>,
pub score: f32,
/// The regions this instance covers, snapped to watershed boundaries.
///
/// Empty without a watershed, which is the default. Kept because it is
/// the mechanism that would give a model outline the image's own edge, if
/// the hierarchy underneath it is ever made to work.
pub regions: Vec<u32>,
/// Coverage at proxy resolution, quantised to a byte.
///
/// **This is what the mask is now made of.** 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 instances would be most of
/// a hundred megabytes for a single photograph.
pub mask: Vec<u8>,
}
impl Segmentation {
@@ -85,7 +111,22 @@ impl Segmentation {
}
pub fn region_count(&self) -> usize {
self.field.region_count
self.field.as_ref().map_or(0, |f| f.region_count)
}
/// Whether a region hierarchy exists to click into.
pub fn has_regions(&self) -> bool {
self.field.is_some()
}
/// The space instance masks and the label field are defined in.
pub fn proxy_size(&self) -> (usize, usize) {
self.proxy
}
/// 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())
}
pub fn instances(&self) -> &[InstanceSummary] {
@@ -97,12 +138,15 @@ impl Segmentation {
/// Cheap — a union-find over a few thousand nodes — which is what makes
/// the scroll wheel a live control rather than a re-segmentation.
pub fn set_level(&mut self, level: u32) {
let level = level.clamp(2, self.field.region_count.max(2) as u32);
let (Some(field), Some(tree)) = (self.field.as_ref(), self.tree.as_ref()) else {
return;
};
let level = level.clamp(2, field.region_count.max(2) as u32);
if level == self.level {
return;
}
self.level = level;
self.grouping = self.tree.cut_to(level as usize);
self.grouping = tree.cut_to(level as usize);
}
/// The group under a point in **normalised image coordinates**.
@@ -111,18 +155,46 @@ impl Segmentation {
/// boxed, zoomed and possibly cropped; converting once at the call site is
/// clearer than passing three transforms in here.
pub fn group_at(&self, x: f32, y: f32) -> Option<u32> {
let field = self.field.as_ref()?;
if !(0.0..1.0).contains(&x) || !(0.0..1.0).contains(&y) {
return None;
}
let px = ((x * self.field.width as f32) as usize).min(self.field.width - 1);
let py = ((y * self.field.height as f32) as usize).min(self.field.height - 1);
let region = self.field.labels[py * self.field.width + px];
let px = ((x * field.width as f32) as usize).min(field.width - 1);
let py = ((y * field.height as f32) as usize).min(field.height - 1);
let region = field.labels[py * field.width + px];
self.grouping.get(region as usize).copied()
}
/// The strongest instance covering a point, for click-to-select.
///
/// Strongest rather than smallest: the detections are score-ordered and
/// overlapping ones are usually the same object found twice, so the more
/// confident is the better guess. A person standing 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;
let px = ((x * w as f32) as usize).min(w.saturating_sub(1));
let py = ((y * h as f32) as usize).min(h.saturating_sub(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)
}
/// Every region belonging to a group — the ids a mask stores.
pub fn regions_in_group(&self, group: u32) -> Vec<u32> {
(0..self.field.region_count as u32)
let Some(field) = self.field.as_ref() else {
return Vec::new();
};
(0..field.region_count as u32)
.filter(|&r| self.grouping.get(r as usize) == Some(&group))
.collect()
}
@@ -134,26 +206,95 @@ impl Segmentation {
.unwrap_or_default()
}
/// A false-coloured RGBA image of the current grouping.
/// A false-coloured RGBA picture of what a click can select.
///
/// The diagnostic that makes the hierarchy arguable instead of a matter of
/// faith — and the same thing a user wants when deciding whether a click
/// will land where they mean. Colours come from a golden-angle walk over
/// hue keyed by group id, so adjacent groups are very unlikely to share
/// one and the same group keeps its colour as the level changes.
/// The diagnostic that makes the segmentation arguable rather than a
/// matter of faith — and the thing a user wants when deciding whether a
/// click will land where they mean.
///
/// **Transparent where nothing is selectable**, which is the difference
/// between this and the region version it replaced. A region map covers
/// every pixel, so it hid the photograph it was drawn over; instances
/// cover the objects and nothing else, so the frame stays visible and the
/// question "does this outline follow the subject" can actually be
/// answered by looking.
pub fn overlay_rgba(&self) -> (Vec<u8>, u32, u32) {
let (w, h) = (self.field.width, self.field.height);
match self.field.as_ref() {
Some(field) => self.region_overlay(field),
None => self.instance_overlay(),
}
}
fn instance_overlay(&self) -> (Vec<u8>, u32, u32) {
let (w, h) = self.proxy;
let mut out = vec![0u8; w * h * 4];
for (p, &region) in self.field.labels.iter().enumerate() {
// Painted strongest-last, so where two detections overlap the more
// confident one is the colour on top — matching which of them 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 (rank, &i) in order.iter().enumerate() {
let instance = &self.instances[i];
// Keyed by the instance's own index rather than by paint order, so
// a colour belongs to an object and does not change when another
// detection appears beside it.
let [r, g, b] = group_colour(i as u32);
let _ = rank;
for (p, &cov) in instance.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 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)
}
fn region_overlay(&self, field: &RegionField) -> (Vec<u8>, u32, u32) {
let (w, h) = (field.width, field.height);
let mut out = vec![0u8; w * h * 4];
for (p, &region) in field.labels.iter().enumerate() {
let group = self.grouping.get(region as usize).copied().unwrap_or(0);
let [r, g, b] = group_colour(group);
out[p * 4] = r;
out[p * 4 + 1] = g;
out[p * 4 + 2] = b;
// Opaque here; the view composites it at whatever strength the
// user asked for. Baking transparency in would mean regenerating
// the whole image to change a slider.
out[p * 4 + 3] = 255;
}
@@ -162,11 +303,9 @@ impl Segmentation {
// exactly the judgement the overlay exists to support.
for y in 0..h {
for x in 0..w {
let here = self.grouping[self.field.labels[y * w + x] as usize];
let right = (x + 1 < w)
.then(|| self.grouping[self.field.labels[y * w + x + 1] as usize]);
let down = (y + 1 < h)
.then(|| self.grouping[self.field.labels[(y + 1) * w + x] as usize]);
let here = self.grouping[field.labels[y * w + x] as usize];
let right = (x + 1 < w).then(|| self.grouping[field.labels[y * w + x + 1] as usize]);
let down = (y + 1 < h).then(|| self.grouping[field.labels[(y + 1) * w + x] as usize]);
if right.is_some_and(|g| g != here) || down.is_some_and(|g| g != here) {
let p = (y * w + x) * 4;
out[p..p + 3].copy_from_slice(&[0, 0, 0]);
@@ -210,12 +349,18 @@ pub struct Options {
pub prior: PriorOptions,
/// Whether to run the model at all.
///
/// Off gives arm A alone: still a full region map with a working
/// granularity ladder, just without the semantic grouping — and about
/// seven times faster. Worth having as a choice rather than a fallback,
/// because on a landscape with no COCO class in it the model contributes
/// nothing and costs the whole inference.
/// This is the one that matters now: with it off there is nothing to
/// select, because the watershed below is not currently usable.
pub semantic: bool,
/// Whether to run the watershed and build the region hierarchy.
///
/// **Off by default.** It costs ~80 ms plus a full-resolution readback and
/// currently produces a ladder that collapses (docs/segmentation.md §15),
/// so paying for it on every photograph buys a control that misleads.
/// Kept switchable rather than deleted: the passes and the hierarchy are
/// correct in themselves and it is the *merge criterion* that fails, which
/// is a change to one function rather than to the subsystem.
pub watershed: bool,
}
impl Default for Options {
@@ -224,16 +369,19 @@ impl Default for Options {
segment: SegmentOptions::default(),
prior: PriorOptions::default(),
semantic: true,
watershed: false,
}
}
}
/// Run the watershed, optionally the model, and combine them.
/// Find what can be selected in this photograph.
///
/// `rgb` is the proxy the model reads — tightly packed RGB floats at
/// `(rgb_width, rgb_height)`. Passed in rather than derived here because the
/// caller already has the decoded image and re-deriving it would mean a second
/// readback of something the CPU is holding.
/// caller already has the decoded image, and re-deriving it would mean a
/// second readback of something the CPU is holding.
///
/// The watershed is optional and off by default; see [`Options::watershed`].
pub fn compute(
ctx: &GpuContext,
source: &dr_gpu::DemosaicedImage,
@@ -242,44 +390,74 @@ pub fn compute(
rgb_height: usize,
options: &Options,
) -> Result<Segmentation, String> {
let pass = SegmentPass::new(ctx).map_err(|e| format!("watershed unavailable: {e}"))?;
let raw = pass
.run(source, options.segment)
.map_err(|e| format!("watershed failed: {e}"))?;
let field = raw
.read_field()
.map_err(|e| format!("could not read the region field: {e}"))?;
// Arm B, and its failure is not fatal. A missing or unreadable model
// leaves a perfectly usable watershed segmentation; refusing to segment at
// all because the semantic half is unavailable would trade a working
// feature for a strict one.
let instances = if options.semantic {
match detect(rgb, rgb_width, rgb_height) {
Ok(found) => found,
Err(e) => {
log::warn!("semantic segmentation unavailable, using watershed alone: {e}");
Vec::new()
}
}
detect(rgb, rgb_width, rgb_height)?
} else {
Vec::new()
};
let (field, summaries) = apply_prior(field, &instances, &options.prior);
let tree = MergeTree::build(&field);
let grouping = tree.cut_to(DEFAULT_LEVEL as usize);
let mut field = None;
let mut tree = None;
let mut labels = None;
let mut grouping = Vec::new();
let mut summaries: Vec<InstanceSummary> = Vec::new();
let (w, h) = (field.width as u32, field.height as u32);
let labels = LabelField::upload(ctx, &field.labels, w, h, field.region_count as u32)
.inspect_err(|e| log::warn!("could not upload the label field: {e}"))
.ok();
if options.watershed {
let pass = SegmentPass::new(ctx).map_err(|e| format!("watershed unavailable: {e}"))?;
let raw = pass
.run(source, options.segment)
.map_err(|e| format!("watershed failed: {e}"))?;
let raw_field = raw
.read_field()
.map_err(|e| format!("could not read the region field: {e}"))?;
let (weighted, snapped) = apply_prior(raw_field, &instances, &options.prior);
let built = MergeTree::build(&weighted);
grouping = built.cut_to(DEFAULT_LEVEL as usize);
let (w, h) = (weighted.width as u32, weighted.height as u32);
labels = LabelField::upload(ctx, &weighted.labels, w, h, weighted.region_count as u32)
.inspect_err(|e| log::warn!("could not upload the label field: {e}"))
.ok();
summaries = snapped;
field = Some(weighted);
tree = Some(built);
}
// Whether or not the watershed ran, the instances themselves are the
// masks. `apply_prior` fills in region ids when there is a hierarchy to
// snap to; this fills in the pixels either way.
if summaries.is_empty() {
summaries = instances
.iter()
.filter(|i| i.score >= options.prior.confidence)
.map(|i| InstanceSummary {
class_name: i.class_name.clone(),
score: i.score,
regions: Vec::new(),
mask: quantise(&i.mask),
})
.collect();
} else {
for (summary, instance) in summaries.iter_mut().zip(
instances
.iter()
.filter(|i| i.score >= options.prior.confidence),
) {
summary.mask = quantise(&instance.mask);
}
}
let proxy = match field.as_ref() {
Some(f) => (f.width, f.height),
None => (rgb_width, rgb_height),
};
let signature = segmentation_signature(
w,
h,
field.region_count as u32,
proxy.0 as u32,
proxy.1 as u32,
summaries.len() as u32,
tuning_hash(options),
);
@@ -291,6 +469,7 @@ pub fn compute(
signature,
grouping,
level: DEFAULT_LEVEL,
proxy,
})
}
@@ -306,6 +485,16 @@ fn detect(rgb: &[f32], width: usize, height: usize) -> Result<Vec<dr_segment::In
.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()
}
/// Arm C: fold the instances into the merge weights, and record what each one
/// selects.
fn apply_prior(
@@ -337,8 +526,10 @@ fn apply_prior(
class_name: i.class_name.clone(),
score: i.score,
regions: prior::regions_for_instance(&field, &i.mask, options),
// Filled in by `compute`, which quantises every instance the same
// way whether or not a hierarchy was built.
mask: Vec::new(),
})
.filter(|s| !s.regions.is_empty())
.collect();
(
@@ -405,8 +596,9 @@ mod tests {
let tree = MergeTree::build(&field);
let grouping = tree.cut_to(level as usize);
Segmentation {
field,
tree,
proxy: (field.width, field.height),
field: Some(field),
tree: Some(tree),
instances: Vec::new(),
labels: None,
signature: 1,
@@ -586,6 +778,9 @@ mod tests {
// point of this test is the watershed half.
let options = Options {
semantic: false,
// Explicit: the watershed is off by default now, and these tests
// exist to exercise it.
watershed: true,
..Options::default()
};
@@ -631,6 +826,9 @@ mod tests {
let options = Options {
semantic: false,
// Explicit: the watershed is off by default now, and these tests
// exist to exercise it.
watershed: true,
..Options::default()
};
let mut seg = compute(&ctx, &source, &rgb, SIZE as usize, SIZE as usize, &options)
@@ -669,6 +867,7 @@ mod tests {
let mut coarse = Options {
semantic: false,
watershed: true,
..Options::default()
};
coarse.segment.max_edge = 96;
@@ -703,6 +902,9 @@ mod tests {
.collect();
let options = Options {
semantic: false,
// Explicit: the watershed is off by default now, and these tests
// exist to exercise it.
watershed: true,
..Options::default()
};