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.
This commit is contained in:
2026-08-22 08:39:17 +02:00
parent b1433ad4a9
commit ee10097435
13 changed files with 963 additions and 101 deletions
+151 -9
View File
@@ -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");
}
}