Put the region map behind the sliders that were already there

A mask layer holds a real develop chain, so the develop panel can edit one
with no new controls: select a layer and the same sliders read and write
its chain instead of the graph's. An operation declared in `ops/` tomorrow
becomes locally adjustable by existing, which is the payoff for making a
layer a chain rather than a handful of special-cased parameters.

`segmentation.rs` joins the two arms into the one thing the view needs.
The model reads the image through a neutral graph rather than the edited
one, so a segmentation survives an exposure change instead of being
invalidated by every slider. Arm B failing is not fatal: a missing or
unreadable model leaves a working watershed map, because refusing to
segment at all would trade a working feature for a strict one.

The overlay colours groups by a golden-angle walk over hue. Deterministic
rather than random, so a region keeps its colour across a level change and
the eye can track it; boundaries drawn black over the fill, because two
adjacent groups landing on near hues read as one region and telling them
apart is the whole reason to look at it.

Clicking the photograph creates the layer if none is selected — that is how
a local adjustment begins, and making the user press "add layer" first
would be a step with no decision in it. Shift-click extends, and clicking a
region already selected removes it, so one gesture both adds and corrects.

`segment-readback` is a new dr-gpu feature and not a loosening of
`readback`. The region-graph transfer is once per image on a worker; the
one AC-8 forbids is per frame in the render loop. Sharing a switch would
have forced a build wanting local masking to unlock the other. F3 still
stands and the feature name says so.
This commit is contained in:
2026-08-22 08:39:17 +02:00
parent 94cfea4748
commit 5ecb35864f
8 changed files with 1080 additions and 21 deletions
+8 -1
View File
@@ -11,7 +11,14 @@ dr-types.workspace = true
# the compositor the texture itself and there is no display round-trip left to
# gate (ARCH §6.1, AC-8). The export path reads pixels back through
# `export_pixels`, which is ungated and always was.
dr-gpu.workspace = true
#
# `segment-readback` *is* on, and it is not a contradiction of the above. It
# gates the region-graph transfer that local masking is built on: once per
# image, on a worker, off the frame path. The display round-trip AC-8 forbids
# stays behind its own switch, which remains off. See `dr-gpu/src/segment.rs`.
dr-gpu = { workspace = true, features = ["segment-readback"] }
# The semantic arm and its weights, for local adjustments (FR-DEV-3, D14).
dr-segment = { workspace = true, features = ["semantic", "embedded-model"] }
dr-decode.workspace = true
serde_json.workspace = true
tokio.workspace = true
+444 -6
View File
@@ -10,7 +10,12 @@
//! new operation appears in the panel with no change here (FR-DEV-3c).
use dr_decode::RawImage;
use dr_gpu::{AdjustPass, DemosaicedImage, Demosaicer, GpuContext, Histogram, HistogramPass};
use dr_gpu::{
AdjustPass, DemosaicedImage, Demosaicer, GpuContext, Histogram, HistogramPass, MaskPass,
};
use dr_pipeline::mask::{MaskLayer, MaskSource};
use crate::segmentation::{self, Segmentation};
use dr_pipeline::ops::curve;
use dr_pipeline::{
CropRect, Edit, EditGraph, History, OpCapability, OpId, ParamId, ParamKind, Presentation,
@@ -40,6 +45,24 @@ pub struct DevelopSession {
/// old driver, a device without the storage-buffer atomics it needs — the
/// photographer loses the histogram and keeps the photograph.
histogram: Option<HistogramPass>,
/// TRACES: FR-DEV-3
/// The region map local masks select from, once it has been computed.
///
/// `None` until the photographer asks for it. Segmentation costs about
/// half a second and most edits never need one, so running it on open
/// would tax every photograph for a feature used on some of them.
segmentation: Option<Segmentation>,
/// Rasterises the mask layers. Built lazily for the same reason.
masks: Option<MaskPass>,
/// Which layer the develop panel is editing, if any.
///
/// This is what lets one panel serve both scopes: with a layer selected,
/// the sliders read and write *its* chain, and the photographer is
/// adjusting a region rather than the frame.
active_mask: Option<String>,
/// Whether to draw the false-coloured region overlay.
show_overlay: bool,
}
impl DevelopSession {
@@ -97,6 +120,10 @@ impl DevelopSession {
histogram: HistogramPass::new(ctx)
.inspect_err(|e| log::warn!("no histogram on this device: {e}"))
.ok(),
segmentation: None,
masks: None,
active_mask: None,
show_overlay: false,
}
}
@@ -105,7 +132,24 @@ impl DevelopSession {
/// Built entirely from the capability list. The `kind` string chooses the
/// widget; nothing switches on a parameter's identity.
pub fn rows(&self) -> Vec<ParamRow> {
rows_from(&self.graph.capabilities())
match self.active_layer() {
// A selected layer takes over the panel. The rows are built from
// the layer's own capability list, so every control the global
// chain offers is offered here too — including ones added later,
// which need no work to become local.
Some(layer) => rows_from(&layer.capabilities()),
None => rows_from(&self.graph.capabilities()),
}
}
fn active_layer(&self) -> Option<&MaskLayer> {
let id = self.active_mask.as_ref()?;
self.graph.masks().get(id)
}
fn active_layer_mut(&mut self) -> Option<&mut MaskLayer> {
let id = self.active_mask.clone()?;
self.graph.masks_mut().get_mut(&id)
}
}
@@ -463,10 +507,24 @@ impl DevelopSession {
/// here is curve-shaped; it walks whatever parameters the operation
/// declares.
pub fn reset_op(&mut self, op_index: i32) {
let caps = self.graph.capabilities();
let caps = match self.active_layer() {
Some(layer) => layer.capabilities(),
None => self.graph.capabilities(),
};
let Some(cap) = usize::try_from(op_index).ok().and_then(|i| caps.get(i)) else {
return;
};
if self.active_mask.is_some() {
let params: Vec<_> = cap.params.iter().map(|p| (p.id, p.default)).collect();
let id = cap.id.0;
if let Some(layer) = self.active_layer_mut() {
for (param, default) in params {
layer.set_param(id, param, default);
}
}
self.history.record(&self.graph, Edit::Discrete);
return;
}
for p in &cap.params {
self.graph.set_param(cap.id, p.id, p.default);
}
@@ -493,6 +551,17 @@ impl DevelopSession {
log::warn!("control at ({op_index}, {param_index}) has no parameter");
return;
};
if self.active_mask.is_some() {
if let Some(layer) = self.active_layer_mut() {
layer.set_param(op.0, param, value);
}
// Coalesced the same way a global drag is: a slider dragged across
// a masked layer is still one gesture and must undo as one.
let edit = Edit::for_param(&self.graph, op, param);
self.history.record(&self.graph, edit);
return;
}
self.graph.set_param(op, param, value);
let edit = Edit::for_param(&self.graph, op, param);
self.history.record(&self.graph, edit);
@@ -520,10 +589,372 @@ impl DevelopSession {
self.history.record(&self.graph, Edit::Discrete);
}
/// Rasterise the current mask stack, if there is one.
///
/// Returns `None` for a stack with no active layers, which is the common
/// case and the one that must cost nothing: the adjust pass then binds its
/// own placeholder and the generated shader has no layer block to read it
/// with.
/// Returns whether the array is now valid for the current stack.
///
/// Split from reading the array back because the render below needs
/// `self.adjust` mutably while holding `self.masks` immutably. Those are
/// disjoint fields and the borrow checker will allow it — but only when
/// each is reached directly rather than through a method taking `self`.
fn rasterise_masks(&mut self, w: u32, h: u32) -> bool {
if self.graph.masks().is_neutral() {
return false;
}
let labels = self.segmentation.as_ref().and_then(|s| s.labels());
let Some(pass) = self.masks.as_mut() else {
return false;
};
pass.render(self.graph.masks(), labels, w, h)
.inspect_err(|e| log::warn!("mask rasterisation failed: {e}"))
.is_ok()
}
// ----------------------------------------------------------------------
// Segmentation (S15, docs/segmentation.md)
// ----------------------------------------------------------------------
/// TRACES: FR-DEV-3
/// Compute the region map this image's local masks select from.
///
/// **Blocking, and roughly half a second.** The caller is responsible for
/// running it off the UI thread — see the worker in `lib.rs`. It is
/// exposed as a plain blocking call rather than something async because
/// what it needs is a GPU context and a CPU core, not a runtime.
pub fn segment(&mut self, ctx: &GpuContext, options: &segmentation::Options) -> Result<(), String> {
// The model reads the photograph as captured, not as edited: the
// segmentation must survive an exposure change, or every slider would
// invalidate the masks that depend on it (docs/segmentation.md §3).
let (rgb, rw, rh) = self.neutral_proxy(ctx, options.segment.max_edge)?;
let seg = segmentation::compute(ctx, &self.demosaiced, &rgb, rw, rh, options)?;
if self.masks.is_none() {
self.masks = MaskPass::new(ctx)
.inspect_err(|e| log::warn!("no mask rasteriser on this device: {e}"))
.ok();
}
self.segmentation = Some(seg);
Ok(())
}
/// Render the *unedited* image to a CPU buffer at proxy size.
///
/// Goes through a throwaway [`AdjustPass`] with a neutral graph rather
/// than the session's own. Reusing `self.adjust` would overwrite the frame
/// the histogram reads and leave the view showing an unedited image until
/// the next redraw — a visible flicker for the sake of not allocating.
///
/// This is `export_pixels`, which is ungated: an export is not the display
/// round-trip AC-8 forbids, and neither is this.
fn neutral_proxy(
&self,
ctx: &GpuContext,
max_edge: u32,
) -> Result<(Vec<f32>, usize, usize), String> {
let (sw, sh) = self.demosaiced.size();
let scale = (max_edge as f32 / sw.max(sh) as f32).min(1.0);
let (w, h) = (
((sw as f32 * scale) as u32).max(1),
((sh as f32 * scale) as u32).max(1),
);
let neutral = EditGraph::default_chain();
let mut pass = AdjustPass::new(ctx);
pass.render(&self.demosaiced, &neutral.compose(), w, h)
.map_err(|e| format!("could not render the segmentation proxy: {e}"))?;
let (rgba, pw, ph) = pass
.export_pixels()
.map_err(|e| format!("could not read the segmentation proxy: {e}"))?;
// Straight to float RGB, dropping alpha. The values stay display-
// encoded because that is what the model was trained on — one of the
// few places in this codebase where not linearising is correct.
let rgb = rgba
.chunks_exact(4)
.flat_map(|p| [p[0] as f32 / 255.0, p[1] as f32 / 255.0, p[2] as f32 / 255.0])
.collect();
Ok((rgb, pw as usize, ph as usize))
}
pub fn has_segmentation(&self) -> bool {
self.segmentation.is_some()
}
pub fn segmentation_level(&self) -> u32 {
self.segmentation.as_ref().map_or(0, |s| s.level())
}
pub fn segmentation_region_count(&self) -> usize {
self.segmentation.as_ref().map_or(0, |s| s.region_count())
}
/// Move the granularity ladder — the scroll wheel over the canvas.
pub fn set_segmentation_level(&mut self, level: u32) {
if let Some(seg) = self.segmentation.as_mut() {
seg.set_level(level);
}
}
/// The subjects the model recognised, as `(label, confidence)`.
///
/// Confidence is shown rather than hidden because the detector is offered
/// as a shortcut, not as an authority: a 0.42 "dog" is worth listing and
/// worth flagging, and a list that presented it identically to a 0.95 one
/// would make the tool look wrong when the guess was merely weak.
pub fn detected_subjects(&self) -> Vec<(String, f32)> {
self.segmentation
.as_ref()
.map(|s| {
s.instances()
.iter()
.map(|i| (i.class_name.to_string(), i.score))
.collect()
})
.unwrap_or_default()
}
// ----------------------------------------------------------------------
// The region overlay
// ----------------------------------------------------------------------
pub fn overlay_enabled(&self) -> bool {
self.show_overlay
}
pub fn set_overlay(&mut self, on: bool) {
self.show_overlay = on;
}
/// TRACES: FR-DEV-3
/// A false-coloured picture of the current grouping, for the canvas.
///
/// Returned as a CPU image rather than a texture, and deliberately: it is
/// regenerated only when the level changes, it is proxy-sized rather than
/// viewport-sized, and Slint scales and composites it for free. Putting it
/// on the GPU would buy nothing and add a second texture to keep in step
/// with the view.
///
/// `None` when the overlay is off or nothing has been segmented, so the
/// caller can bind this straight to an image source.
pub fn overlay_image(&self) -> Option<slint::Image> {
if !self.show_overlay {
return None;
}
let (rgba, w, h) = self.segmentation.as_ref()?.overlay_rgba();
let buffer = slint::SharedPixelBuffer::<slint::Rgba8Pixel>::clone_from_slice(&rgba, w, h);
Some(slint::Image::from_rgba8(buffer))
}
// ----------------------------------------------------------------------
// Mask layers
// ----------------------------------------------------------------------
/// The layers, as `(id, name, enabled, is_active_selection)`.
pub fn mask_layers(&self) -> Vec<(String, String, bool, bool)> {
self.graph
.masks()
.layers()
.iter()
.map(|l| {
(
l.id.clone(),
l.display_name().to_string(),
l.enabled,
Some(&l.id) == self.active_mask.as_ref(),
)
})
.collect()
}
pub fn active_mask(&self) -> Option<&str> {
self.active_mask.as_deref()
}
/// Select a layer for editing, or `None` to return the panel to the
/// global chain.
pub fn set_active_mask(&mut self, id: Option<&str>) {
self.active_mask = id
.filter(|id| self.graph.masks().get(id).is_some())
.map(|id| id.to_string());
}
/// Select the region under a normalised image point.
///
/// `add` extends the selected layer instead of replacing its selection,
/// which is the shift-click every selection tool has. With no layer
/// selected a new one is created, because clicking the photograph is how a
/// local adjustment begins and requiring "add layer" first would be a step
/// with no decision in it.
///
/// Returns the layer that now holds the selection.
pub fn select_region_at(&mut self, x: f32, y: f32, add: bool) -> Option<String> {
let seg = self.segmentation.as_ref()?;
let picked = seg.regions_at(x, y);
if picked.is_empty() {
return None;
}
let (signature, level) = (seg.signature(), seg.level());
let id = match self.active_mask.clone() {
Some(id) => id,
None => {
let id = self.graph.masks().next_id();
let layer = MaskLayer::new(
id.clone(),
MaskSource::Regions {
signature,
level,
ids: Vec::new(),
},
);
if !self.graph.masks_mut().push(layer) {
return None;
}
self.active_mask = Some(id.clone());
id
}
};
let layer = self.graph.masks_mut().get_mut(&id)?;
let mut ids = match (&layer.source, add) {
(MaskSource::Regions { ids, .. }, true) => ids.clone(),
_ => Vec::new(),
};
// Clicking a region already in the selection removes it, so one
// gesture both adds and corrects — the alternative is a modifier for
// subtract that nobody remembers.
if add && picked.iter().all(|r| ids.contains(r)) {
ids.retain(|r| !picked.contains(r));
} else {
ids.extend(picked);
}
ids.sort_unstable();
ids.dedup();
layer.source = MaskSource::Regions {
signature,
level,
ids,
};
self.history.record(&self.graph, Edit::Discrete);
Some(id)
}
/// Add a layer selecting one detected subject.
pub fn add_subject_mask(&mut self, index: usize) -> Option<String> {
let seg = self.segmentation.as_ref()?;
let instance = seg.instances().get(index)?;
let (signature, level) = (seg.signature(), seg.level());
let (name, ids) = (instance.class_name.to_string(), instance.regions.clone());
let id = self.graph.masks().next_id();
let mut layer = MaskLayer::new(
id.clone(),
MaskSource::Regions {
signature,
level,
ids,
},
);
layer.name = name;
if !self.graph.masks_mut().push(layer) {
return None;
}
self.active_mask = Some(id.clone());
self.history.record(&self.graph, Edit::Discrete);
Some(id)
}
/// Add a gradient layer, which needs no segmentation.
pub fn add_gradient_mask(&mut self, radial: bool) -> Option<String> {
let id = self.graph.masks().next_id();
let source = if radial {
MaskSource::Radial {
centre: (0.5, 0.5),
radii: (0.35, 0.35),
angle: 0.0,
feather: 0.5,
}
} else {
MaskSource::Linear {
centre: (0.5, 0.5),
angle: std::f32::consts::FRAC_PI_2,
width: 0.3,
}
};
if !self.graph.masks_mut().push(MaskLayer::new(id.clone(), source)) {
return None;
}
self.active_mask = Some(id.clone());
self.history.record(&self.graph, Edit::Discrete);
Some(id)
}
pub fn remove_mask(&mut self, id: &str) {
if self.graph.masks_mut().remove(id).is_some() {
if self.active_mask.as_deref() == Some(id) {
self.active_mask = None;
}
self.history.record(&self.graph, Edit::Discrete);
}
}
pub fn set_mask_enabled(&mut self, id: &str, enabled: bool) {
if let Some(layer) = self.graph.masks_mut().get_mut(id) {
layer.enabled = enabled;
self.history.record(&self.graph, Edit::Discrete);
}
}
pub fn set_mask_invert(&mut self, id: &str, invert: bool) {
if let Some(layer) = self.graph.masks_mut().get_mut(id) {
layer.invert = invert;
self.history.record(&self.graph, Edit::Discrete);
}
}
pub fn set_mask_opacity(&mut self, id: &str, opacity: f32) {
if let Some(layer) = self.graph.masks_mut().get_mut(id) {
layer.opacity = opacity.clamp(0.0, 1.0);
// `Op` rather than `Discrete`: opacity is dragged, and a drag is
// one decision however many values it passes through. `Discrete`
// would put every intermediate position on the undo stack.
self.history.record(&self.graph, Edit::Op(OpId("mask-opacity")));
}
}
/// Whether a layer's region ids belong to a segmentation other than the
/// one currently loaded — a mask restored from a sidecar written under
/// different tuning.
pub fn mask_is_stale(&self, id: &str) -> bool {
let Some(layer) = self.graph.masks().get(id) else {
return false;
};
match self.segmentation.as_ref() {
Some(seg) => layer.is_stale(seg.signature()),
// Nothing loaded to compare against. Not stale, just unrenderable
// — the distinction matters because "stale" invites the user to
// recompute the selection and this only needs the segmentation
// running.
None => false,
}
}
fn lookup(&self, op_index: i32, param_index: i32) -> Option<(OpId, ParamId)> {
// Rows are emitted in capability order, so the flat index is the sum
// of preceding parameter counts.
let caps = self.graph.capabilities();
// of preceding parameter counts. Taken from whichever scope `rows`
// last described — the indices the interface is holding are positions
// in *that* list, and reading the global chain while a layer is
// selected would map a slider onto a different operation.
let caps = match self.active_layer() {
Some(layer) => layer.capabilities(),
None => self.graph.capabilities(),
};
let op = caps.get(usize::try_from(op_index).ok()?)?;
let param = op.params.get(usize::try_from(param_index).ok()?)?;
Some((op.id, param.id))
@@ -560,9 +991,16 @@ impl DevelopSession {
let (w, h) = fit(fw, fh, width.max(1), height.max(1));
let shader = self.graph.compose();
// Rasterise the masks first: the shader addresses array slices by
// index, so the array has to describe *this* stack before it is bound.
let masks = self
.rasterise_masks(w, h)
.then(|| self.masks.as_ref().and_then(|p| p.array()))
.flatten();
let texture = self
.adjust
.render(&self.demosaiced, &shader, w, h)
.render_masked(&self.demosaiced, &shader, w, h, masks)
.map_err(|e| e.to_string())?;
// The import is fallible on format and usage only, and both are fixed
+1
View File
@@ -32,6 +32,7 @@ mod library_ui;
mod live_style;
mod net_runtime;
mod presets;
mod segmentation;
mod settings_store;
mod settings_ui;
mod sidecar_cache;
+538
View File
@@ -0,0 +1,538 @@
//! Running the segmentation for the image on screen.
//!
//! Joins the two arms into the one thing the develop view needs: a region map
//! whose boundaries came from the watershed and whose grouping came from the
//! model (docs/segmentation.md §5).
//!
//! # This is a precompute, and it is slow
//!
//! Roughly 70 ms of watershed plus ~470 ms of inference plus a full-resolution
//! readback, so it runs **once per image on a worker** and never on the frame
//! path. Every interaction it enables — click a subject, widen a selection,
//! toggle the overlay — reads its cached output and costs a lookup.
//!
//! The consequence for the interface is that segmentation is a thing the user
//! *starts*, and the view must stay usable while it runs. There is no version
//! of this that is fast enough to do silently when a photograph opens.
use std::sync::Arc;
use dr_gpu::{GpuContext, LabelField, SegmentOptions, SegmentPass};
use dr_pipeline::mask::segmentation_signature;
use dr_segment::prior::{self, Membership, PriorOptions};
use dr_segment::{MergeTree, RegionField};
/// How many regions a click's selection starts from.
///
/// Not the finest level: at the bottom of the ladder a smooth sky bands into
/// strips and a click lands on one of them, which reads as the tool being
/// broken rather than as precision. A few hundred regions over a proxy is
/// roughly "objects and their large parts", which is where a first click
/// should land — and the scroll wheel walks from there in both directions.
pub const DEFAULT_LEVEL: u32 = 300;
/// One image's region map, ready to select from.
pub struct Segmentation {
/// The watershed's partition, with the semantic prior already applied to
/// its edge weights.
field: RegionField,
/// The merge order over `field`, which is the granularity ladder.
tree: MergeTree,
/// What the model found, for "select the subject" and for the layer names
/// a user actually recognises.
instances: Vec<InstanceSummary>,
/// The labels on the GPU, for the mask rasteriser.
///
/// Optional, and not merely for testing: the region map is complete and
/// every selection decision is made on the CPU half above, so an upload
/// that fails costs the *rendering* of a mask and not the ability to
/// compute one. It also lets the selection logic be asserted with no
/// adapter present, which is where the behaviour that matters lives.
labels: Option<LabelField>,
/// Identifies this segmentation, so a stored mask can tell whether its
/// region ids still mean what they meant.
signature: u64,
/// The grouping at the current level: region id -> group id.
grouping: Vec<u32>,
level: u32,
}
/// A detected object, reduced to what the interface needs.
///
/// The full [`dr_segment::Instance`] carries a soft mask the size of the
/// image; keeping several of those alive per photograph is tens of megabytes
/// for data already folded into the region field. The regions it covers are
/// the durable part.
#[derive(Debug, Clone)]
pub struct InstanceSummary {
pub class_name: Arc<str>,
pub score: f32,
/// The regions this instance covers, snapped to watershed boundaries.
pub regions: Vec<u32>,
}
impl Segmentation {
pub fn signature(&self) -> u64 {
self.signature
}
pub fn labels(&self) -> Option<&LabelField> {
self.labels.as_ref()
}
pub fn level(&self) -> u32 {
self.level
}
pub fn region_count(&self) -> usize {
self.field.region_count
}
pub fn instances(&self) -> &[InstanceSummary] {
&self.instances
}
/// Move the granularity ladder, recomputing the grouping.
///
/// 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);
if level == self.level {
return;
}
self.level = level;
self.grouping = self.tree.cut_to(level as usize);
}
/// The group under a point in **normalised image coordinates**.
///
/// Normalised because the caller has a click in a widget that is letter-
/// 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> {
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];
self.grouping.get(region as usize).copied()
}
/// 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)
.filter(|&r| self.grouping.get(r as usize) == Some(&group))
.collect()
}
/// The regions under a click, at the current level.
pub fn regions_at(&self, x: f32, y: f32) -> Vec<u32> {
self.group_at(x, y)
.map(|g| self.regions_in_group(g))
.unwrap_or_default()
}
/// A false-coloured RGBA image of the current grouping.
///
/// 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.
pub fn overlay_rgba(&self) -> (Vec<u8>, u32, u32) {
let (w, h) = (self.field.width, self.field.height);
let mut out = vec![0u8; w * h * 4];
for (p, &region) in self.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;
}
// Boundaries drawn black over the fill. Without them two adjacent
// groups that happen to land on near hues read as one region, which is
// 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]);
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]);
}
}
}
(out, w as u32, h as u32)
}
}
/// A distinct colour per group.
///
/// Golden-angle hue stepping: consecutive ids land far apart on the wheel, so
/// neighbouring regions — which usually *do* have consecutive ids, since the
/// field is compacted in raster order — are visually separated. Deterministic
/// rather than random, so the same region is the same colour every time the
/// overlay is drawn and the eye can track it across a level change.
fn group_colour(group: u32) -> [u8; 3] {
let h = (group as f32 * 137.508) % 360.0;
// Fixed saturation and value: a fully saturated palette is easier to
// separate than one that also varies in lightness, and the image beneath
// supplies all the tonal variation the eye needs.
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 {
pub segment: SegmentOptions,
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.
pub semantic: bool,
}
impl Default for Options {
fn default() -> Self {
Self {
segment: SegmentOptions::default(),
prior: PriorOptions::default(),
semantic: true,
}
}
}
/// Run the watershed, optionally the model, and combine them.
///
/// `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.
pub fn compute(
ctx: &GpuContext,
source: &dr_gpu::DemosaicedImage,
rgb: &[f32],
rgb_width: usize,
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()
}
}
} 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 (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();
let signature = segmentation_signature(
w,
h,
field.region_count as u32,
tuning_hash(options),
);
Ok(Segmentation {
field,
tree,
instances: summaries,
labels,
signature,
grouping,
level: DEFAULT_LEVEL,
})
}
/// Arm B. Loads the model each call rather than holding a session open.
///
/// ~24 ms against the ~470 ms of inference that follows it, and it runs once
/// per image — so caching the session would keep 11 MB of weights resident for
/// the whole life of the app to save five percent of a background task.
fn detect(rgb: &[f32], width: usize, height: usize) -> Result<Vec<dr_segment::Instance>, String> {
let mut model = dr_segment::SemanticModel::embedded().map_err(|e| e.to_string())?;
model
.detect(rgb, width, height, &dr_segment::SemanticOptions::default())
.map_err(|e| e.to_string())
}
/// Arm C: fold the instances into the merge weights, and record what each one
/// selects.
fn apply_prior(
field: RegionField,
instances: &[dr_segment::Instance],
options: &PriorOptions,
) -> (RegionField, Vec<InstanceSummary>) {
let strong: Vec<&dr_segment::Instance> = instances
.iter()
.filter(|i| i.score >= options.confidence)
.collect();
if strong.is_empty() {
return (field, Vec::new());
}
let masks: Vec<&[f32]> = strong.iter().map(|i| i.mask.as_slice()).collect();
let Ok(membership) = Membership::compute(&field, &masks, options) else {
// Only reachable if a mask disagrees with the field's size, which
// would be a bug rather than a user-facing condition. Arm A alone is
// the right degradation.
log::warn!("instance masks do not match the region field; using watershed alone");
return (field, Vec::new());
};
let summaries = strong
.iter()
.map(|i| InstanceSummary {
class_name: i.class_name.clone(),
score: i.score,
regions: prior::regions_for_instance(&field, &i.mask, options),
})
.filter(|s| !s.regions.is_empty())
.collect();
(
prior::apply_semantic_prior(&field, &membership, options),
summaries,
)
}
/// Everything about the run that changes what a region id means.
///
/// Folded into the signature so a mask stored under one tuning is not silently
/// reinterpreted under another. Not a general-purpose hash of `Options` — the
/// prior's strength changes the *grouping* but not the underlying regions, and
/// including it would invalidate masks that are still perfectly valid.
fn tuning_hash(options: &Options) -> u64 {
let s = &options.segment;
let mut h: u64 = 0xcbf2_9ce4_8422_2325;
for word in [
s.max_edge as u64,
s.blur_radius as u64,
s.w_luma.to_bits() as u64,
s.w_chroma.to_bits() as u64,
s.plateau_iterations as u64,
u64::from(options.semantic),
] {
for byte in word.to_le_bytes() {
h ^= byte as u64;
h = h.wrapping_mul(0x1000_0000_01b3);
}
}
h
}
#[cfg(test)]
mod tests {
use super::*;
/// A 4x2 field split down the middle.
///
/// Regions 0 (top-left) and 1 (bottom-left) make the left half; 2 and 3
/// the right. The seam between the halves is the strong boundary and the
/// within-half seams are weak, so coarsening merges each half before it
/// ever joins the two — which is the structure a real picture of a subject
/// against a background has, and the one every assertion below depends on.
fn field() -> RegionField {
RegionField {
width: 4,
height: 2,
labels: vec![0, 0, 2, 2, 1, 1, 3, 3],
region_count: 4,
adjacency: vec![
dr_segment::Edge { a: 0, b: 1, saddle: 0.1 },
dr_segment::Edge { a: 2, b: 3, saddle: 0.1 },
dr_segment::Edge { a: 0, b: 2, saddle: 1.0 },
dr_segment::Edge { a: 1, b: 3, saddle: 1.0 },
],
}
}
/// A segmentation with no GPU field, which is all the selection logic
/// needs — every decision below is made on the CPU half.
fn segmentation(level: u32) -> Segmentation {
let field = field();
let tree = MergeTree::build(&field);
let grouping = tree.cut_to(level as usize);
Segmentation {
field,
tree,
instances: Vec::new(),
labels: None,
signature: 1,
grouping,
level,
}
}
#[test]
fn a_click_lands_on_the_group_under_it() {
// Four regions cut to two groups: the field's equal saddles merge in
// (a, b) order, so 0-1 and 2-3 pair up — left half and right half.
let seg = segmentation(2);
let left = seg.group_at(0.1, 0.5).expect("inside the frame");
let right = seg.group_at(0.9, 0.5).expect("inside the frame");
assert_ne!(left, right, "the two halves are different groups");
// The same group whichever row is clicked, which is what "region"
// means and what a per-pixel lookup would not give.
assert_eq!(seg.group_at(0.1, 0.1), Some(left));
assert_eq!(seg.group_at(0.1, 0.9), Some(left));
}
#[test]
fn a_click_outside_the_frame_selects_nothing() {
let seg = segmentation(2);
assert_eq!(seg.group_at(-0.1, 0.5), None);
assert_eq!(seg.group_at(1.5, 0.5), None);
assert_eq!(seg.group_at(0.5, 1.0), None, "the far edge is exclusive");
}
#[test]
fn selecting_a_group_returns_every_region_in_it() {
let seg = segmentation(2);
let regions = seg.regions_at(0.1, 0.5);
assert_eq!(regions.len(), 2, "each half is two merged regions");
assert!(regions.windows(2).all(|w| w[0] < w[1]), "sorted for the mask");
}
/// The scroll wheel. Coarsening must not re-segment, and must actually
/// change what a click selects.
#[test]
fn the_level_changes_what_a_click_selects() {
let mut seg = segmentation(4);
assert_eq!(seg.regions_at(0.1, 0.5).len(), 1, "finest: one region");
seg.set_level(2);
assert_eq!(seg.level(), 2);
assert_eq!(seg.regions_at(0.1, 0.5).len(), 2, "coarser: the pair");
}
#[test]
fn the_level_is_clamped_to_what_exists() {
let mut seg = segmentation(4);
seg.set_level(9999);
assert!(seg.level() <= seg.region_count() as u32);
seg.set_level(0);
assert!(seg.level() >= 2, "one group would select the whole frame");
}
#[test]
fn the_overlay_covers_every_pixel_and_draws_boundaries() {
let seg = segmentation(2);
let (rgba, w, h) = seg.overlay_rgba();
assert_eq!((w, h), (4, 2));
assert_eq!(rgba.len(), 4 * 2 * 4);
assert!(rgba.chunks_exact(4).all(|p| p[3] == 255), "fully opaque");
// The vertical seam between the halves must be drawn.
let at = |x: usize, y: usize| {
let p = (y * 4 + x) * 4;
[rgba[p], rgba[p + 1], rgba[p + 2]]
};
assert_eq!(at(1, 0), [0, 0, 0], "the boundary column is black");
assert_ne!(at(0, 0), [0, 0, 0], "the interior is not");
}
#[test]
fn the_two_halves_get_different_colours() {
let seg = segmentation(2);
let (rgba, _, _) = seg.overlay_rgba();
let left = [rgba[0], rgba[1], rgba[2]];
let right = {
let p = 3 * 4;
[rgba[p], rgba[p + 1], rgba[p + 2]]
};
assert_ne!(left, right);
}
#[test]
fn colours_are_stable_and_distinct() {
assert_eq!(group_colour(7), group_colour(7), "same id, same colour");
assert_ne!(group_colour(0), group_colour(1));
assert_ne!(group_colour(1), group_colour(2));
}
#[test]
fn every_colour_is_visible_against_the_image() {
// The overlay is composited over a photograph, so a colour that is
// nearly black is a region the user cannot see they have selected.
for id in 0..64u32 {
let [r, g, b] = group_colour(id);
let brightest = r.max(g).max(b);
assert!(brightest >= 200, "group {id} is too dark: {r},{g},{b}");
}
}
#[test]
fn tuning_changes_the_signature_but_the_prior_does_not() {
let base = Options::default();
let mut coarser = base;
coarser.segment.max_edge = 800;
assert_ne!(tuning_hash(&base), tuning_hash(&coarser), "proxy size");
let mut no_model = base;
no_model.semantic = false;
assert_ne!(tuning_hash(&base), tuning_hash(&no_model), "arm B on/off");
// The prior reweights the merge order over the *same* regions, so a
// mask's ids still mean what they meant. Invalidating them here would
// discard work for no reason.
let mut stronger = base;
stronger.prior.strength = 0.1;
assert_eq!(tuning_hash(&base), tuning_hash(&stronger));
}
}