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;