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

Put the region map behind the sliders that were already there

Browse Source 
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:
dtourolle
2026-08-22 08:39:17 +02:00
parent 94cfea4748
commit 5ecb35864f
8 changed files with 1080 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ui/dr-ui/src/segmentation.rs
+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));
}
}