diff --git a/core/dr-gpu/examples/local.rs b/core/dr-gpu/examples/local.rs index 89323b3..e8f94b2 100644 --- a/core/dr-gpu/examples/local.rs +++ b/core/dr-gpu/examples/local.rs @@ -30,10 +30,10 @@ use dr_gpu::{ AdjustPass, DemosaicedImage, Demosaicer, GpuContext, MaskPass, SubjectMasks, }; use dr_pipeline::descriptor::ParamId; -use dr_pipeline::mask::{MaskLayer, MaskSource, MaskStack}; +use dr_pipeline::mask::{MaskLayer, MaskSource, MaskStack, Morphology}; use dr_pipeline::operation::compose_full; use dr_pipeline::{ops, EditGraph, Framing}; -use dr_segment::{SemanticModel, SemanticOptions}; +use dr_segment::{SemanticModel, SemanticOptions, Shaped}; use dr_types::ColourSpace; /// Longest edge the mask and the model work at. @@ -127,7 +127,6 @@ fn main() { .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"); @@ -153,7 +152,7 @@ fn main() { drain.feather = 0.02; pop.push(drain); - render_stack(&ctx, &source, &mut masks, &mut adjust, &pop, &subjects, pw, ph, ow, oh); + render_stack(&ctx, &source, &mut masks, &mut adjust, &pop, &alpha, pw, ph, ow, oh); write(&format!("{prefix}-colour-pop.ppm"), &adjust); // ---- 2. lift the subject out of its background ------------------------ @@ -171,15 +170,39 @@ fn main() { darker.feather = 0.03; lift.push(darker); - render_stack(&ctx, &source, &mut masks, &mut adjust, &lift, &subjects, pw, ph, ow, oh); + render_stack(&ctx, &source, &mut masks, &mut adjust, &lift, &alpha, pw, ph, ow, oh); write(&format!("{prefix}-subject-lift.ppm"), &adjust); + // ---- 3. the same edit, grown and shrunk ------------------------------- + // + // The model's outline is approximately right and slightly soft, so the + // everyday correction is to move it: grow to catch a halo the detector + // stopped short of, shrink to pull off one it caught. Both are a threshold + // of the distance field, which is why they cost a uniform. + for (name, morphology, radius) in [ + ("grown", Morphology::Dilate, 0.012), + ("shrunk", Morphology::Erode, 0.012), + ] { + let mut stack = MaskStack::new(); + let mut layer = subject_layer("m1", index, subject); + layer.invert = true; + layer.set_param("saturation", ParamId("saturation"), -100.0); + layer.feather = 0.004; + layer.morphology = morphology; + layer.morph_radius = radius; + stack.push(layer); + + render_stack(&ctx, &source, &mut masks, &mut adjust, &stack, &alpha, pw, ph, ow, oh); + write(&format!("{prefix}-{name}.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}-grown.ppm, {prefix}-shrunk.ppm"); println!(" {prefix}-mask.ppm"); } @@ -221,18 +244,44 @@ fn render_stack( masks: &mut MaskPass, adjust: &mut AdjustPass, stack: &MaskStack, - subjects: &SubjectMasks, + coverage: &[u8], pw: u32, ph: u32, ow: u32, oh: u32, ) { - let _ = ctx; + // One signed distance field per active layer, in that order — the order + // the rasteriser indexes them by. Built here rather than once up front + // because a compound morphology rebuilds the field, so it belongs to the + // layer that shaped it rather than to the object. + let fields: Vec> = stack + .active() + .map(|layer| { + Shaped::build( + coverage, + pw as usize, + ph as usize, + 128, + match layer.morphology { + Morphology::None => dr_segment::Morphology::None, + Morphology::Dilate => dr_segment::Morphology::Dilate, + Morphology::Erode => dr_segment::Morphology::Erode, + Morphology::Close => dr_segment::Morphology::Close, + Morphology::Open => dr_segment::Morphology::Open, + }, + layer.morph_radius * pw.min(ph) as f32, + ) + .distance + }) + .collect(); + let refs: Vec<&[f32]> = fields.iter().map(|f| f.as_slice()).collect(); + let subjects = SubjectMasks::upload(ctx, &refs, pw, ph).expect("upload fields"); + // 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) + .render(stack, None, Some(&subjects), pw, ph) .expect("rasterise masks"); let shader = compose_full(&ops::chain(), &Framing::new(), ColourSpace::Srgb, stack); diff --git a/core/dr-gpu/src/mask.rs b/core/dr-gpu/src/mask.rs index b338666..c65f857 100644 --- a/core/dr-gpu/src/mask.rs +++ b/core/dr-gpu/src/mask.rs @@ -43,7 +43,9 @@ struct MaskParams { mode: u32, region_count: u32, feather: f32, - _pad0: f32, + /// Which falloff curve a subject layer uses. Kept in step with the + /// `switch` in `mask.wgsl` by `falloff_code`. + falloff: u32, centre: [f32; 2], axis: [f32; 2], @@ -110,10 +112,17 @@ impl LabelField { } } -/// The recognised objects' coverage, resident on the GPU. +/// Signed distance fields for the subject layers, 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. +/// **One per active layer, in that order** — not one per detected object. Two +/// layers can mask the same subject with different morphology, and closing or +/// opening rebuilds the field rather than offsetting it, so the field belongs +/// to the layer that shaped it. +/// +/// `R32Float`, because the values are signed distances in pixels and the +/// controls read them at sub-pixel precision. That is four bytes a pixel: +/// ~7 MB per layer at a 1600 px proxy, which is the price of making grow, +/// shrink and feather cost nothing per frame. pub struct SubjectMasks { views: Vec, width: u32, @@ -121,33 +130,28 @@ pub struct SubjectMasks { } 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. + /// Upload one distance field per active subject layer. pub fn upload( ctx: &GpuContext, - masks: &[&[u8]], + fields: &[&[f32]], width: u32, height: u32, ) -> Result { let expected = (width * height) as usize; - let mut views = Vec::with_capacity(masks.len()); + let mut views = Vec::with_capacity(fields.len()); - for (i, mask) in masks.iter().enumerate() { - if mask.len() != expected { + for (i, field) in fields.iter().enumerate() { + if field.len() != expected { return Err(GpuError::InvalidMask(format!( - "subject {i} mask is {} bytes, expected {width}x{height}", - mask.len() + "subject field {i} is {} values, expected {width}x{height}", + field.len() ))); } let texture = ctx.device.create_texture_with_data( &ctx.queue, &wgpu::TextureDescriptor { - label: Some("subject-mask"), + label: Some("subject-distance"), size: wgpu::Extent3d { width, height, @@ -156,12 +160,12 @@ impl SubjectMasks { mip_level_count: 1, sample_count: 1, dimension: wgpu::TextureDimension::D2, - format: wgpu::TextureFormat::R8Unorm, + format: wgpu::TextureFormat::R32Float, usage: wgpu::TextureUsages::TEXTURE_BINDING, view_formats: &[], }, wgpu::util::TextureDataOrder::LayerMajor, - mask, + bytemuck::cast_slice(field), ); views.push(texture.create_view(&wgpu::TextureViewDescriptor::default())); } @@ -190,7 +194,7 @@ impl SubjectMasks { } } -/// The rasterised masks for one edit. +/// The rasterised masks for one edit./// The rasterised masks for one edit. pub struct MaskArray { texture: wgpu::Texture, view: wgpu::TextureView, @@ -266,7 +270,10 @@ impl MaskPass { binding: 3, visibility: wgpu::ShaderStages::FRAGMENT, ty: wgpu::BindingType::Texture { - sample_type: wgpu::TextureSampleType::Float { filterable: true }, + // `filterable: false`: R32Float cannot be filtered + // without an optional feature, and the shader loads + // texels and interpolates them itself anyway. + sample_type: wgpu::TextureSampleType::Float { filterable: false }, view_dimension: wgpu::TextureViewDimension::D2, multisampled: false, }, @@ -312,7 +319,9 @@ impl MaskPass { } let placeholder = LabelField::upload(ctx, &[0], 1, 1, 0)?; - let empty_subject = SubjectMasks::upload(ctx, &[&[0u8][..]], 1, 1)?; + // Everywhere outside, so a layer that somehow reaches this masks + // nothing rather than everything. + let empty_subject = SubjectMasks::upload(ctx, &[&[-1.0f32][..]], 1, 1)?; Ok(Self { ctx: ctx.clone(), @@ -367,20 +376,20 @@ impl MaskPass { // 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. + // Indexed by *slot*, not by the instance the layer names: the + // fields are built per layer, in this same order, because two + // layers over one subject can carry different morphology. 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; - } + MaskSource::Subject { .. } => match subjects.filter(|s| slot < s.len()) { + Some(s) => (s, slot), + None => { + log::warn!( + "mask layer {} has no distance field; skipping", + layer.id + ); + continue; } - } + }, _ => (&self.empty_subject, 0), }; @@ -418,7 +427,7 @@ impl MaskPass { mode: MODE_REGIONS, region_count: field.region_count, feather: 0.0, - _pad0: 0.0, + falloff: 0, centre: [0.5, 0.5], axis: [1.0, 0.0], softness: 0.0, @@ -431,11 +440,24 @@ impl MaskPass { // 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 - }, + // Feather and morphology are in fractions of the frame's shorter + // edge; the field is in proxy pixels. Converting here keeps the + // stored edit resolution-independent while the shader works in the + // units its texture is actually measured in. + MaskSource::Subject { .. } => { + let short = field_short_edge(width, height); + MaskParams { + mode: MODE_SUBJECT, + // `softness` is the feather half-width in pixels. + softness: (layer.feather * short).max(0.0), + // `angle` carries the morphology offset — reused rather + // than padded, since a subject layer has no ellipse to + // rotate. + angle: morph_offset(layer) * short, + falloff: falloff_code(layer.falloff), + ..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. @@ -619,6 +641,40 @@ impl MaskPass { } } +/// The shorter edge of the space the mask is rasterised in. +/// +/// Feather and morphology are stored as fractions of it, so the same edit is +/// the same edge whether it renders to a viewport or to a 24 MP export. +fn field_short_edge(width: u32, height: u32) -> f32 { + width.min(height).max(1) as f32 +} + +/// How far the boundary moves, in fractions of the shorter edge. +/// +/// Zero for closing and opening: those are folded into the field itself when +/// it is built, because their second half acts on a shape the original field +/// does not describe. +fn morph_offset(layer: &dr_pipeline::mask::MaskLayer) -> f32 { + use dr_pipeline::mask::Morphology; + match layer.morphology { + Morphology::Dilate => layer.morph_radius, + Morphology::Erode => -layer.morph_radius, + Morphology::None | Morphology::Close | Morphology::Open => 0.0, + } +} + +/// Kept in step with the `switch` in `mask.wgsl`. +fn falloff_code(falloff: dr_pipeline::mask::Falloff) -> u32 { + use dr_pipeline::mask::Falloff; + match falloff { + Falloff::Hard => 0, + Falloff::Linear => 1, + Falloff::Smooth => 2, + Falloff::Gaussian => 3, + Falloff::Exponential => 4, + } +} + fn uniform_entry(binding: u32) -> wgpu::BindGroupLayoutEntry { wgpu::BindGroupLayoutEntry { binding, diff --git a/core/dr-gpu/src/shaders/mask.wgsl b/core/dr-gpu/src/shaders/mask.wgsl index 422a029..a108b4d 100644 --- a/core/dr-gpu/src/shaders/mask.wgsl +++ b/core/dr-gpu/src/shaders/mask.wgsl @@ -35,7 +35,9 @@ struct MaskParams { region_count: u32, // Softening applied to a region mask, in output pixels. feather: f32, - _pad0: f32, + // 0 hard, 1 linear, 2 smooth, 3 gaussian, 4 exponential. Kept in step with + // `falloff_code` on the Rust side. + falloff: u32, // Geometry, in normalised output coordinates. Meaning depends on `mode`. centre: vec2, @@ -57,9 +59,14 @@ 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. +// The **signed distance** from one subject's boundary, in proxy pixels: +// positive inside, negative outside. A 1x1 placeholder when the layer is not a +// subject — the binding is fixed, and a second pipeline differing only in what +// it ignores would be worse than a wasted texel. +// +// A distance field rather than a finished alpha is what makes growing, +// shrinking and feathering free: each is arithmetic on this, so a slider moves +// a uniform instead of rebuilding a mask. @group(0) @binding(3) var subject: texture_2d; // A full-screen triangle rather than a quad: three vertices instead of six, @@ -142,12 +149,12 @@ 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. +// Coverage for one object, from its distance field. // -// 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. +// Bilinear on the *distance*, which is the reason this is a distance field at +// all: distance varies smoothly across the boundary where coverage does not, +// so interpolating it gives a clean sub-pixel edge even though the model's +// own mask was quarter-resolution. fn subject_mask(uv: vec2) -> f32 { let dims = vec2(textureDimensions(subject)); let last = vec2(dims) - vec2(1); @@ -164,15 +171,34 @@ fn subject_mask(uv: vec2) -> f32 { 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); + // `angle` carries the morphology offset in pixels: positive grows the + // mask, negative shrinks it. Adding it before the falloff is what makes + // dilation move the boundary rather than merely brighten the edge. + let dist = mix(mix(a, b, f.x), mix(c, d, f.x), f.y) + p.angle; - // `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. + // `softness` is the feather half-width, also in pixels. if (p.softness <= 0.0) { - return cov; + return select(0.0, 1.0, dist >= 0.0); + } + let t_norm = dist / p.softness; + + // Every curve is 0.5 at the boundary, so changing the falloff changes how + // the transition looks and never where it sits. + switch p.falloff { + case 0u: { return select(0.0, 1.0, dist >= 0.0); } + case 1u: { return clamp(t_norm * 0.5 + 0.5, 0.0, 1.0); } + case 3u: { return 1.0 / (1.0 + exp(-3.0 * t_norm)); } + case 4u: { + if (t_norm >= 0.0) { + return 1.0 - 0.5 * exp(-3.0 * t_norm); + } + return 0.5 * exp(3.0 * t_norm); + } + default: { + let x = clamp(t_norm * 0.5 + 0.5, 0.0, 1.0); + return x * x * (3.0 - 2.0 * x); + } } - return smoothstep(0.5 - p.softness, 0.5 + p.softness, cov); } @fragment diff --git a/core/dr-segment/src/distance.rs b/core/dr-segment/src/distance.rs new file mode 100644 index 0000000..0c10d22 --- /dev/null +++ b/core/dr-segment/src/distance.rs @@ -0,0 +1,568 @@ +//! Exact Euclidean distance from a mask's boundary, and the shaping built on +//! it. +//! +//! # One measurement, four controls +//! +//! Feathering, growing, shrinking, closing and opening are the same number +//! read differently. Given the **signed** distance from the boundary — +//! positive inside, negative outside — dilation by `r` is the set where +//! `d >= -r`, erosion is `d >= +r`, and a feather of any shape is a function +//! of `d`. So the field is computed once and the controls are arithmetic on +//! it. +//! +//! That is also why the *field* is what gets uploaded to the GPU rather than a +//! finished alpha: growing a mask or changing its falloff then costs a uniform +//! upload and no recomputation at all, which is what makes those live +//! controls rather than ones that stall on every drag. +//! +//! Closing and opening are the exception. After the first threshold the shape +//! has changed, so the old distances describe the old one and a second field +//! is needed — [`Shaped::needs_recompute`] says so, and it is the only edge +//! control that is not free. +//! +//! # Why this is on the CPU +//! +//! ARCH §5.4 says masks rasterise on the GPU and never exist in CPU memory, +//! and the reason it says so is brush lag: a stroke rasterised per-frame on +//! the CPU is what makes darktable's drawn masks unusable. This is a different +//! operation with different economics. +//! +//! - It runs **once per mask edit**, not once per frame. +//! - Its input is already CPU-side — the model's coverage was produced here. +//! - Its output is a field the GPU then samples for free, forever after. +//! +//! And doing it here buys two things a shader could not. It is **exactly +//! deterministic**, which matters because masks reach the sidecar as indices +//! and a field that varied by vendor would mean a mask meaning one thing on +//! the desktop and another on the phone (docs/segmentation.md §6, M5). And it +//! is testable against hand-computed distances with no adapter present. +//! +//! # The transform +//! +//! Felzenszwalb and Huttenlocher's separable exact transform: the lower +//! envelope of parabolas along every row, then along every column. Linear in +//! the number of pixels, exactly Euclidean — not the chamfer approximation +//! that leaves a mask visibly octagonal when grown by more than a few pixels. + +/// Larger than any squared distance in an image anyone will render. +const FAR: f32 = 1e20; + +/// Signed Euclidean distance from the mask's boundary, in pixels. +/// +/// Positive inside, negative outside. `coverage` is one byte per pixel; +/// `threshold` is the value at or above which a pixel counts as inside. +/// +/// The sign convention is the one that makes the controls read naturally: a +/// positive offset grows the mask, matching "dilate by 3 pixels". +/// +/// # The half-pixel, which is not a detail +/// +/// The transform measures to the nearest pixel *centre* of the other class, so +/// the closest an inside pixel can be to an outside one is exactly 1. Reported +/// raw, that puts the boundary nowhere — no pixel is at zero, the smallest +/// magnitude either side is 1, and **eroding by anything under a pixel removes +/// nothing at all**. A control whose first notch does nothing is a broken +/// control. +/// +/// So half a pixel comes off each side, which puts the boundary where it +/// physically is: between the last inside pixel and the first outside one. +/// A pixel against the edge then reads `+0.5` inside and `-0.5` outside, +/// symmetric, and eroding by 1 takes exactly the outermost ring. +pub fn signed_distance(coverage: &[u8], width: usize, height: usize, threshold: u8) -> Vec { + assert_eq!( + coverage.len(), + width * height, + "coverage must cover every pixel" + ); + + let inside: Vec = coverage.iter().map(|&c| c >= threshold).collect(); + + // Two transforms, because a pixel's distance is to the nearest pixel of + // the *opposite* class and that is a different seed set on each side. + let to_outside = euclidean(&inside, width, height, false); + let to_inside = euclidean(&inside, width, height, true); + + inside + .iter() + .enumerate() + .map(|(i, &is_in)| { + if is_in { + to_outside[i] - 0.5 + } else { + -(to_inside[i] - 0.5) + } + }) + .collect() +} + +/// Distance from every pixel to the nearest pixel whose class is `seed`. +fn euclidean(inside: &[bool], width: usize, height: usize, seed: bool) -> Vec { + let mut f: Vec = inside + .iter() + .map(|&v| if v == seed { 0.0 } else { FAR }) + .collect(); + + // Scratch, allocated once and reused by every line: the parabola vertices + // still on the lower envelope, and the crossings between consecutive ones. + let longest = width.max(height); + let mut v = vec![0usize; longest]; + let mut z = vec![0.0f32; longest + 1]; + let mut out = vec![0.0f32; longest]; + + for y in 0..height { + let row: Vec = f[y * width..(y + 1) * width].to_vec(); + transform(&row, &mut out[..width], &mut v, &mut z); + f[y * width..(y + 1) * width].copy_from_slice(&out[..width]); + } + + let mut column = vec![0.0f32; height]; + for x in 0..width { + for y in 0..height { + column[y] = f[y * width + x]; + } + transform(&column, &mut out[..height], &mut v, &mut z); + for y in 0..height { + f[y * width + x] = out[y]; + } + } + + // Squared until here — the transform works in squares because that is what + // makes the parabolas parabolas. + f.iter().map(|d| d.max(0.0).sqrt()).collect() +} + +/// One-dimensional squared distance transform. +/// +/// `out[q] = min over p of ( f[p] + (q - p)^2 )`, computed by walking the +/// lower envelope of those parabolas. +fn transform(f: &[f32], out: &mut [f32], v: &mut [usize], z: &mut [f32]) { + let n = f.len(); + if n == 0 { + return; + } + + let mut k = 0usize; + v[0] = 0; + z[0] = -FAR; + z[1] = FAR; + + for q in 1..n { + if f[q] >= FAR { + // An infinite parabola is never the lowest anywhere, and + // including it would divide one infinity by another. + continue; + } + + loop { + let p = v[k]; + // Where this parabola crosses the one currently on top. + let s = ((f[q] + sq(q)) - (f[p] + sq(p))) / (2.0 * q as f32 - 2.0 * p as f32); + if s > z[k] { + k += 1; + v[k] = q; + z[k] = s; + z[k + 1] = FAR; + break; + } + if k == 0 { + // It dominates everything before it: start the envelope again + // from this vertex. + v[0] = q; + z[0] = -FAR; + z[1] = FAR; + break; + } + k -= 1; + } + } + + // Every parabola was infinite, so every distance is. + if f.iter().all(|&x| x >= FAR) { + out[..n].fill(FAR); + return; + } + + k = 0; + for q in 0..n { + while z[k + 1] < q as f32 { + k += 1; + } + let p = v[k]; + out[q] = (q as f32 - p as f32).powi(2) + f[p]; + } +} + +fn sq(x: usize) -> f32 { + (x as f32) * (x as f32) +} + +/// How coverage falls away from the boundary. +/// +/// Mirrors `dr_pipeline::mask::Falloff`, which this crate cannot see — the +/// pipeline depends on nothing here and inverting that to share one enum would +/// be a dependency edge for five variants. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum Falloff { + Hard, + Linear, + #[default] + Smooth, + Gaussian, + Exponential, +} + +impl Falloff { + /// Coverage at a signed distance, given a feather half-width. + /// + /// `t` is the distance normalised to the feather: `-1` is a feather-width + /// outside, `+1` a feather-width inside. Every curve returns `0.5` at the + /// boundary, which is what keeps the *edge* where the mask says it is + /// whichever curve is chosen — changing the falloff should change how the + /// transition looks, never where it sits. + pub fn coverage(self, t: f32) -> f32 { + match self { + Self::Hard => { + if t >= 0.0 { + 1.0 + } else { + 0.0 + } + } + Self::Linear => (t * 0.5 + 0.5).clamp(0.0, 1.0), + Self::Smooth => { + let x = (t * 0.5 + 0.5).clamp(0.0, 1.0); + x * x * (3.0 - 2.0 * x) + } + // A logistic curve rather than a true Gaussian integral: it is the + // same shape to the eye, has a closed form, and is exactly 0.5 at + // the boundary by construction. + Self::Gaussian => 1.0 / (1.0 + (-3.0 * t).exp()), + // Reaches full coverage quickly inside and trails off slowly + // outside, for blending an adjustment away without moving its edge. + Self::Exponential => { + if t >= 0.0 { + 1.0 - 0.5 * (-3.0 * t).exp() + } else { + 0.5 * (3.0 * t).exp() + } + } + } + } +} + +/// Growing, shrinking and tidying. Mirrors `dr_pipeline::mask::Morphology`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum Morphology { + #[default] + None, + Dilate, + Erode, + Close, + Open, +} + +impl Morphology { + /// Whether this needs the distance field rebuilt after a threshold. + pub fn needs_recompute(self) -> bool { + matches!(self, Self::Close | Self::Open) + } + + /// The offset applied to the distance before the falloff, in pixels. + /// + /// Zero for the compound pair: they are applied by [`apply_morphology`] + /// rather than by shifting the field, because their second half operates + /// on a shape the field does not describe. + pub fn offset(self, radius: f32) -> f32 { + match self { + Self::Dilate => radius, + Self::Erode => -radius, + Self::None | Self::Close | Self::Open => 0.0, + } + } +} + +/// Apply a compound morphology, returning a new distance field. +/// +/// Closing is a dilation followed by an erosion, opening the reverse. Each +/// half is a threshold of a field, and the second half needs a field of the +/// *thresholded* shape — so this recomputes once in the middle and is the one +/// edge control that is not free. +/// +/// Returns `None` for the operations that need no rebuild, so a caller can use +/// the field it already has. +pub fn apply_morphology( + distance: &[f32], + width: usize, + height: usize, + morphology: Morphology, + radius: f32, +) -> Option> { + if !morphology.needs_recompute() || radius <= 0.0 { + return None; + } + + // First half: threshold the existing field. + let first = match morphology { + Morphology::Close => -radius, // dilate + Morphology::Open => radius, // erode + _ => return None, + }; + let intermediate: Vec = distance + .iter() + .map(|&d| if d >= first { 255 } else { 0 }) + .collect(); + + // Second half: measure the new shape, and shift so the caller's threshold + // at zero performs the opposite operation. + let rebuilt = signed_distance(&intermediate, width, height, 128); + let second = match morphology { + Morphology::Close => radius, // erode + Morphology::Open => -radius, // dilate + _ => unreachable!("guarded above"), + }; + Some(rebuilt.iter().map(|&d| d - second).collect()) +} + +/// A distance field ready to be sampled, with the offset already folded in. +#[derive(Debug, Clone, PartialEq)] +pub struct Shaped { + pub distance: Vec, + pub width: usize, + pub height: usize, +} + +impl Shaped { + /// Build from coverage, applying whatever morphology needs a rebuild. + /// + /// The simple operations are *not* folded in here: they are an offset the + /// shader adds when it samples, so changing "grow by 4px" to "grow by 6px" + /// costs a uniform rather than a transform. + pub fn build( + coverage: &[u8], + width: usize, + height: usize, + threshold: u8, + morphology: Morphology, + radius_px: f32, + ) -> Self { + let distance = signed_distance(coverage, width, height, threshold); + let distance = apply_morphology(&distance, width, height, morphology, radius_px) + .unwrap_or(distance); + Self { + distance, + width, + height, + } + } + + /// Coverage at one pixel, for tests and for the CPU export path. + pub fn coverage_at(&self, index: usize, offset: f32, feather: f32, falloff: Falloff) -> f32 { + let d = self.distance[index] + offset; + if feather <= 0.0 { + return if d >= 0.0 { 1.0 } else { 0.0 }; + } + falloff.coverage(d / feather) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A 2px-wide vertical bar down the middle of a 9x1 strip. + fn bar() -> (Vec, usize, usize) { + let mut m = vec![0u8; 9]; + m[4] = 255; + (m, 9, 1) + } + + #[test] + fn distance_is_exact_along_a_line() { + let (m, w, h) = bar(); + let d = signed_distance(&m, w, h, 128); + + // The lone inside pixel sits half a pixel from the boundary on each + // side, and the pixels either side of it are half a pixel out. + assert_eq!(d[4], 0.5); + assert_eq!(d[3], -0.5); + assert_eq!(d[5], -0.5); + // Then one per pixel from there. + assert_eq!(d[2], -1.5); + assert_eq!(d[0], -3.5); + assert_eq!(d[8], -3.5); + } + + /// The property that separates an exact transform from a chamfer one: a + /// diagonal neighbour is √2 away, not 1 and not 2. + #[test] + fn diagonals_are_euclidean_not_chamfer() { + let mut m = vec![0u8; 25]; + m[12] = 255; // centre of 5x5 + let d = signed_distance(&m, 5, 5, 128); + + // Distance from the boundary, so the half-pixel comes back off to + // compare against the centre-to-centre figures. + let at = |x: usize, y: usize| -d[y * 5 + x] + 0.5; + assert!((at(1, 1) - std::f32::consts::SQRT_2).abs() < 1e-4, "{}", at(1, 1)); + assert!((at(0, 0) - (8.0f32).sqrt()).abs() < 1e-4, "{}", at(0, 0)); + assert_eq!(at(2, 0), 2.0, "straight up is exactly two"); + } + + #[test] + fn an_empty_mask_is_everywhere_outside() { + let d = signed_distance(&vec![0u8; 16], 4, 4, 128); + assert!(d.iter().all(|&v| v < 0.0), "no pixel can be inside"); + } + + #[test] + fn a_full_mask_is_everywhere_inside() { + let d = signed_distance(&vec![255u8; 16], 4, 4, 128); + assert!(d.iter().all(|&v| v > 0.0), "no pixel can be outside"); + } + + /// Every curve must cross at the boundary, or changing the falloff would + /// move the edge rather than soften it. + #[test] + fn every_falloff_is_half_at_the_boundary() { + for f in [ + Falloff::Linear, + Falloff::Smooth, + Falloff::Gaussian, + Falloff::Exponential, + ] { + let v = f.coverage(0.0); + assert!((v - 0.5).abs() < 1e-5, "{f:?} gave {v} at the boundary"); + } + } + + #[test] + fn every_falloff_is_monotone_and_bounded() { + for f in Falloff::ALL_FOR_TEST { + let mut previous = -1.0; + for i in -20..=20 { + let v = f.coverage(i as f32 / 10.0); + assert!((0.0..=1.0).contains(&v), "{f:?} left the range: {v}"); + assert!(v >= previous - 1e-6, "{f:?} went backwards at {i}"); + previous = v; + } + } + } + + #[test] + fn a_hard_falloff_has_no_transition() { + assert_eq!(Falloff::Hard.coverage(-0.01), 0.0); + assert_eq!(Falloff::Hard.coverage(0.0), 1.0); + } + + #[test] + fn dilating_grows_and_eroding_shrinks() { + let mut m = vec![0u8; 81]; + for y in 3..6 { + for x in 3..6 { + m[y * 9 + x] = 255; + } + } + let d = signed_distance(&m, 9, 9, 128); + + let area = |offset: f32| d.iter().filter(|&&v| v + offset >= 0.0).count(); + let plain = area(0.0); + assert_eq!(plain, 9, "the 3x3 block itself"); + assert!(area(Morphology::Dilate.offset(1.0)) > plain, "dilate grows"); + assert!(area(Morphology::Erode.offset(1.0)) < plain, "erode shrinks"); + } + + /// Closing fills a hole without growing the outline — the thing it is for. + #[test] + fn closing_fills_a_pinhole() { + // Generous margin on purpose. Outside the image is not "outside the + // mask", so a dilation that reaches the border cannot erode back from + // it, and a tight frame measures that artefact instead of the + // operation. + let (w, h) = (21, 21); + let mut m = vec![0u8; w * h]; + for y in 7..14 { + for x in 7..14 { + m[y * w + x] = 255; + } + } + // One missing pixel inside, as a soft mask commonly has. + m[10 * w + 10] = 0; + + let before = signed_distance(&m, w, h, 128); + assert!(before[10 * w + 10] < 0.0, "the hole starts outside the mask"); + + let after = apply_morphology(&before, w, h, Morphology::Close, 2.0).expect("recomputed"); + assert!(after[10 * w + 10] >= 0.0, "closing should have filled it"); + + // The outline may round at a corner — closing does that, and it is the + // price of filling — but it must not march outward across the frame. + let grew = after + .iter() + .zip(&before) + .filter(|(a, b)| **a >= 0.0 && **b < 0.0) + .count(); + assert!(grew <= 5, "closing should fill the hole, not grow: {grew}"); + } + + /// Opening removes a speck without shrinking the body. + #[test] + fn opening_removes_a_speck() { + let (w, h) = (13, 13); + let mut m = vec![0u8; w * h]; + for y in 3..10 { + for x in 3..10 { + m[y * w + x] = 255; + } + } + m[0] = 255; // an isolated speck in the corner + + let before = signed_distance(&m, w, h, 128); + assert!(before[0] >= 0.0, "the speck starts inside the mask"); + + let after = apply_morphology(&before, w, h, Morphology::Open, 1.5).expect("recomputed"); + assert!(after[0] < 0.0, "opening should have removed it"); + assert!( + after[6 * w + 6] >= 0.0, + "and left the body of the mask alone" + ); + } + + #[test] + fn the_simple_operations_need_no_rebuild() { + assert!(!Morphology::None.needs_recompute()); + assert!(!Morphology::Dilate.needs_recompute()); + assert!(!Morphology::Erode.needs_recompute()); + assert!(Morphology::Close.needs_recompute()); + assert!(Morphology::Open.needs_recompute()); + } + + #[test] + fn the_same_mask_gives_the_same_field_twice() { + // M5: masks reach the sidecar as indices, so the field they are shaped + // by has to be reproducible. + let mut m = vec![0u8; 64]; + for i in 20..30 { + m[i] = 255; + } + let a = signed_distance(&m, 8, 8, 128); + let b = signed_distance(&m, 8, 8, 128); + assert_eq!(a, b); + } + + #[test] + fn zero_feather_is_a_step() { + let shaped = Shaped::build(&[0, 255, 255, 0], 4, 1, 128, Morphology::None, 0.0); + assert_eq!(shaped.coverage_at(1, 0.0, 0.0, Falloff::Smooth), 1.0); + assert_eq!(shaped.coverage_at(3, 0.0, 0.0, Falloff::Smooth), 0.0); + } + + impl Falloff { + const ALL_FOR_TEST: [Falloff; 5] = [ + Falloff::Hard, + Falloff::Linear, + Falloff::Smooth, + Falloff::Gaussian, + Falloff::Exponential, + ]; + } +} diff --git a/core/dr-segment/src/lib.rs b/core/dr-segment/src/lib.rs index dc6476e..4909f5b 100644 --- a/core/dr-segment/src/lib.rs +++ b/core/dr-segment/src/lib.rs @@ -29,11 +29,13 @@ //! foliage or wall (`models/LICENCE.md`). Selecting those falls to arm A, //! which never needed a vocabulary to begin with. +pub mod distance; pub mod hierarchy; pub mod prior; #[cfg(feature = "semantic")] pub mod semantic; +pub use distance::{signed_distance, Falloff, Morphology, Shaped}; pub use hierarchy::{Edge, Merge, MergeTree, RegionField}; pub use prior::{Membership, PriorOptions}; #[cfg(feature = "semantic")] diff --git a/ui/dr-ui/src/develop.rs b/ui/dr-ui/src/develop.rs index ffaab24..9583996 100644 --- a/ui/dr-ui/src/develop.rs +++ b/ui/dr-ui/src/develop.rs @@ -27,6 +27,12 @@ use crate::ParamRow; /// A loaded image plus its edit state. pub struct DevelopSession { + /// Kept so the session can build GPU resources after construction. + /// + /// The distance fields behind a subject mask are made when a layer is + /// *shaped*, not when the image opens, and cloning a `GpuContext` is two + /// `Arc` bumps. + ctx: GpuContext, graph: EditGraph, /// TRACES: FR-DEV-5 /// Undo, kept beside the graph rather than in the window. @@ -55,8 +61,16 @@ pub struct DevelopSession { segmentation: Option, /// Rasterises the mask layers. Built lazily for the same reason. masks: Option, - /// The recognised objects' coverage, on the GPU. + /// One signed distance field per active subject layer, on the GPU. subjects: Option, + /// What `subjects` was built from. + /// + /// The fields are expensive — an exact distance transform over the proxy + /// for each layer — and almost nothing changes them. Feather, falloff and + /// simple growing are arithmetic the shader does on the field it already + /// has, so this deliberately does *not* include them: dragging those + /// sliders must not rebuild anything. + subject_key: u64, /// Which layer the develop panel is editing, if any. /// /// This is what lets one panel serve both scopes: with a layer selected, @@ -115,6 +129,7 @@ impl DevelopSession { graph.set_orientation(orientation); let history = History::new(&graph); Self { + ctx: ctx.clone(), graph, history, demosaiced, @@ -125,6 +140,7 @@ impl DevelopSession { segmentation: None, masks: None, subjects: None, + subject_key: 0, active_mask: None, show_overlay: false, } @@ -604,6 +620,96 @@ 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`. + /// What the uploaded distance fields depend on. + /// + /// The instance each layer names, and the morphology that *rebuilds* a + /// field rather than offsetting it. Nothing else: see `subject_key`. + fn subject_signature(&self) -> u64 { + use dr_pipeline::mask::{MaskSource, Morphology}; + let mut h: u64 = 0xcbf2_9ce4_8422_2325; + let mut mix = |v: u64| { + for byte in v.to_le_bytes() { + h ^= byte as u64; + h = h.wrapping_mul(0x1000_0000_01b3); + } + }; + + for layer in self.graph.masks().active() { + match &layer.source { + MaskSource::Subject { index, .. } => { + mix(1); + mix(*index as u64); + // Only the compound operations change the field itself. + if layer.morphology.is_compound() { + mix(match layer.morphology { + Morphology::Close => 2, + Morphology::Open => 3, + _ => 0, + }); + mix(layer.morph_radius.to_bits() as u64); + } + } + _ => mix(0), + } + } + h + } + + /// Rebuild the distance fields if anything they depend on moved. + fn ensure_subject_fields(&mut self, ctx: &GpuContext) { + use dr_pipeline::mask::MaskSource; + + let key = self.subject_signature(); + if key == self.subject_key && self.subjects.is_some() { + return; + } + + let Some(seg) = self.segmentation.as_ref() else { + return; + }; + let (pw, ph) = seg.proxy_size(); + + // In `active()` order, because that is the order the rasteriser walks + // and the order it indexes these by. + let mut fields: Vec> = Vec::new(); + for layer in self.graph.masks().active() { + let field = match &layer.source { + MaskSource::Subject { index, .. } => seg + .instance_mask(*index as usize) + .map(|coverage| { + dr_segment::Shaped::build( + coverage, + pw, + ph, + 128, + morphology_for(layer.morphology), + // Radii are fractions of the shorter edge; the + // field is in proxy pixels. + layer.morph_radius * pw.min(ph) as f32, + ) + .distance + }) + .unwrap_or_default(), + // A placeholder of the right size, so the slot indices line up + // with `active()` whatever mix of sources the stack holds. + _ => vec![-1.0; pw * ph], + }; + fields.push(field); + } + + if fields.is_empty() { + self.subjects = None; + self.subject_key = key; + return; + } + + let refs: Vec<&[f32]> = fields.iter().map(|f| f.as_slice()).collect(); + self.subjects = dr_gpu::SubjectMasks::upload(ctx, &refs, pw as u32, ph as u32) + .inspect_err(|e| log::warn!("could not upload the subject fields: {e}")) + .ok(); + self.subject_key = key; + } + fn rasterise_masks(&mut self) -> bool { if self.graph.masks().is_neutral() { return false; @@ -661,19 +767,12 @@ impl DevelopSession { .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() - }; - + // The fields themselves are built per *layer*, on demand — there are + // none yet, and building one per detected object would transform + // several megapixels for masks the user may never make. self.segmentation = Some(seg); + self.subjects = None; + self.subject_key = 0; Ok(()) } @@ -1082,6 +1181,11 @@ 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. + // Any layer whose *shape* changed needs its field rebuilt before the + // rasteriser reads it. Keyed, so a feather drag reaches neither. + let ctx = self.ctx.clone(); + self.ensure_subject_fields(&ctx); + let masks = self .rasterise_masks() .then(|| self.masks.as_ref().and_then(|p| p.array())) @@ -2789,3 +2893,21 @@ mod tests { } } } + +/// `dr_pipeline`'s morphology, as `dr_segment` names it. +/// +/// Two enums for one idea, and deliberately: `dr-pipeline` describes the +/// *edit* and `dr-segment` implements the *transform*, and neither depends on +/// the other. The crossing is this function, which the compiler makes +/// exhaustive on both sides. +fn morphology_for(m: dr_pipeline::mask::Morphology) -> dr_segment::Morphology { + use dr_pipeline::mask::Morphology as Edit; + use dr_segment::Morphology as Transform; + match m { + Edit::None => Transform::None, + Edit::Dilate => Transform::Dilate, + Edit::Erode => Transform::Erode, + Edit::Close => Transform::Close, + Edit::Open => Transform::Open, + } +} diff --git a/ui/dr-ui/src/segmentation.rs b/ui/dr-ui/src/segmentation.rs index 445c280..f7ff178 100644 --- a/ui/dr-ui/src/segmentation.rs +++ b/ui/dr-ui/src/segmentation.rs @@ -84,9 +84,13 @@ pub struct InstanceSummary { 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. + /// Empty without a watershed, which is the default, and unread by anything + /// today. Kept — with the lint silenced rather than the field deleted — + /// because it is the mechanism that would give a model's outline the + /// image's own edge, and that is the point of the whole exercise if the + /// hierarchy underneath it is ever made to work. Deleting it would mean + /// rediscovering `regions_for_instance` from the git history. + #[allow(dead_code)] pub regions: Vec, /// Coverage at proxy resolution, quantised to a byte. ///