From ee1009743504f2aedc755d677980f7ed27ec0516 Mon Sep 17 00:00:00 2001 From: Duncan Tourolle Date: Fri, 21 Aug 2026 23:13:20 +0200 Subject: [PATCH] Mask the subject the model found, not the regions underneath it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- core/dr-gpu/Cargo.toml | 11 + core/dr-gpu/examples/local.rs | 265 ++++++++++++++++++++ core/dr-gpu/src/lib.rs | 2 +- core/dr-gpu/src/mask.rs | 160 +++++++++++- core/dr-gpu/src/shaders/mask.wgsl | 40 ++- core/dr-gpu/tests/local_adjustments.rs | 2 +- core/dr-pipeline/src/mask.rs | 103 +++++++- core/dr-pipeline/src/sidecar.rs | 26 ++ ui/dr-ui/src/develop.rs | 81 +++++- ui/dr-ui/src/masks_ui.rs | 3 + ui/dr-ui/src/segmentation.rs | 334 ++++++++++++++++++++----- ui/dr-ui/ui/app.slint | 2 + ui/dr-ui/ui/masks.slint | 35 ++- 13 files changed, 963 insertions(+), 101 deletions(-) create mode 100644 core/dr-gpu/examples/local.rs diff --git a/core/dr-gpu/Cargo.toml b/core/dr-gpu/Cargo.toml index bcbf9a2..143614b 100644 --- a/core/dr-gpu/Cargo.toml +++ b/core/dr-gpu/Cargo.toml @@ -25,6 +25,10 @@ pollster.workspace = true [dev-dependencies] env_logger.workspace = true +# The local-adjustment example needs the model, which the library half of this +# crate deliberately does not: `dr-gpu` holds the shaders, and the inference +# runtime belongs to whoever is asking a question about the picture. +dr-segment = { workspace = true, features = ["semantic", "embedded-model"] } [[example]] name = "bench" @@ -49,3 +53,10 @@ segment-readback = [] [[example]] name = "segment" required-features = ["segment-readback"] + +[[example]] +name = "local" +# No `segment-readback`: this reads the *rendered* proxy back through +# `export_pixels` to feed the model, which is the ungated export path. The +# watershed, and the region-graph transfer that needs the gate, is not involved. + diff --git a/core/dr-gpu/examples/local.rs b/core/dr-gpu/examples/local.rs new file mode 100644 index 0000000..89323b3 --- /dev/null +++ b/core/dr-gpu/examples/local.rs @@ -0,0 +1,265 @@ +//! Local adjustments end to end, on a real photograph. +//! +//! Two edits a photographer actually makes, both driven by the model finding +//! the subject rather than by anyone drawing a shape: +//! +//! - **The subject in colour, everything else monochrome.** One layer, the +//! subject's mask inverted, saturation at −100. +//! - **The subject lifted out of its background.** Two layers over the same +//! mask: the subject brightened, the background pulled down. +//! +//! ```sh +//! cargo run -p dr-gpu --example local --release \ +//! --features segment-readback -- photo.CR2 out +//! ``` +//! +//! Writes `-original.ppm`, `-colour-pop.ppm`, +//! `-subject-lift.ppm` and `-mask.ppm`. PPM for the reason +//! every other example here uses it: no encoder dependency, and every viewer +//! reads it. +//! +//! # What this is really testing +//! +//! That the whole chain agrees with itself. The mask is rasterised in *source* +//! space at proxy resolution and sampled by the composed shader after the +//! framing map, so a fault anywhere in that handoff — a transposed axis, a +//! mask pinned to the viewport, a slice read from the wrong layer — shows up +//! here as an adjustment in the wrong place, and nowhere else. + +use dr_gpu::{ + AdjustPass, DemosaicedImage, Demosaicer, GpuContext, MaskPass, SubjectMasks, +}; +use dr_pipeline::descriptor::ParamId; +use dr_pipeline::mask::{MaskLayer, MaskSource, MaskStack}; +use dr_pipeline::operation::compose_full; +use dr_pipeline::{ops, EditGraph, Framing}; +use dr_segment::{SemanticModel, SemanticOptions}; +use dr_types::ColourSpace; + +/// Longest edge the mask and the model work at. +const PROXY: u32 = 1600; +/// Longest edge of the written frames. +const OUT: u32 = 1400; + +fn main() { + env_logger::init(); + + let mut args = std::env::args().skip(1); + let Some(path) = args.next() else { + eprintln!("usage: local [out-prefix]"); + std::process::exit(2); + }; + let prefix = args.next().unwrap_or_else(|| "local".into()); + + let ctx = pollster::block_on(GpuContext::new_headless()).expect("gpu context"); + println!("gpu {}", ctx.adapter_name()); + + // ---- the photograph --------------------------------------------------- + let bytes = std::fs::read(&path).expect("read file"); + let raw = dr_decode::decode(&bytes).expect("decode"); + println!("source {} × {}", raw.crop.width, raw.crop.height); + let source = Demosaicer::new(&ctx) + .expect("demosaicer") + .run(&raw) + .expect("demosaic"); + + // ---- what the model sees ---------------------------------------------- + // + // The *unedited* image, so the detection does not shift when the edit + // does. Through `export_pixels`, which is ungated: an export is not the + // display round-trip AC-8 forbids, and neither is this. + let (sw, sh) = source.size(); + let scale = (PROXY as f32 / sw.max(sh) as f32).min(1.0); + let (pw, ph) = ( + ((sw as f32 * scale) as u32).max(1), + ((sh as f32 * scale) as u32).max(1), + ); + + let neutral = EditGraph::default_chain(); + let mut proxy_pass = AdjustPass::new(&ctx); + proxy_pass + .render(&source, &neutral.compose(), pw, ph) + .expect("proxy render"); + let (rgba, pw, ph) = proxy_pass.export_pixels().expect("proxy readback"); + println!("proxy {pw} × {ph}"); + + let rgb: Vec = 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(); + + // ---- find the subject ------------------------------------------------- + let t = std::time::Instant::now(); + let mut model = SemanticModel::embedded().expect("model"); + let instances = model + .detect(&rgb, pw as usize, ph as usize, &SemanticOptions::default()) + .expect("detect"); + println!( + "detect {} found in {:.0} ms", + instances.len(), + t.elapsed().as_secs_f32() * 1000.0 + ); + for (i, inst) in instances.iter().enumerate() { + println!(" [{i}] {:<14} {:.2}", inst.class_name, inst.score); + } + + let Some((index, subject)) = pick_subject(&instances) else { + eprintln!("\nNothing recognised in this frame — nothing to adjust locally."); + eprintln!("The model knows COCO's 80 classes; a landscape with no person,"); + eprintln!("animal or vehicle in it has no subject for it to find."); + std::process::exit(1); + }; + println!( + "subject [{index}] {} at {:.2}", + subject.class_name, subject.score + ); + + // Quantised exactly as the develop session does, so this example exercises + // the shipping path rather than a shortcut around it. + let alpha: Vec = subject + .mask + .iter() + .map(|&v| (v.clamp(0.0, 1.0) * 255.0).round() as u8) + .collect(); + let subjects = SubjectMasks::upload(&ctx, &[&alpha], pw, ph).expect("upload subject"); + + let (ow, oh) = fit(sw, sh, OUT); + let mut masks = MaskPass::new(&ctx).expect("mask pass"); + let mut adjust = AdjustPass::new(&ctx); + + // ---- the original, for comparison ------------------------------------- + adjust + .render(&source, &neutral.compose(), ow, oh) + .expect("render"); + write(&format!("{prefix}-original.ppm"), &adjust); + + // ---- 1. the subject in colour, the rest monochrome -------------------- + // + // One layer, inverted. Inverting rather than making a second mask for the + // background is the whole point of having one: there is exactly one + // boundary, so there is exactly one thing to get right. + let mut pop = MaskStack::new(); + let mut drain = subject_layer("m1", index, subject); + drain.invert = true; + drain.set_param("saturation", ParamId("saturation"), -100.0); + // A touch of feather, or the colour stops dead on the model's outline and + // the eye goes straight to the edge instead of to the subject. + drain.feather = 0.02; + pop.push(drain); + + render_stack(&ctx, &source, &mut masks, &mut adjust, &pop, &subjects, pw, ph, ow, oh); + write(&format!("{prefix}-colour-pop.ppm"), &adjust); + + // ---- 2. lift the subject out of its background ------------------------ + let mut lift = MaskStack::new(); + + let mut brighter = subject_layer("m1", index, subject); + brighter.set_param("exposure", ParamId("exposure"), 0.45); + brighter.feather = 0.015; + lift.push(brighter); + + let mut darker = subject_layer("m2", index, subject); + darker.invert = true; + darker.set_param("exposure", ParamId("exposure"), -0.55); + darker.set_param("saturation", ParamId("saturation"), -25.0); + darker.feather = 0.03; + lift.push(darker); + + render_stack(&ctx, &source, &mut masks, &mut adjust, &lift, &subjects, pw, ph, ow, oh); + write(&format!("{prefix}-subject-lift.ppm"), &adjust); + + // ---- the mask itself, to check the outline ---------------------------- + write_mask(&format!("{prefix}-mask.ppm"), &alpha, pw, ph); + + println!("\nwrote {prefix}-original.ppm"); + println!(" {prefix}-colour-pop.ppm"); + println!(" {prefix}-subject-lift.ppm"); + println!(" {prefix}-mask.ppm"); +} + +/// A layer masked to one detected object. +fn subject_layer(id: &str, index: usize, subject: &dr_segment::Instance) -> MaskLayer { + let mut layer = MaskLayer::new( + id, + MaskSource::Subject { + // One segmentation in this process, so any signature agrees with + // itself; the session computes a real one. + signature: 0, + index: index as u32, + class: subject.class_name.to_string(), + score: subject.score, + }, + ); + layer.name = subject.class_name.to_string(); + layer +} + +/// The most promising thing to adjust. +/// +/// Prefers a person, then falls back to the strongest detection of anything. +/// Not because people are special to the pipeline, but because they are what a +/// local adjustment is usually *for*, and an example that picks the parked car +/// behind the subject demonstrates the mechanism while missing the point. +fn pick_subject(instances: &[dr_segment::Instance]) -> Option<(usize, &dr_segment::Instance)> { + instances + .iter() + .enumerate() + .find(|(_, i)| &*i.class_name == "person") + .or_else(|| instances.iter().enumerate().next()) +} + +#[allow(clippy::too_many_arguments)] +fn render_stack( + ctx: &GpuContext, + source: &DemosaicedImage, + masks: &mut MaskPass, + adjust: &mut AdjustPass, + stack: &MaskStack, + subjects: &SubjectMasks, + pw: u32, + ph: u32, + ow: u32, + oh: u32, +) { + let _ = ctx; + // Rasterised at *proxy* size in source space, then sampled by the shader + // after the framing map — which is what makes one mask correct at every + // output size, zoom and crop. + let array = masks + .render(stack, None, Some(subjects), pw, ph) + .expect("rasterise masks"); + + let shader = compose_full(&ops::chain(), &Framing::new(), ColourSpace::Srgb, stack); + adjust + .render_masked(source, &shader, ow, oh, Some(array)) + .expect("render"); +} + +fn fit(w: u32, h: u32, longest: u32) -> (u32, u32) { + let s = (longest as f32 / w.max(h) as f32).min(1.0); + (((w as f32 * s) as u32).max(1), ((h as f32 * s) as u32).max(1)) +} + +fn write(path: &str, adjust: &AdjustPass) { + let (rgba, w, h) = adjust.export_pixels().expect("readback"); + let rgb: Vec = rgba.chunks_exact(4).flat_map(|p| [p[0], p[1], p[2]]).collect(); + write_ppm(path, &rgb, w, h); +} + +fn write_mask(path: &str, alpha: &[u8], w: u32, h: u32) { + let rgb: Vec = alpha.iter().flat_map(|&a| [a, a, a]).collect(); + write_ppm(path, &rgb, w, h); +} + +fn write_ppm(path: &str, rgb: &[u8], w: u32, h: u32) { + use std::io::Write as _; + let mut f = std::io::BufWriter::new(std::fs::File::create(path).expect("create")); + write!(f, "P6\n{w} {h}\n255\n").expect("header"); + f.write_all(rgb).expect("body"); +} diff --git a/core/dr-gpu/src/lib.rs b/core/dr-gpu/src/lib.rs index ee5d298..9c59e9a 100644 --- a/core/dr-gpu/src/lib.rs +++ b/core/dr-gpu/src/lib.rs @@ -30,7 +30,7 @@ pub use error::GpuError; // Renamed on the way out: `BINS` says enough inside `histogram`, and nothing // at all at a crate root shared with demosaic and segmentation. pub use histogram::{Histogram, HistogramPass, BINS as HISTOGRAM_BINS}; -pub use mask::{LabelField, MaskArray, MaskPass}; +pub use mask::{LabelField, MaskArray, MaskPass, SubjectMasks}; pub use segment::{SegmentOptions, SegmentPass, Segmentation}; /// Owns the wgpu device and queue. diff --git a/core/dr-gpu/src/mask.rs b/core/dr-gpu/src/mask.rs index b06e8f5..b338666 100644 --- a/core/dr-gpu/src/mask.rs +++ b/core/dr-gpu/src/mask.rs @@ -30,6 +30,7 @@ use crate::{GpuContext, GpuError}; const MODE_REGIONS: u32 = 0; const MODE_LINEAR: u32 = 1; const MODE_RADIAL: u32 = 2; +const MODE_SUBJECT: u32 = 3; #[repr(C)] #[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)] @@ -109,6 +110,86 @@ impl LabelField { } } +/// The recognised objects' coverage, resident on the GPU. +/// +/// Uploaded once per segmentation, indexed exactly as the detection list is, +/// so a layer storing "instance 3" finds instance 3 here. +pub struct SubjectMasks { + views: Vec, + width: u32, + height: u32, +} + +impl SubjectMasks { + /// Upload one `R8Unorm` texture per instance. + /// + /// A byte per pixel, which is what the model's coverage was quantised to + /// on the way out of inference: 256 levels is finer than any edge a person + /// can see, and four bytes would make a handful of objects most of a + /// hundred megabytes for one photograph. + pub fn upload( + ctx: &GpuContext, + masks: &[&[u8]], + width: u32, + height: u32, + ) -> Result { + let expected = (width * height) as usize; + let mut views = Vec::with_capacity(masks.len()); + + for (i, mask) in masks.iter().enumerate() { + if mask.len() != expected { + return Err(GpuError::InvalidMask(format!( + "subject {i} mask is {} bytes, expected {width}x{height}", + mask.len() + ))); + } + + let texture = ctx.device.create_texture_with_data( + &ctx.queue, + &wgpu::TextureDescriptor { + label: Some("subject-mask"), + size: wgpu::Extent3d { + width, + height, + depth_or_array_layers: 1, + }, + mip_level_count: 1, + sample_count: 1, + dimension: wgpu::TextureDimension::D2, + format: wgpu::TextureFormat::R8Unorm, + usage: wgpu::TextureUsages::TEXTURE_BINDING, + view_formats: &[], + }, + wgpu::util::TextureDataOrder::LayerMajor, + mask, + ); + views.push(texture.create_view(&wgpu::TextureViewDescriptor::default())); + } + + Ok(Self { + views, + width, + height, + }) + } + + pub fn len(&self) -> usize { + self.views.len() + } + + pub fn is_empty(&self) -> bool { + self.views.is_empty() + } + + pub fn size(&self) -> (u32, u32) { + (self.width, self.height) + } + + fn view(&self, index: usize) -> Option<&wgpu::TextureView> { + self.views.get(index) + } +} + /// The rasterised masks for one edit. pub struct MaskArray { texture: wgpu::Texture, @@ -152,6 +233,8 @@ pub struct MaskPass { /// nowhere in the output, so the cheap reuse path is worth a test that /// can actually see it. allocations: usize, + /// Bound at the subject slot for any layer that is not a subject. + empty_subject: SubjectMasks, /// A one-region, always-unselected field, for a stack with no region mask. /// /// The shader's bindings are fixed, so *something* must be bound at the @@ -175,7 +258,21 @@ impl MaskPass { .device .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { label: Some("mask-bgl"), - entries: &[uniform_entry(0), storage_entry(1), storage_entry(2)], + entries: &[ + uniform_entry(0), + storage_entry(1), + storage_entry(2), + wgpu::BindGroupLayoutEntry { + binding: 3, + visibility: wgpu::ShaderStages::FRAGMENT, + ty: wgpu::BindingType::Texture { + sample_type: wgpu::TextureSampleType::Float { filterable: true }, + view_dimension: wgpu::TextureViewDimension::D2, + multisampled: false, + }, + count: None, + }, + ], }); let pipeline_layout = ctx @@ -215,6 +312,7 @@ impl MaskPass { } let placeholder = LabelField::upload(ctx, &[0], 1, 1, 0)?; + let empty_subject = SubjectMasks::upload(ctx, &[&[0u8][..]], 1, 1)?; Ok(Self { ctx: ctx.clone(), @@ -223,6 +321,7 @@ impl MaskPass { array: None, allocations: 0, placeholder, + empty_subject, }) } @@ -235,6 +334,7 @@ impl MaskPass { &mut self, stack: &MaskStack, labels: Option<&LabelField>, + subjects: Option<&SubjectMasks>, width: u32, height: u32, ) -> Result<&MaskArray, GpuError> { @@ -263,9 +363,30 @@ impl MaskPass { (_, _) => &self.placeholder, }; + // A subject layer whose instance is missing is skipped for the + // same reason a region layer without a segmentation is: an absent + // mask that defaults to "everything" would apply the adjustment to + // the whole photograph, which is a much louder failure than none. + let subject = match &layer.source { + MaskSource::Subject { index, .. } => { + match subjects.filter(|s| (*index as usize) < s.len()) { + Some(s) => (s, *index as usize), + None => { + log::warn!( + "mask layer {} names subject {index}, which this segmentation \ + does not have; skipping", + layer.id + ); + continue; + } + } + } + _ => (&self.empty_subject, 0), + }; + let params = self.params(layer, field, width, height); let selected = self.selection_buffer(layer, field); - self.draw(&mut encoder, slot as u32, ¶ms, field, &selected); + self.draw(&mut encoder, slot as u32, ¶ms, field, &selected, subject); } self.ctx.queue.submit([encoder.finish()]); @@ -306,6 +427,15 @@ impl MaskPass { }; match &layer.source { + // `softness` carries the layer's feather. The model's coverage is + // already a soft sigmoid, so zero means "use the edge the model + // drew" rather than "hard edge" — the one place in this shader + // where zero softness is not a step. + MaskSource::Subject { .. } => MaskParams { + mode: MODE_SUBJECT, + softness: layer.feather.clamp(0.0, 0.5), + ..base + }, MaskSource::Regions { .. } => MaskParams { // A pixel of softening at the proxy-to-output ratio, so the // edge is equally soft whatever size the render is. @@ -371,6 +501,7 @@ impl MaskPass { params: &MaskParams, field: &LabelField, selected: &wgpu::Buffer, + subject: (&SubjectMasks, usize), ) { let params_buf = self .ctx @@ -400,6 +531,17 @@ impl MaskPass { binding: 2, resource: selected.as_entire_binding(), }, + wgpu::BindGroupEntry { + binding: 3, + resource: wgpu::BindingResource::TextureView( + subject + .0 + .view(subject.1) + .unwrap_or_else(|| { + self.empty_subject.view(0).expect("placeholder exists") + }), + ), + }, ], }); @@ -550,7 +692,7 @@ mod tests { })); let mut pass = MaskPass::new(&ctx).expect("mask pass"); - let array = pass.render(&stack, Some(&field), w, h).expect("render"); + let array = pass.render(&stack, Some(&field), None, w, h).expect("render"); assert_eq!(array.size(), (w, h)); assert_eq!(array.layers(), 1); } @@ -571,7 +713,7 @@ mod tests { })); let mut pass = MaskPass::new(&ctx).expect("mask pass"); - assert!(pass.render(&stack, None, 8, 8).is_ok()); + assert!(pass.render(&stack, None, None, 8, 8).is_ok()); } #[test] @@ -594,7 +736,7 @@ mod tests { })); let mut pass = MaskPass::new(&ctx).expect("mask pass"); - let array = pass.render(&stack, None, 16, 16).expect("render"); + let array = pass.render(&stack, None, None, 16, 16).expect("render"); assert_eq!(array.layers(), 2, "one slice per active layer"); } @@ -606,7 +748,7 @@ mod tests { }; let mut pass = MaskPass::new(&ctx).expect("mask pass"); let array = pass - .render(&MaskStack::new(), None, 8, 8) + .render(&MaskStack::new(), None, None, 8, 8) .expect("render"); assert_eq!( array.layers(), @@ -629,17 +771,17 @@ mod tests { })); let mut pass = MaskPass::new(&ctx).expect("mask pass"); - pass.render(&stack, None, 32, 32).expect("render"); + pass.render(&stack, None, None, 32, 32).expect("render"); assert_eq!(pass.allocations(), 1); - pass.render(&stack, None, 32, 32).expect("render"); + pass.render(&stack, None, None, 32, 32).expect("render"); assert_eq!( pass.allocations(), 1, "same size and layer count should not reallocate" ); - pass.render(&stack, None, 64, 64).expect("render"); + pass.render(&stack, None, None, 64, 64).expect("render"); assert_eq!(pass.allocations(), 2, "a resize must reallocate"); } } diff --git a/core/dr-gpu/src/shaders/mask.wgsl b/core/dr-gpu/src/shaders/mask.wgsl index 5f33c47..422a029 100644 --- a/core/dr-gpu/src/shaders/mask.wgsl +++ b/core/dr-gpu/src/shaders/mask.wgsl @@ -28,7 +28,7 @@ struct MaskParams { label_width: u32, label_height: u32, - // 0 = regions, 1 = linear, 2 = radial. + // 0 = regions, 1 = linear, 2 = radial, 3 = subject. mode: u32, // How many regions the label field holds, so an out-of-range label is // caught rather than read past the end of `selected`. @@ -57,6 +57,10 @@ struct MaskParams { // One entry per region: non-zero if the region is in this mask. Small — a few // thousand bytes — which is what makes changing a selection cheap. @group(0) @binding(2) var selected: array; +// One recognised object's coverage, for a subject mask. A 1x1 placeholder +// when the layer is not one — the binding is fixed, and a second pipeline +// differing only in what it ignores would be worse than a wasted texel. +@group(0) @binding(3) var subject: texture_2d; // A full-screen triangle rather than a quad: three vertices instead of six, // no shared edge for the rasteriser to crack along, and no vertex buffer. @@ -138,6 +142,39 @@ fn radial_mask(uv: vec2) -> f32 { return 1.0 - smoothstep(1.0 - edge, 1.0, r); } +// The model's coverage for one object, resampled to the mask's own grid. +// +// Bilinear, unlike the region lookup above: this is a *quantity*, not a name, +// so the value between two samples is meaningful. The model's own mask is a +// quarter-resolution sigmoid, and interpolating it is what stops the outline +// stair-stepping in blocks of four. +fn subject_mask(uv: vec2) -> f32 { + let dims = vec2(textureDimensions(subject)); + let last = vec2(dims) - vec2(1); + + let t = uv * dims - vec2(0.5); + let base = vec2(floor(t)); + let f = fract(t); + + let p0 = clamp(base, vec2(0), last); + let p1 = clamp(base + vec2(1), vec2(0), last); + + let a = textureLoad(subject, vec2(p0.x, p0.y), 0).r; + let b = textureLoad(subject, vec2(p1.x, p0.y), 0).r; + let c = textureLoad(subject, vec2(p0.x, p1.y), 0).r; + let d = textureLoad(subject, vec2(p1.x, p1.y), 0).r; + + let cov = mix(mix(a, b, f.x), mix(c, d, f.x), f.y); + + // `softness` carries the layer's feather here. Zero gives the model's own + // soft edge untouched, which is a perfectly good mask edge and a better + // default than imposing a ramp on top of one that already exists. + if (p.softness <= 0.0) { + return cov; + } + return smoothstep(0.5 - p.softness, 0.5 + p.softness, cov); +} + @fragment fn fs(@builtin(position) pos: vec4) -> @location(0) vec4 { let px = vec2(i32(pos.x), i32(pos.y)); @@ -151,6 +188,7 @@ fn fs(@builtin(position) pos: vec4) -> @location(0) vec4 { case 0u: { m = region_mask(px); } case 1u: { m = linear_mask(uv); } case 2u: { m = radial_mask(uv); } + case 3u: { m = subject_mask(uv); } default: { m = 0.0; } } diff --git a/core/dr-gpu/tests/local_adjustments.rs b/core/dr-gpu/tests/local_adjustments.rs index 52db997..8a930c8 100644 --- a/core/dr-gpu/tests/local_adjustments.rs +++ b/core/dr-gpu/tests/local_adjustments.rs @@ -59,7 +59,7 @@ fn render(ctx: &GpuContext, stack: &MaskStack, field: Option<&LabelField>) -> Ve ); let mut masks = MaskPass::new(ctx).expect("mask pass"); - let array = masks.render(stack, field, SIZE, SIZE).expect("rasterise"); + let array = masks.render(stack, field, None, SIZE, SIZE).expect("rasterise"); let mut adjust = AdjustPass::new(ctx); adjust diff --git a/core/dr-pipeline/src/mask.rs b/core/dr-pipeline/src/mask.rs index 13e76b8..17c7b50 100644 --- a/core/dr-pipeline/src/mask.rs +++ b/core/dr-pipeline/src/mask.rs @@ -44,6 +44,40 @@ use crate::ops; /// be a literal in one. pub const DEFAULT_FEATHER: f32 = 0.004; +/// Bilinear sampling of one slice of the mask array, in **source** space. +/// +/// Hand-rolled rather than done with a sampler, matching how the source +/// texture is read: adding a sampler would change a bind group layout every +/// pass shares. +/// +/// Bilinear rather than a straight load because the array is at proxy +/// resolution and the view may be zoomed well past it. A nearest-neighbour +/// mask shows as visible stair-stepping along the edge of the adjustment at +/// 100%, which is exactly where a mask is judged. +pub(crate) const MASK_SAMPLER: crate::operation::Helper = crate::operation::Helper { + name: "sample_mask", + source: "fn sample_mask(uv: vec2, layer: i32) -> f32 { + let dims = vec2(textureDimensions(masks)); + let last = vec2(dims) - vec2(1); + + // Sample positions are texel centres, so the half-texel offset is what + // keeps the interpolated edge where the rasteriser drew it. + let t = uv * dims - vec2(0.5); + let base = vec2(floor(t)); + let f = fract(t); + + let p0 = clamp(base, vec2(0), last); + let p1 = clamp(base + vec2(1), vec2(0), last); + + let a = textureLoad(masks, vec2(p0.x, p0.y), layer, 0).r; + let b = textureLoad(masks, vec2(p1.x, p0.y), layer, 0).r; + let c = textureLoad(masks, vec2(p0.x, p1.y), layer, 0).r; + let d = textureLoad(masks, vec2(p1.x, p1.y), layer, 0).r; + + return mix(mix(a, b, f.x), mix(c, d, f.x), f.y); +}", +}; + /// Per-layer uniforms the generated shader reads: `invert`, then `opacity`. pub const LAYER_UNIFORM_FIELDS: usize = 2; @@ -215,6 +249,38 @@ pub enum MaskSource { ids: Vec, }, + /// One object the model recognised, used as the mask directly. + /// + /// **The primary way a local adjustment is made.** The watershed hierarchy + /// this crate was first built around does not survive a photograph: its + /// saddles are near zero almost everywhere, so a global cut collapses the + /// frame into one region plus noise (docs/segmentation.md §15). A model + /// instance is a whole object, found as one thing, and needs no ladder. + /// + /// The trade is that the boundary is the model's — a quarter-resolution + /// sigmoid — rather than the image's own gradient. That is what the edge + /// treatment on [`MaskLayer`] is for: the mask arrives approximately + /// right and soft, and dilation, erosion and a chosen falloff are how it + /// is made to fit. + /// + /// Stored as *identity*, not as pixels. The mask itself is several + /// megabytes and is reproducible by running the same model over the same + /// image, so the sidecar carries what is needed to find it again and the + /// session carries the pixels. + Subject { + /// Which segmentation run produced it, so a layer can tell whether + /// the index below still means what it meant. + signature: u64, + /// Position in that run's detection list, strongest first. + index: u32, + /// The class name, for display and as a sanity check on re-detection: + /// if instance 3 is a "car" where it was a "dog", the model or the + /// image changed and the layer should be treated as stale rather than + /// silently masking something else. + class: String, + score: f32, + }, + /// A linear gradient — the graduated-filter mask. /// /// Geometry is in **normalised output coordinates**, so it survives a crop @@ -247,6 +313,7 @@ impl MaskSource { pub fn kind(&self) -> &'static str { match self { Self::Regions { .. } => "regions", + Self::Subject { .. } => "subject", Self::Linear { .. } => "linear", Self::Radial { .. } => "radial", } @@ -405,7 +472,15 @@ impl MaskLayer { /// Applying it anyway would produce a confidently wrong mask, so callers /// should offer to recompute rather than render it. pub fn is_stale(&self, current: u64) -> bool { - matches!(self.source, MaskSource::Regions { signature, .. } if signature != current) + match self.source { + MaskSource::Regions { signature, .. } | MaskSource::Subject { signature, .. } => { + signature != current + } + // A gradient is geometry in normalised coordinates. It means the + // same thing whatever was or was not detected, so nothing about a + // new run can invalidate it. + MaskSource::Linear { .. } | MaskSource::Radial { .. } => false, + } } pub fn descriptors(&self) -> Vec<&'static OpDescriptor> { @@ -599,6 +674,10 @@ pub(crate) fn compose_layers(stack: &MaskStack) -> LayerShader { helpers: Vec::new(), }; + if stack.active().next().is_some() { + out.helpers.push(MASK_SAMPLER); + } + for (slot, layer) in stack.active().enumerate() { let prefix = format!("mask{slot}"); @@ -616,9 +695,21 @@ pub(crate) fn compose_layers(stack: &MaskStack) -> LayerShader { layer.source.kind() ); let _ = writeln!(out.body, " {{"); + // **`uv_src`, not `gid.xy`.** The mask array is rasterised in *source* + // space, and `uv_src` is the source position this output pixel came + // from — after the crop, the zoom, the pan, the straightening and the + // flips. Sampling by output pixel instead, as this once did, pins the + // mask to the viewport: zooming in slides the photograph under a mask + // that stays where it was, and cropping moves the adjustment to a + // different part of the picture. + // + // Doing it this way also means the framing map exists in exactly one + // place. A second copy here would be a second thing to keep in step + // with `Framing::wgsl_prologue`, and the failure would be a mask that + // is subtly wrong only when straightened. let _ = writeln!( out.body, - " var m = textureLoad(masks, vec2(gid.xy), {slot}, 0).r;" + " var m = sample_mask(uv_src, {slot});" ); let _ = writeln!( out.body, @@ -763,8 +854,8 @@ mod tests { stack.push(lit_layer("m2", -1.0)); let shader = compose_layers(&stack); - assert!(shader.body.contains("textureLoad(masks, vec2(gid.xy), 0, 0)")); - assert!(shader.body.contains("textureLoad(masks, vec2(gid.xy), 1, 0)")); + assert!(shader.body.contains("sample_mask(uv_src, 0)")); + assert!(shader.body.contains("sample_mask(uv_src, 1)")); assert!(shader.body.contains("u.mask0_opacity")); assert!(shader.body.contains("u.mask1_opacity")); } @@ -781,10 +872,10 @@ mod tests { let shader = compose_layers(&stack); assert!( - shader.body.contains("gid.xy), 0, 0"), + shader.body.contains("sample_mask(uv_src, 0)"), "the one active layer must use slot 0, not slot 1" ); - assert!(!shader.body.contains("gid.xy), 1, 0")); + assert!(!shader.body.contains("sample_mask(uv_src, 1)")); } #[test] diff --git a/core/dr-pipeline/src/sidecar.rs b/core/dr-pipeline/src/sidecar.rs index 4ea41cc..2f21750 100644 --- a/core/dr-pipeline/src/sidecar.rs +++ b/core/dr-pipeline/src/sidecar.rs @@ -689,6 +689,17 @@ fn write_mask(out: &mut String, version: &str, layer: &MaskLayer) { let list: Vec = ids.iter().map(|i| i.to_string()).collect(); let _ = writeln!(out, "regions = {}", list.join(" ")); } + MaskSource::Subject { + signature, + index, + class, + score, + } => { + let _ = writeln!(out, "signature = {signature}"); + let _ = writeln!(out, "index = {index}"); + let _ = writeln!(out, "class = {class}"); + let _ = writeln!(out, "score = {}", format_value(*score)); + } MaskSource::Linear { centre, angle, @@ -766,6 +777,9 @@ struct PartialMask { signature: u64, level: u32, ids: Vec, + index: u32, + class: String, + score: f32, centre: (f32, f32), radii: (f32, f32), angle: f32, @@ -793,6 +807,9 @@ impl PartialMask { signature: 0, level: 0, ids: Vec::new(), + index: 0, + class: String::new(), + score: 0.0, centre: (0.5, 0.5), radii: (0.25, 0.25), angle: 0.0, @@ -827,6 +844,9 @@ impl PartialMask { self.ids.sort_unstable(); self.ids.dedup(); } + "index" => self.index = value.parse().unwrap_or(0), + "class" => self.class = value.to_string(), + "score" => self.score = value.parse().unwrap_or(0.0), "centre" => self.centre = pair(value).unwrap_or(self.centre), "radii" => self.radii = pair(value).unwrap_or(self.radii), "angle" => self.angle = value.parse().unwrap_or(0.0), @@ -881,6 +901,12 @@ impl PartialMask { level: self.level, ids: self.ids, }, + "subject" => MaskSource::Subject { + signature: self.signature, + index: self.index, + class: self.class, + score: self.score, + }, "linear" => MaskSource::Linear { centre: self.centre, angle: self.angle, diff --git a/ui/dr-ui/src/develop.rs b/ui/dr-ui/src/develop.rs index b20abd4..ffaab24 100644 --- a/ui/dr-ui/src/develop.rs +++ b/ui/dr-ui/src/develop.rs @@ -55,6 +55,8 @@ pub struct DevelopSession { segmentation: Option, /// Rasterises the mask layers. Built lazily for the same reason. masks: Option, + /// The recognised objects' coverage, on the GPU. + subjects: Option, /// Which layer the develop panel is editing, if any. /// /// This is what lets one panel serve both scopes: with a layer selected, @@ -122,6 +124,7 @@ impl DevelopSession { .ok(), segmentation: None, masks: None, + subjects: None, active_mask: None, show_overlay: false, } @@ -601,17 +604,37 @@ impl DevelopSession { /// `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 { + fn rasterise_masks(&mut self) -> bool { if self.graph.masks().is_neutral() { return false; } - let labels = self.segmentation.as_ref().and_then(|s| s.labels()); + // **Source space, at the segmentation's proxy size** — not the + // viewport's. The generated shader samples this after the framing map, + // so a mask drawn here stays on the photograph through a zoom, a pan + // and a crop. Rasterising at viewport size, as this first did, pinned + // the mask to the screen instead: zooming slid the picture underneath + // one that stayed put. + // + // It also means the array does not reallocate when the window + // resizes, and does not need redrawing when the view moves. + let Some(seg) = self.segmentation.as_ref() else { + return false; + }; + let (pw, ph) = seg.proxy_size(); + let labels = seg.labels(); + let subjects = self.subjects.as_ref(); 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() + pass.render( + self.graph.masks(), + labels, + subjects, + pw as u32, + ph as u32, + ) + .inspect_err(|e| log::warn!("mask rasterisation failed: {e}")) + .is_ok() } // ---------------------------------------------------------------------- @@ -637,6 +660,19 @@ impl DevelopSession { .inspect_err(|e| log::warn!("no mask rasteriser on this device: {e}")) .ok(); } + + let (pw, ph) = seg.proxy_size(); + let alphas: Vec<&[u8]> = (0..seg.instances().len()) + .filter_map(|i| seg.instance_mask(i)) + .collect(); + self.subjects = if alphas.is_empty() { + None + } else { + dr_gpu::SubjectMasks::upload(ctx, &alphas, pw as u32, ph as u32) + .inspect_err(|e| log::warn!("could not upload the subject masks: {e}")) + .ok() + }; + self.segmentation = Some(seg); Ok(()) } @@ -688,6 +724,11 @@ impl DevelopSession { self.segmentation.as_ref().map_or(0, |s| s.level()) } + /// Whether a region hierarchy exists to click into. + pub fn has_regions(&self) -> bool { + self.segmentation.as_ref().is_some_and(|s| s.has_regions()) + } + pub fn segmentation_region_count(&self) -> usize { self.segmentation.as_ref().map_or(0, |s| s.region_count()) } @@ -821,6 +862,15 @@ impl DevelopSession { /// /// Returns the layer that now holds the selection. pub fn select_region_at(&mut self, x: f32, y: f32, add: bool) -> Option { + // With no region hierarchy — the default — a click means "the object + // under the cursor", which is the interaction the model can actually + // support. `add` has no meaning for a whole object, so it is ignored + // rather than quietly doing something else. + if self.segmentation.as_ref().is_some_and(|s| !s.has_regions()) { + let index = self.segmentation.as_ref()?.instance_at(x, y)?; + return self.add_subject_mask(index); + } + let seg = self.segmentation.as_ref()?; let picked = seg.regions_at(x, y); if picked.is_empty() { @@ -875,19 +925,28 @@ impl DevelopSession { } /// Add a layer selecting one detected subject. + /// + /// The instance's own coverage is the mask, rather than the watershed + /// regions it overlaps. Snapping to regions was the original design and + /// it is not currently worth doing: the hierarchy those ids index into + /// collapses on a photograph (docs/segmentation.md §15), so snapping + /// would trade the model's approximately-right outline for a + /// confidently-wrong one. pub fn add_subject_mask(&mut self, index: usize) -> Option { 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 signature = seg.signature(); + let name = instance.class_name.to_string(); + let score = instance.score; let id = self.graph.masks().next_id(); let mut layer = MaskLayer::new( id.clone(), - MaskSource::Regions { + MaskSource::Subject { signature, - level, - ids, + index: index as u32, + class: name.clone(), + score, }, ); layer.name = name; @@ -1024,7 +1083,7 @@ impl DevelopSession { // 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) + .rasterise_masks() .then(|| self.masks.as_ref().and_then(|p| p.array())) .flatten(); let texture = self diff --git a/ui/dr-ui/src/masks_ui.rs b/ui/dr-ui/src/masks_ui.rs index d54341e..8eff319 100644 --- a/ui/dr-ui/src/masks_ui.rs +++ b/ui/dr-ui/src/masks_ui.rs @@ -31,6 +31,7 @@ pub(crate) fn sync(window: &AppWindow, session: &Rc::default())); window.set_segmented(false); window.set_segmentation_regions(0); + window.set_has_regions(false); window.set_editing_mask(false); window.set_overlay_on(false); return; @@ -69,6 +70,7 @@ pub(crate) fn sync(window: &AppWindow, session: &Rc::default())); window.set_subject_rows(ModelRc::new(VecModel::::default())); window.set_editing_mask(false); diff --git a/ui/dr-ui/src/segmentation.rs b/ui/dr-ui/src/segmentation.rs index e1ad3ef..445c280 100644 --- a/ui/dr-ui/src/segmentation.rs +++ b/ui/dr-ui/src/segmentation.rs @@ -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, /// The merge order over `field`, which is the granularity ladder. - tree: MergeTree, + tree: Option, /// What the model found, for "select the subject" and for the layer names /// a user actually recognises. instances: Vec, @@ -55,6 +63,13 @@ pub struct Segmentation { /// The grouping at the current level: region id -> group id. grouping: Vec, 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, 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, + /// 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, } 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 { + 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 { + 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 { - (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, 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, u32, u32) { + let (w, h) = self.proxy; let mut out = vec![0u8; w * h * 4]; - for (p, ®ion) 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 = (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, u32, u32) { + let (w, h) = (field.width, field.height); + let mut out = vec![0u8; w * h * 4]; + + for (p, ®ion) 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 { - 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 = 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 { + 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() }; diff --git a/ui/dr-ui/ui/app.slint b/ui/dr-ui/ui/app.slint index 2a35dcb..7cfb336 100644 --- a/ui/dr-ui/ui/app.slint +++ b/ui/dr-ui/ui/app.slint @@ -821,6 +821,7 @@ export component AppWindow inherits Window { in property segmenting: false; in property segmentation-level: 300; in property segmentation-regions: 0; + in property has-regions: false; in property editing-mask: false; callback segment-image(); @@ -1903,6 +1904,7 @@ in property panel-visible: true; segmenting: root.segmenting; level: root.segmentation-level; region-count: root.segmentation-regions; + has-regions: root.has-regions; overlay: root.overlay-on; picking: root.region-picking; editing-mask: root.editing-mask; diff --git a/ui/dr-ui/ui/masks.slint b/ui/dr-ui/ui/masks.slint index b0cdb9b..408ad0e 100644 --- a/ui/dr-ui/ui/masks.slint +++ b/ui/dr-ui/ui/masks.slint @@ -171,6 +171,13 @@ export component MaskPanel inherits Rectangle { /// How many regions the ladder is currently cut to. in property level: 300; in property region-count: 0; + /// Whether a region hierarchy exists to click into. + /// + /// Off in the ordinary case. The watershed's ladder collapses on a + /// photograph, so the controls that drive it are hidden rather than shown + /// doing nothing — a slider that changes no outcome is worse than an + /// absent one, because it invites the user to blame themselves. + in property has-regions: false; /// Draw the false-coloured region map over the photograph. in property overlay: false; @@ -213,7 +220,12 @@ export component MaskPanel inherits Rectangle { HorizontalLayout { PanelHeading { text: "LOCAL"; } Rectangle { horizontal-stretch: 1; } - if root.segmented: Value { text: root.level + " / " + root.region-count; } + if root.segmented && root.has-regions: Value { + text: root.level + " / " + root.region-count; + } + if root.segmented && !root.has-regions: Value { + text: root.subjects.length + (root.subjects.length == 1 ? " subject" : " subjects"); + } } if !root.enabled: Caption { text: "No image"; } @@ -224,7 +236,7 @@ export component MaskPanel inherits Rectangle { // and inference is too long to spend on every photograph opened, and // most edits never need one. if root.enabled && !root.segmented: Caption { - text: "Find the regions in this photograph, so a mask can snap to them."; + text: "Find the subjects in this photograph, so a mask can follow one."; // Wrapped, or this sentence's single-line width becomes the // panel's minimum and levers the fixed 280px develop column open, // taking every other panel's controls off the right edge with it. @@ -232,7 +244,7 @@ export component MaskPanel inherits Rectangle { } if root.enabled && !root.segmented: Button { - text: root.segmenting ? "Finding regions…" : "Find regions"; + text: root.segmenting ? "Looking…" : "Find subjects"; enabled: !root.segmenting; primary: true; clicked => { root.segment(); } @@ -256,11 +268,16 @@ export component MaskPanel inherits Rectangle { } } - if root.enabled && root.segmented && root.picking: Caption { + if root.enabled && root.segmented && root.picking && root.has-regions: Caption { text: "Click the photograph to select a region. Shift-click to add or remove."; wrap: word-wrap; } + if root.enabled && root.segmented && root.picking && !root.has-regions: Caption { + text: "Click a subject in the photograph to mask it."; + wrap: word-wrap; + } + // Granularity, labelled by what it does rather than by its number: // "detail" is what a photographer is choosing between, where "300 // regions" is an implementation detail they would have to learn. @@ -268,7 +285,7 @@ export component MaskPanel inherits Rectangle { // and a photograph that segments into more than that had the finest // part of its own ladder unreachable — the slider simply stopped // before the regions did. - if root.enabled && root.segmented: SliderRow { + if root.enabled && root.segmented && root.has-regions: SliderRow { label: "Detail"; hint: "How finely a click divides the picture, out of " + root.region-count + " regions the watershed found."; @@ -286,7 +303,7 @@ export component MaskPanel inherits Rectangle { // A list of four beside a count of three thousand invites exactly // one question, and the panel should answer it rather than // provoke it. - text: "Subjects the model recognised. Click one to select the whole thing."; + text: "Click one to mask it. The outline is the model's, so soften or grow it below."; wrap: word-wrap; } @@ -311,6 +328,12 @@ export component MaskPanel inherits Rectangle { } } + if root.enabled && root.segmented && root.subjects.length == 0: Caption { + text: "Nothing recognised. The model knows people, animals and vehicles — " + + "a landscape has no subject for it to find. Add a gradient instead."; + wrap: word-wrap; + } + Rectangle { height: 1px; background: Theme.rule;