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

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/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.

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<wgpu::TextureView>,
+    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<Self, GpuError> {
+        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, &params, field, &selected);
+            self.draw(&mut encoder, slot as u32, &params, 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");
     }
 }

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<storage, read> selected: array<u32>;
+// 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<f32>;
 
 // 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>) -> 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>) -> f32 {
+    let dims = vec2<f32>(textureDimensions(subject));
+    let last = vec2<i32>(dims) - vec2<i32>(1);
+
+    let t = uv * dims - vec2<f32>(0.5);
+    let base = vec2<i32>(floor(t));
+    let f = fract(t);
+
+    let p0 = clamp(base, vec2<i32>(0), last);
+    let p1 = clamp(base + vec2<i32>(1), vec2<i32>(0), last);
+
+    let a = textureLoad(subject, vec2<i32>(p0.x, p0.y), 0).r;
+    let b = textureLoad(subject, vec2<i32>(p1.x, p0.y), 0).r;
+    let c = textureLoad(subject, vec2<i32>(p0.x, p1.y), 0).r;
+    let d = textureLoad(subject, vec2<i32>(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<f32>) -> @location(0) vec4<f32> {
     let px = vec2<i32>(i32(pos.x), i32(pos.y));
@@ -151,6 +188,7 @@ fn fs(@builtin(position) pos: vec4<f32>) -> @location(0) vec4<f32> {
         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; }
     }