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:
@@ -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.
|
||||
|
||||
|
||||
@@ -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 `<prefix>-original.ppm`, `<prefix>-colour-pop.ppm`,
|
||||
//! `<prefix>-subject-lift.ppm` and `<prefix>-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 <photo.CR2|photo.RAF> [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<f32> = 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<u8> = 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<u8> = 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<u8> = 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");
|
||||
}
|
||||
@@ -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.
|
||||
|
||||
+151
-9
@@ -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, ¶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");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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; }
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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<f32>, layer: i32) -> f32 {
|
||||
let dims = vec2<f32>(textureDimensions(masks));
|
||||
let last = vec2<i32>(dims) - vec2<i32>(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<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(masks, vec2<i32>(p0.x, p0.y), layer, 0).r;
|
||||
let b = textureLoad(masks, vec2<i32>(p1.x, p0.y), layer, 0).r;
|
||||
let c = textureLoad(masks, vec2<i32>(p0.x, p1.y), layer, 0).r;
|
||||
let d = textureLoad(masks, vec2<i32>(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<u32>,
|
||||
},
|
||||
|
||||
/// 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<i32>(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<i32>(gid.xy), 0, 0)"));
|
||||
assert!(shader.body.contains("textureLoad(masks, vec2<i32>(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]
|
||||
|
||||
@@ -689,6 +689,17 @@ fn write_mask(out: &mut String, version: &str, layer: &MaskLayer) {
|
||||
let list: Vec<String> = 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<u32>,
|
||||
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,
|
||||
|
||||
Reference in New Issue
Block a user