diff --git a/Cargo.lock b/Cargo.lock index dc3519d..a897932 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1568,6 +1568,7 @@ dependencies = [ "dr-gpu", "dr-pipeline", "dr-plat", + "dr-segment", "dr-sync", "dr-sync-nextcloud", "dr-thumbs", diff --git a/core/dr-gpu/Cargo.toml b/core/dr-gpu/Cargo.toml index adcc8a4..bcbf9a2 100644 --- a/core/dr-gpu/Cargo.toml +++ b/core/dr-gpu/Cargo.toml @@ -40,7 +40,12 @@ name = "develop" default = [] # Exposes read_pixels outside tests. Production must not enable this. readback = [] +# Exposes `Segmentation::read_field`, which builds the region adjacency graph +# on the CPU. Separate from `readback` on purpose — see `segment.rs`. Once per +# image on a worker, not the per-frame display round-trip AC-8 forbids; still a +# full-resolution transfer, and still F3's open gap. +segment-readback = [] [[example]] name = "segment" -required-features = ["readback"] +required-features = ["segment-readback"] diff --git a/core/dr-gpu/src/segment.rs b/core/dr-gpu/src/segment.rs index b265460..8134de6 100644 --- a/core/dr-gpu/src/segment.rs +++ b/core/dr-gpu/src/segment.rs @@ -14,16 +14,17 @@ //! # The open question this leaves //! //! [`Segmentation::read_field`] copies the label and gradient buffers back to -//! the CPU to build the region adjacency graph, and is gated behind the -//! `readback` feature for the same reason `read_pixels` is. That gate is not -//! ceremony: a shipping build cannot take this path (ARCH §6.1, AC-8), so the -//! RAG would have to be accumulated GPU-side with atomics instead. +//! the CPU to build the region adjacency graph, behind the `segment-readback` +//! feature. That is deliberately **not** the `readback` switch guarding the +//! display round-trip: this transfer is once per image on a worker, where the +//! one AC-8 forbids is per frame in the render loop, and sharing a switch +//! would force a build wanting local masking to unlock the other. //! -//! For a spike that trade is the right way round — the readback is once per -//! image and off the frame path, and building the GPU-side RAG before knowing -//! whether the granularity ladder is any good would be work spent on a -//! question not yet asked. But it is a real gap between this and something -//! shippable, and it should be read as one. +//! It is still a real cost and still unfinished. F3 in docs/segmentation.md +//! §12 stands: the adjacency accumulation belongs GPU-side with atomics, and +//! until it moves there every segmentation pays a full-resolution transfer. +//! Read the feature name as a description of a known gap rather than as +//! permission. use wgpu::util::DeviceExt; @@ -472,9 +473,26 @@ impl Segmentation { /// Build the region adjacency graph, reading the labels back to the CPU. /// - /// **Not a shipping path** — see this module's header. Gated so it cannot - /// be reached from a production build by accident. - #[cfg(any(test, feature = "readback"))] + /// # Why this has its own feature rather than sharing `readback` + /// + /// `readback` gates [`crate::AdjustPass::read_pixels`], which is the + /// per-frame display round-trip AC-8 exists to forbid. This is a different + /// transfer with different economics, and sharing one switch would have + /// forced a build wanting local masking to also unlock the one thing the + /// architecture is built around never doing. + /// + /// What this transfer actually is: **once per image, on a worker, off the + /// frame path.** Nothing in the render loop waits on it, and the result is + /// a region graph of a few thousand nodes that every later interaction + /// reads from the CPU anyway. + /// + /// What it is *not* is finished. F3 in docs/segmentation.md §12 stands: + /// the adjacency accumulation belongs on the GPU with atomics, and until + /// it moves there a segmentation costs one full-resolution transfer of the + /// label and gradient buffers. That is a real cost on a phone and the + /// reason this is named for what it does rather than hidden behind the + /// general switch. + #[cfg(any(test, feature = "segment-readback"))] pub fn read_field(&self) -> Result { let n = (self.width * self.height) as usize; let roots: Vec = read_buffer(&self.ctx, &self.labels, n)?; @@ -551,7 +569,7 @@ fn proxy_size(src_w: u32, src_h: u32, max_edge: u32) -> (u32, u32) { ) } -#[cfg(any(test, feature = "readback"))] +#[cfg(any(test, feature = "segment-readback"))] fn read_buffer( ctx: &GpuContext, buffer: &wgpu::Buffer, diff --git a/core/dr-pipeline/src/mask.rs b/core/dr-pipeline/src/mask.rs index a5d6cf1..7ab3de3 100644 --- a/core/dr-pipeline/src/mask.rs +++ b/core/dr-pipeline/src/mask.rs @@ -234,6 +234,57 @@ impl MaskLayer { self.ops.iter().map(|o| o.descriptor()).collect() } + /// TRACES: FR-DEV-3a | FR-DEV-3c + /// The controls for this layer's adjustments. + /// + /// The same shape [`crate::EditGraph::capabilities`] returns, so a panel + /// that can render the global chain renders a mask layer with no new code + /// — which is the practical payoff of a layer holding a real chain rather + /// than a handful of special-cased sliders. + /// + /// Framing is absent, and that is the one real difference: a crop changes + /// the output's dimensions, so it is a property of the photograph and not + /// of a region within it. There is no such thing as cropping part of an + /// image. + pub fn capabilities(&self) -> Vec { + self.ops + .iter() + .map(|op| { + let desc = op.descriptor(); + crate::graph::OpCapability { + id: desc.id, + label: desc.label, + active: op.is_active(), + params: desc + .params + .iter() + .map(|p| crate::graph::ParamCapability { + id: p.id, + label: p.label, + kind: p.kind.clone(), + default: p.default, + value: op.param(p.id), + facet: p.facet, + }) + .collect(), + presentation: op.presentation(), + } + }) + .collect() + } + + /// Reset every adjustment, keeping the selection. + /// + /// The selection is the expensive half — it took a click and a scroll to + /// arrive at — so "start this layer's edit again" must not throw it away. + pub fn reset_adjustments(&mut self) { + for op in &mut self.ops { + for p in op.descriptor().params { + op.set_param(p.id, p.default); + } + } + } + pub fn set_param(&mut self, op: &str, param: ParamId, value: f32) { if let Some(o) = self.ops.iter_mut().find(|o| o.descriptor().id.0 == op) { let clamped = o diff --git a/ui/dr-ui/Cargo.toml b/ui/dr-ui/Cargo.toml index 8382119..8acb901 100644 --- a/ui/dr-ui/Cargo.toml +++ b/ui/dr-ui/Cargo.toml @@ -11,7 +11,14 @@ dr-types.workspace = true # the compositor the texture itself and there is no display round-trip left to # gate (ARCH §6.1, AC-8). The export path reads pixels back through # `export_pixels`, which is ungated and always was. -dr-gpu.workspace = true +# +# `segment-readback` *is* on, and it is not a contradiction of the above. It +# gates the region-graph transfer that local masking is built on: once per +# image, on a worker, off the frame path. The display round-trip AC-8 forbids +# stays behind its own switch, which remains off. See `dr-gpu/src/segment.rs`. +dr-gpu = { workspace = true, features = ["segment-readback"] } +# The semantic arm and its weights, for local adjustments (FR-DEV-3, D14). +dr-segment = { workspace = true, features = ["semantic", "embedded-model"] } dr-decode.workspace = true serde_json.workspace = true tokio.workspace = true diff --git a/ui/dr-ui/src/develop.rs b/ui/dr-ui/src/develop.rs index 0b062f4..555c9a8 100644 --- a/ui/dr-ui/src/develop.rs +++ b/ui/dr-ui/src/develop.rs @@ -10,7 +10,12 @@ //! new operation appears in the panel with no change here (FR-DEV-3c). use dr_decode::RawImage; -use dr_gpu::{AdjustPass, DemosaicedImage, Demosaicer, GpuContext, Histogram, HistogramPass}; +use dr_gpu::{ + AdjustPass, DemosaicedImage, Demosaicer, GpuContext, Histogram, HistogramPass, MaskPass, +}; +use dr_pipeline::mask::{MaskLayer, MaskSource}; + +use crate::segmentation::{self, Segmentation}; use dr_pipeline::ops::curve; use dr_pipeline::{ CropRect, Edit, EditGraph, History, OpCapability, OpId, ParamId, ParamKind, Presentation, @@ -40,6 +45,24 @@ pub struct DevelopSession { /// old driver, a device without the storage-buffer atomics it needs — the /// photographer loses the histogram and keeps the photograph. histogram: Option, + + /// TRACES: FR-DEV-3 + /// The region map local masks select from, once it has been computed. + /// + /// `None` until the photographer asks for it. Segmentation costs about + /// half a second and most edits never need one, so running it on open + /// would tax every photograph for a feature used on some of them. + segmentation: Option, + /// Rasterises the mask layers. Built lazily for the same reason. + masks: Option, + /// Which layer the develop panel is editing, if any. + /// + /// This is what lets one panel serve both scopes: with a layer selected, + /// the sliders read and write *its* chain, and the photographer is + /// adjusting a region rather than the frame. + active_mask: Option, + /// Whether to draw the false-coloured region overlay. + show_overlay: bool, } impl DevelopSession { @@ -97,6 +120,10 @@ impl DevelopSession { histogram: HistogramPass::new(ctx) .inspect_err(|e| log::warn!("no histogram on this device: {e}")) .ok(), + segmentation: None, + masks: None, + active_mask: None, + show_overlay: false, } } @@ -105,7 +132,24 @@ impl DevelopSession { /// Built entirely from the capability list. The `kind` string chooses the /// widget; nothing switches on a parameter's identity. pub fn rows(&self) -> Vec { - rows_from(&self.graph.capabilities()) + match self.active_layer() { + // A selected layer takes over the panel. The rows are built from + // the layer's own capability list, so every control the global + // chain offers is offered here too — including ones added later, + // which need no work to become local. + Some(layer) => rows_from(&layer.capabilities()), + None => rows_from(&self.graph.capabilities()), + } + } + + fn active_layer(&self) -> Option<&MaskLayer> { + let id = self.active_mask.as_ref()?; + self.graph.masks().get(id) + } + + fn active_layer_mut(&mut self) -> Option<&mut MaskLayer> { + let id = self.active_mask.clone()?; + self.graph.masks_mut().get_mut(&id) } } @@ -463,10 +507,24 @@ impl DevelopSession { /// here is curve-shaped; it walks whatever parameters the operation /// declares. pub fn reset_op(&mut self, op_index: i32) { - let caps = self.graph.capabilities(); + let caps = match self.active_layer() { + Some(layer) => layer.capabilities(), + None => self.graph.capabilities(), + }; let Some(cap) = usize::try_from(op_index).ok().and_then(|i| caps.get(i)) else { return; }; + if self.active_mask.is_some() { + let params: Vec<_> = cap.params.iter().map(|p| (p.id, p.default)).collect(); + let id = cap.id.0; + if let Some(layer) = self.active_layer_mut() { + for (param, default) in params { + layer.set_param(id, param, default); + } + } + self.history.record(&self.graph, Edit::Discrete); + return; + } for p in &cap.params { self.graph.set_param(cap.id, p.id, p.default); } @@ -493,6 +551,17 @@ impl DevelopSession { log::warn!("control at ({op_index}, {param_index}) has no parameter"); return; }; + if self.active_mask.is_some() { + if let Some(layer) = self.active_layer_mut() { + layer.set_param(op.0, param, value); + } + // Coalesced the same way a global drag is: a slider dragged across + // a masked layer is still one gesture and must undo as one. + let edit = Edit::for_param(&self.graph, op, param); + self.history.record(&self.graph, edit); + return; + } + self.graph.set_param(op, param, value); let edit = Edit::for_param(&self.graph, op, param); self.history.record(&self.graph, edit); @@ -520,10 +589,372 @@ impl DevelopSession { self.history.record(&self.graph, Edit::Discrete); } + /// Rasterise the current mask stack, if there is one. + /// + /// Returns `None` for a stack with no active layers, which is the common + /// case and the one that must cost nothing: the adjust pass then binds its + /// own placeholder and the generated shader has no layer block to read it + /// with. + /// Returns whether the array is now valid for the current stack. + /// + /// Split from reading the array back because the render below needs + /// `self.adjust` mutably while holding `self.masks` immutably. Those are + /// disjoint fields and the borrow checker will allow it — but only when + /// each is reached directly rather than through a method taking `self`. + fn rasterise_masks(&mut self, w: u32, h: u32) -> bool { + if self.graph.masks().is_neutral() { + return false; + } + let labels = self.segmentation.as_ref().and_then(|s| s.labels()); + let Some(pass) = self.masks.as_mut() else { + return false; + }; + pass.render(self.graph.masks(), labels, w, h) + .inspect_err(|e| log::warn!("mask rasterisation failed: {e}")) + .is_ok() + } + + // ---------------------------------------------------------------------- + // Segmentation (S15, docs/segmentation.md) + // ---------------------------------------------------------------------- + + /// TRACES: FR-DEV-3 + /// Compute the region map this image's local masks select from. + /// + /// **Blocking, and roughly half a second.** The caller is responsible for + /// running it off the UI thread — see the worker in `lib.rs`. It is + /// exposed as a plain blocking call rather than something async because + /// what it needs is a GPU context and a CPU core, not a runtime. + pub fn segment(&mut self, ctx: &GpuContext, options: &segmentation::Options) -> Result<(), String> { + // The model reads the photograph as captured, not as edited: the + // segmentation must survive an exposure change, or every slider would + // invalidate the masks that depend on it (docs/segmentation.md §3). + let (rgb, rw, rh) = self.neutral_proxy(ctx, options.segment.max_edge)?; + let seg = segmentation::compute(ctx, &self.demosaiced, &rgb, rw, rh, options)?; + + if self.masks.is_none() { + self.masks = MaskPass::new(ctx) + .inspect_err(|e| log::warn!("no mask rasteriser on this device: {e}")) + .ok(); + } + self.segmentation = Some(seg); + Ok(()) + } + + /// Render the *unedited* image to a CPU buffer at proxy size. + /// + /// Goes through a throwaway [`AdjustPass`] with a neutral graph rather + /// than the session's own. Reusing `self.adjust` would overwrite the frame + /// the histogram reads and leave the view showing an unedited image until + /// the next redraw — a visible flicker for the sake of not allocating. + /// + /// This is `export_pixels`, which is ungated: an export is not the display + /// round-trip AC-8 forbids, and neither is this. + fn neutral_proxy( + &self, + ctx: &GpuContext, + max_edge: u32, + ) -> Result<(Vec, usize, usize), String> { + let (sw, sh) = self.demosaiced.size(); + let scale = (max_edge as f32 / sw.max(sh) as f32).min(1.0); + let (w, h) = ( + ((sw as f32 * scale) as u32).max(1), + ((sh as f32 * scale) as u32).max(1), + ); + + let neutral = EditGraph::default_chain(); + let mut pass = AdjustPass::new(ctx); + pass.render(&self.demosaiced, &neutral.compose(), w, h) + .map_err(|e| format!("could not render the segmentation proxy: {e}"))?; + let (rgba, pw, ph) = pass + .export_pixels() + .map_err(|e| format!("could not read the segmentation proxy: {e}"))?; + + // Straight to float RGB, dropping alpha. The values stay display- + // encoded because that is what the model was trained on — one of the + // few places in this codebase where not linearising is correct. + let rgb = 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(); + Ok((rgb, pw as usize, ph as usize)) + } + + pub fn has_segmentation(&self) -> bool { + self.segmentation.is_some() + } + + pub fn segmentation_level(&self) -> u32 { + self.segmentation.as_ref().map_or(0, |s| s.level()) + } + + pub fn segmentation_region_count(&self) -> usize { + self.segmentation.as_ref().map_or(0, |s| s.region_count()) + } + + /// Move the granularity ladder — the scroll wheel over the canvas. + pub fn set_segmentation_level(&mut self, level: u32) { + if let Some(seg) = self.segmentation.as_mut() { + seg.set_level(level); + } + } + + /// The subjects the model recognised, as `(label, confidence)`. + /// + /// Confidence is shown rather than hidden because the detector is offered + /// as a shortcut, not as an authority: a 0.42 "dog" is worth listing and + /// worth flagging, and a list that presented it identically to a 0.95 one + /// would make the tool look wrong when the guess was merely weak. + pub fn detected_subjects(&self) -> Vec<(String, f32)> { + self.segmentation + .as_ref() + .map(|s| { + s.instances() + .iter() + .map(|i| (i.class_name.to_string(), i.score)) + .collect() + }) + .unwrap_or_default() + } + + // ---------------------------------------------------------------------- + // The region overlay + // ---------------------------------------------------------------------- + + pub fn overlay_enabled(&self) -> bool { + self.show_overlay + } + + pub fn set_overlay(&mut self, on: bool) { + self.show_overlay = on; + } + + /// TRACES: FR-DEV-3 + /// A false-coloured picture of the current grouping, for the canvas. + /// + /// Returned as a CPU image rather than a texture, and deliberately: it is + /// regenerated only when the level changes, it is proxy-sized rather than + /// viewport-sized, and Slint scales and composites it for free. Putting it + /// on the GPU would buy nothing and add a second texture to keep in step + /// with the view. + /// + /// `None` when the overlay is off or nothing has been segmented, so the + /// caller can bind this straight to an image source. + pub fn overlay_image(&self) -> Option { + if !self.show_overlay { + return None; + } + let (rgba, w, h) = self.segmentation.as_ref()?.overlay_rgba(); + let buffer = slint::SharedPixelBuffer::::clone_from_slice(&rgba, w, h); + Some(slint::Image::from_rgba8(buffer)) + } + + // ---------------------------------------------------------------------- + // Mask layers + // ---------------------------------------------------------------------- + + /// The layers, as `(id, name, enabled, is_active_selection)`. + pub fn mask_layers(&self) -> Vec<(String, String, bool, bool)> { + self.graph + .masks() + .layers() + .iter() + .map(|l| { + ( + l.id.clone(), + l.display_name().to_string(), + l.enabled, + Some(&l.id) == self.active_mask.as_ref(), + ) + }) + .collect() + } + + pub fn active_mask(&self) -> Option<&str> { + self.active_mask.as_deref() + } + + /// Select a layer for editing, or `None` to return the panel to the + /// global chain. + pub fn set_active_mask(&mut self, id: Option<&str>) { + self.active_mask = id + .filter(|id| self.graph.masks().get(id).is_some()) + .map(|id| id.to_string()); + } + + /// Select the region under a normalised image point. + /// + /// `add` extends the selected layer instead of replacing its selection, + /// which is the shift-click every selection tool has. With no layer + /// selected a new one is created, because clicking the photograph is how a + /// local adjustment begins and requiring "add layer" first would be a step + /// with no decision in it. + /// + /// Returns the layer that now holds the selection. + pub fn select_region_at(&mut self, x: f32, y: f32, add: bool) -> Option { + let seg = self.segmentation.as_ref()?; + let picked = seg.regions_at(x, y); + if picked.is_empty() { + return None; + } + let (signature, level) = (seg.signature(), seg.level()); + + let id = match self.active_mask.clone() { + Some(id) => id, + None => { + let id = self.graph.masks().next_id(); + let layer = MaskLayer::new( + id.clone(), + MaskSource::Regions { + signature, + level, + ids: Vec::new(), + }, + ); + if !self.graph.masks_mut().push(layer) { + return None; + } + self.active_mask = Some(id.clone()); + id + } + }; + + let layer = self.graph.masks_mut().get_mut(&id)?; + let mut ids = match (&layer.source, add) { + (MaskSource::Regions { ids, .. }, true) => ids.clone(), + _ => Vec::new(), + }; + + // Clicking a region already in the selection removes it, so one + // gesture both adds and corrects — the alternative is a modifier for + // subtract that nobody remembers. + if add && picked.iter().all(|r| ids.contains(r)) { + ids.retain(|r| !picked.contains(r)); + } else { + ids.extend(picked); + } + ids.sort_unstable(); + ids.dedup(); + + layer.source = MaskSource::Regions { + signature, + level, + ids, + }; + self.history.record(&self.graph, Edit::Discrete); + Some(id) + } + + /// Add a layer selecting one detected subject. + pub fn add_subject_mask(&mut self, index: usize) -> Option { + let seg = self.segmentation.as_ref()?; + let instance = seg.instances().get(index)?; + let (signature, level) = (seg.signature(), seg.level()); + let (name, ids) = (instance.class_name.to_string(), instance.regions.clone()); + + let id = self.graph.masks().next_id(); + let mut layer = MaskLayer::new( + id.clone(), + MaskSource::Regions { + signature, + level, + ids, + }, + ); + layer.name = name; + if !self.graph.masks_mut().push(layer) { + return None; + } + self.active_mask = Some(id.clone()); + self.history.record(&self.graph, Edit::Discrete); + Some(id) + } + + /// Add a gradient layer, which needs no segmentation. + pub fn add_gradient_mask(&mut self, radial: bool) -> Option { + let id = self.graph.masks().next_id(); + let source = if radial { + MaskSource::Radial { + centre: (0.5, 0.5), + radii: (0.35, 0.35), + angle: 0.0, + feather: 0.5, + } + } else { + MaskSource::Linear { + centre: (0.5, 0.5), + angle: std::f32::consts::FRAC_PI_2, + width: 0.3, + } + }; + if !self.graph.masks_mut().push(MaskLayer::new(id.clone(), source)) { + return None; + } + self.active_mask = Some(id.clone()); + self.history.record(&self.graph, Edit::Discrete); + Some(id) + } + + pub fn remove_mask(&mut self, id: &str) { + if self.graph.masks_mut().remove(id).is_some() { + if self.active_mask.as_deref() == Some(id) { + self.active_mask = None; + } + self.history.record(&self.graph, Edit::Discrete); + } + } + + pub fn set_mask_enabled(&mut self, id: &str, enabled: bool) { + if let Some(layer) = self.graph.masks_mut().get_mut(id) { + layer.enabled = enabled; + self.history.record(&self.graph, Edit::Discrete); + } + } + + pub fn set_mask_invert(&mut self, id: &str, invert: bool) { + if let Some(layer) = self.graph.masks_mut().get_mut(id) { + layer.invert = invert; + self.history.record(&self.graph, Edit::Discrete); + } + } + + pub fn set_mask_opacity(&mut self, id: &str, opacity: f32) { + if let Some(layer) = self.graph.masks_mut().get_mut(id) { + layer.opacity = opacity.clamp(0.0, 1.0); + // `Op` rather than `Discrete`: opacity is dragged, and a drag is + // one decision however many values it passes through. `Discrete` + // would put every intermediate position on the undo stack. + self.history.record(&self.graph, Edit::Op(OpId("mask-opacity"))); + } + } + + /// Whether a layer's region ids belong to a segmentation other than the + /// one currently loaded — a mask restored from a sidecar written under + /// different tuning. + pub fn mask_is_stale(&self, id: &str) -> bool { + let Some(layer) = self.graph.masks().get(id) else { + return false; + }; + match self.segmentation.as_ref() { + Some(seg) => layer.is_stale(seg.signature()), + // Nothing loaded to compare against. Not stale, just unrenderable + // — the distinction matters because "stale" invites the user to + // recompute the selection and this only needs the segmentation + // running. + None => false, + } + } + fn lookup(&self, op_index: i32, param_index: i32) -> Option<(OpId, ParamId)> { // Rows are emitted in capability order, so the flat index is the sum - // of preceding parameter counts. - let caps = self.graph.capabilities(); + // of preceding parameter counts. Taken from whichever scope `rows` + // last described — the indices the interface is holding are positions + // in *that* list, and reading the global chain while a layer is + // selected would map a slider onto a different operation. + let caps = match self.active_layer() { + Some(layer) => layer.capabilities(), + None => self.graph.capabilities(), + }; let op = caps.get(usize::try_from(op_index).ok()?)?; let param = op.params.get(usize::try_from(param_index).ok()?)?; Some((op.id, param.id)) @@ -560,9 +991,16 @@ impl DevelopSession { let (w, h) = fit(fw, fh, width.max(1), height.max(1)); let shader = self.graph.compose(); + + // Rasterise the masks first: the shader addresses array slices by + // index, so the array has to describe *this* stack before it is bound. + let masks = self + .rasterise_masks(w, h) + .then(|| self.masks.as_ref().and_then(|p| p.array())) + .flatten(); let texture = self .adjust - .render(&self.demosaiced, &shader, w, h) + .render_masked(&self.demosaiced, &shader, w, h, masks) .map_err(|e| e.to_string())?; // The import is fallible on format and usage only, and both are fixed diff --git a/ui/dr-ui/src/lib.rs b/ui/dr-ui/src/lib.rs index 73c6ac4..cf23dd6 100644 --- a/ui/dr-ui/src/lib.rs +++ b/ui/dr-ui/src/lib.rs @@ -32,6 +32,7 @@ mod library_ui; mod live_style; mod net_runtime; mod presets; +mod segmentation; mod settings_store; mod settings_ui; mod sidecar_cache; diff --git a/ui/dr-ui/src/segmentation.rs b/ui/dr-ui/src/segmentation.rs new file mode 100644 index 0000000..15d829a --- /dev/null +++ b/ui/dr-ui/src/segmentation.rs @@ -0,0 +1,538 @@ +//! Running the segmentation for the image on screen. +//! +//! Joins the two arms into the one thing the develop view needs: a region map +//! whose boundaries came from the watershed and whose grouping came from the +//! model (docs/segmentation.md §5). +//! +//! # This is a precompute, and it is slow +//! +//! Roughly 70 ms of watershed plus ~470 ms of inference plus a full-resolution +//! readback, so it runs **once per image on a worker** and never on the frame +//! path. Every interaction it enables — click a subject, widen a selection, +//! toggle the overlay — reads its cached output and costs a lookup. +//! +//! The consequence for the interface is that segmentation is a thing the user +//! *starts*, and the view must stay usable while it runs. There is no version +//! of this that is fast enough to do silently when a photograph opens. + +use std::sync::Arc; + +use dr_gpu::{GpuContext, LabelField, SegmentOptions, SegmentPass}; +use dr_pipeline::mask::segmentation_signature; +use dr_segment::prior::{self, Membership, PriorOptions}; +use dr_segment::{MergeTree, RegionField}; + +/// How many regions a click's selection starts from. +/// +/// Not the finest level: at the bottom of the ladder a smooth sky bands into +/// strips and a click lands on one of them, which reads as the tool being +/// broken rather than as precision. A few hundred regions over a proxy is +/// roughly "objects and their large parts", which is where a first click +/// should land — and the scroll wheel walks from there in both directions. +pub const DEFAULT_LEVEL: u32 = 300; + +/// One image's region map, ready to select from. +pub struct Segmentation { + /// The watershed's partition, with the semantic prior already applied to + /// its edge weights. + field: RegionField, + /// The merge order over `field`, which is the granularity ladder. + tree: MergeTree, + /// What the model found, for "select the subject" and for the layer names + /// a user actually recognises. + instances: Vec, + /// The labels on the GPU, for the mask rasteriser. + /// + /// Optional, and not merely for testing: the region map is complete and + /// every selection decision is made on the CPU half above, so an upload + /// that fails costs the *rendering* of a mask and not the ability to + /// compute one. It also lets the selection logic be asserted with no + /// adapter present, which is where the behaviour that matters lives. + labels: Option, + /// Identifies this segmentation, so a stored mask can tell whether its + /// region ids still mean what they meant. + signature: u64, + /// The grouping at the current level: region id -> group id. + grouping: Vec, + level: u32, +} + +/// A detected object, reduced to what the interface needs. +/// +/// The full [`dr_segment::Instance`] carries a soft mask the size of the +/// image; keeping several of those alive per photograph is tens of megabytes +/// for data already folded into the region field. The regions it covers are +/// the durable part. +#[derive(Debug, Clone)] +pub struct InstanceSummary { + pub class_name: Arc, + pub score: f32, + /// The regions this instance covers, snapped to watershed boundaries. + pub regions: Vec, +} + +impl Segmentation { + pub fn signature(&self) -> u64 { + self.signature + } + + pub fn labels(&self) -> Option<&LabelField> { + self.labels.as_ref() + } + + pub fn level(&self) -> u32 { + self.level + } + + pub fn region_count(&self) -> usize { + self.field.region_count + } + + pub fn instances(&self) -> &[InstanceSummary] { + &self.instances + } + + /// Move the granularity ladder, recomputing the grouping. + /// + /// Cheap — a union-find over a few thousand nodes — which is what makes + /// the scroll wheel a live control rather than a re-segmentation. + pub fn set_level(&mut self, level: u32) { + let level = level.clamp(2, self.field.region_count.max(2) as u32); + if level == self.level { + return; + } + self.level = level; + self.grouping = self.tree.cut_to(level as usize); + } + + /// The group under a point in **normalised image coordinates**. + /// + /// Normalised because the caller has a click in a widget that is letter- + /// boxed, zoomed and possibly cropped; converting once at the call site is + /// clearer than passing three transforms in here. + pub fn group_at(&self, x: f32, y: f32) -> Option { + if !(0.0..1.0).contains(&x) || !(0.0..1.0).contains(&y) { + return None; + } + let px = ((x * self.field.width as f32) as usize).min(self.field.width - 1); + let py = ((y * self.field.height as f32) as usize).min(self.field.height - 1); + let region = self.field.labels[py * self.field.width + px]; + self.grouping.get(region as usize).copied() + } + + /// Every region belonging to a group — the ids a mask stores. + pub fn regions_in_group(&self, group: u32) -> Vec { + (0..self.field.region_count as u32) + .filter(|&r| self.grouping.get(r as usize) == Some(&group)) + .collect() + } + + /// The regions under a click, at the current level. + pub fn regions_at(&self, x: f32, y: f32) -> Vec { + self.group_at(x, y) + .map(|g| self.regions_in_group(g)) + .unwrap_or_default() + } + + /// A false-coloured RGBA image of the current grouping. + /// + /// The diagnostic that makes the hierarchy arguable instead of a matter of + /// faith — and the same thing a user wants when deciding whether a click + /// will land where they mean. Colours come from a golden-angle walk over + /// hue keyed by group id, so adjacent groups are very unlikely to share + /// one and the same group keeps its colour as the level changes. + pub fn overlay_rgba(&self) -> (Vec, u32, u32) { + let (w, h) = (self.field.width, self.field.height); + let mut out = vec![0u8; w * h * 4]; + + for (p, ®ion) in self.field.labels.iter().enumerate() { + let group = self.grouping.get(region as usize).copied().unwrap_or(0); + let [r, g, b] = group_colour(group); + out[p * 4] = r; + out[p * 4 + 1] = g; + out[p * 4 + 2] = b; + // Opaque here; the view composites it at whatever strength the + // user asked for. Baking transparency in would mean regenerating + // the whole image to change a slider. + out[p * 4 + 3] = 255; + } + + // Boundaries drawn black over the fill. Without them two adjacent + // groups that happen to land on near hues read as one region, which is + // exactly the judgement the overlay exists to support. + for y in 0..h { + for x in 0..w { + let here = self.grouping[self.field.labels[y * w + x] as usize]; + let right = (x + 1 < w) + .then(|| self.grouping[self.field.labels[y * w + x + 1] as usize]); + let down = (y + 1 < h) + .then(|| self.grouping[self.field.labels[(y + 1) * w + x] as usize]); + if right.is_some_and(|g| g != here) || down.is_some_and(|g| g != here) { + let p = (y * w + x) * 4; + out[p..p + 3].copy_from_slice(&[0, 0, 0]); + } + } + } + + (out, w as u32, h as u32) + } +} + +/// A distinct colour per group. +/// +/// Golden-angle hue stepping: consecutive ids land far apart on the wheel, so +/// neighbouring regions — which usually *do* have consecutive ids, since the +/// field is compacted in raster order — are visually separated. Deterministic +/// rather than random, so the same region is the same colour every time the +/// overlay is drawn and the eye can track it across a level change. +fn group_colour(group: u32) -> [u8; 3] { + let h = (group as f32 * 137.508) % 360.0; + // Fixed saturation and value: a fully saturated palette is easier to + // separate than one that also varies in lightness, and the image beneath + // supplies all the tonal variation the eye needs. + let c = 230.0; + let x = c * (1.0 - ((h / 60.0) % 2.0 - 1.0).abs()); + let (r, g, b) = match (h / 60.0) as u32 { + 0 => (c, x, 0.0), + 1 => (x, c, 0.0), + 2 => (0.0, c, x), + 3 => (0.0, x, c), + 4 => (x, 0.0, c), + _ => (c, 0.0, x), + }; + [r as u8 + 25, g as u8 + 25, b as u8 + 25] +} + +/// What to run. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct Options { + pub segment: SegmentOptions, + pub prior: PriorOptions, + /// Whether to run the model at all. + /// + /// Off gives arm A alone: still a full region map with a working + /// granularity ladder, just without the semantic grouping — and about + /// seven times faster. Worth having as a choice rather than a fallback, + /// because on a landscape with no COCO class in it the model contributes + /// nothing and costs the whole inference. + pub semantic: bool, +} + +impl Default for Options { + fn default() -> Self { + Self { + segment: SegmentOptions::default(), + prior: PriorOptions::default(), + semantic: true, + } + } +} + +/// Run the watershed, optionally the model, and combine them. +/// +/// `rgb` is the proxy the model reads — tightly packed RGB floats at +/// `(rgb_width, rgb_height)`. Passed in rather than derived here because the +/// caller already has the decoded image and re-deriving it would mean a second +/// readback of something the CPU is holding. +pub fn compute( + ctx: &GpuContext, + source: &dr_gpu::DemosaicedImage, + rgb: &[f32], + rgb_width: usize, + rgb_height: usize, + options: &Options, +) -> Result { + let pass = SegmentPass::new(ctx).map_err(|e| format!("watershed unavailable: {e}"))?; + let raw = pass + .run(source, options.segment) + .map_err(|e| format!("watershed failed: {e}"))?; + + let field = raw + .read_field() + .map_err(|e| format!("could not read the region field: {e}"))?; + + // Arm B, and its failure is not fatal. A missing or unreadable model + // leaves a perfectly usable watershed segmentation; refusing to segment at + // all because the semantic half is unavailable would trade a working + // feature for a strict one. + let instances = if options.semantic { + match detect(rgb, rgb_width, rgb_height) { + Ok(found) => found, + Err(e) => { + log::warn!("semantic segmentation unavailable, using watershed alone: {e}"); + Vec::new() + } + } + } else { + Vec::new() + }; + + let (field, summaries) = apply_prior(field, &instances, &options.prior); + let tree = MergeTree::build(&field); + let grouping = tree.cut_to(DEFAULT_LEVEL as usize); + + let (w, h) = (field.width as u32, field.height as u32); + let labels = LabelField::upload(ctx, &field.labels, w, h, field.region_count as u32) + .inspect_err(|e| log::warn!("could not upload the label field: {e}")) + .ok(); + + let signature = segmentation_signature( + w, + h, + field.region_count as u32, + tuning_hash(options), + ); + + Ok(Segmentation { + field, + tree, + instances: summaries, + labels, + signature, + grouping, + level: DEFAULT_LEVEL, + }) +} + +/// Arm B. Loads the model each call rather than holding a session open. +/// +/// ~24 ms against the ~470 ms of inference that follows it, and it runs once +/// per image — so caching the session would keep 11 MB of weights resident for +/// the whole life of the app to save five percent of a background task. +fn detect(rgb: &[f32], width: usize, height: usize) -> Result, String> { + let mut model = dr_segment::SemanticModel::embedded().map_err(|e| e.to_string())?; + model + .detect(rgb, width, height, &dr_segment::SemanticOptions::default()) + .map_err(|e| e.to_string()) +} + +/// Arm C: fold the instances into the merge weights, and record what each one +/// selects. +fn apply_prior( + field: RegionField, + instances: &[dr_segment::Instance], + options: &PriorOptions, +) -> (RegionField, Vec) { + let strong: Vec<&dr_segment::Instance> = instances + .iter() + .filter(|i| i.score >= options.confidence) + .collect(); + + if strong.is_empty() { + return (field, Vec::new()); + } + + let masks: Vec<&[f32]> = strong.iter().map(|i| i.mask.as_slice()).collect(); + let Ok(membership) = Membership::compute(&field, &masks, options) else { + // Only reachable if a mask disagrees with the field's size, which + // would be a bug rather than a user-facing condition. Arm A alone is + // the right degradation. + log::warn!("instance masks do not match the region field; using watershed alone"); + return (field, Vec::new()); + }; + + let summaries = strong + .iter() + .map(|i| InstanceSummary { + class_name: i.class_name.clone(), + score: i.score, + regions: prior::regions_for_instance(&field, &i.mask, options), + }) + .filter(|s| !s.regions.is_empty()) + .collect(); + + ( + prior::apply_semantic_prior(&field, &membership, options), + summaries, + ) +} + +/// Everything about the run that changes what a region id means. +/// +/// Folded into the signature so a mask stored under one tuning is not silently +/// reinterpreted under another. Not a general-purpose hash of `Options` — the +/// prior's strength changes the *grouping* but not the underlying regions, and +/// including it would invalidate masks that are still perfectly valid. +fn tuning_hash(options: &Options) -> u64 { + let s = &options.segment; + let mut h: u64 = 0xcbf2_9ce4_8422_2325; + for word in [ + s.max_edge as u64, + s.blur_radius as u64, + s.w_luma.to_bits() as u64, + s.w_chroma.to_bits() as u64, + s.plateau_iterations as u64, + u64::from(options.semantic), + ] { + for byte in word.to_le_bytes() { + h ^= byte as u64; + h = h.wrapping_mul(0x1000_0000_01b3); + } + } + h +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A 4x2 field split down the middle. + /// + /// Regions 0 (top-left) and 1 (bottom-left) make the left half; 2 and 3 + /// the right. The seam between the halves is the strong boundary and the + /// within-half seams are weak, so coarsening merges each half before it + /// ever joins the two — which is the structure a real picture of a subject + /// against a background has, and the one every assertion below depends on. + fn field() -> RegionField { + RegionField { + width: 4, + height: 2, + labels: vec![0, 0, 2, 2, 1, 1, 3, 3], + region_count: 4, + adjacency: vec![ + dr_segment::Edge { a: 0, b: 1, saddle: 0.1 }, + dr_segment::Edge { a: 2, b: 3, saddle: 0.1 }, + dr_segment::Edge { a: 0, b: 2, saddle: 1.0 }, + dr_segment::Edge { a: 1, b: 3, saddle: 1.0 }, + ], + } + } + + /// A segmentation with no GPU field, which is all the selection logic + /// needs — every decision below is made on the CPU half. + fn segmentation(level: u32) -> Segmentation { + let field = field(); + let tree = MergeTree::build(&field); + let grouping = tree.cut_to(level as usize); + Segmentation { + field, + tree, + instances: Vec::new(), + labels: None, + signature: 1, + grouping, + level, + } + } + + #[test] + fn a_click_lands_on_the_group_under_it() { + // Four regions cut to two groups: the field's equal saddles merge in + // (a, b) order, so 0-1 and 2-3 pair up — left half and right half. + let seg = segmentation(2); + + let left = seg.group_at(0.1, 0.5).expect("inside the frame"); + let right = seg.group_at(0.9, 0.5).expect("inside the frame"); + assert_ne!(left, right, "the two halves are different groups"); + + // The same group whichever row is clicked, which is what "region" + // means and what a per-pixel lookup would not give. + assert_eq!(seg.group_at(0.1, 0.1), Some(left)); + assert_eq!(seg.group_at(0.1, 0.9), Some(left)); + } + + #[test] + fn a_click_outside_the_frame_selects_nothing() { + let seg = segmentation(2); + assert_eq!(seg.group_at(-0.1, 0.5), None); + assert_eq!(seg.group_at(1.5, 0.5), None); + assert_eq!(seg.group_at(0.5, 1.0), None, "the far edge is exclusive"); + } + + #[test] + fn selecting_a_group_returns_every_region_in_it() { + let seg = segmentation(2); + let regions = seg.regions_at(0.1, 0.5); + assert_eq!(regions.len(), 2, "each half is two merged regions"); + assert!(regions.windows(2).all(|w| w[0] < w[1]), "sorted for the mask"); + } + + /// The scroll wheel. Coarsening must not re-segment, and must actually + /// change what a click selects. + #[test] + fn the_level_changes_what_a_click_selects() { + let mut seg = segmentation(4); + assert_eq!(seg.regions_at(0.1, 0.5).len(), 1, "finest: one region"); + + seg.set_level(2); + assert_eq!(seg.level(), 2); + assert_eq!(seg.regions_at(0.1, 0.5).len(), 2, "coarser: the pair"); + } + + #[test] + fn the_level_is_clamped_to_what_exists() { + let mut seg = segmentation(4); + seg.set_level(9999); + assert!(seg.level() <= seg.region_count() as u32); + + seg.set_level(0); + assert!(seg.level() >= 2, "one group would select the whole frame"); + } + + #[test] + fn the_overlay_covers_every_pixel_and_draws_boundaries() { + let seg = segmentation(2); + let (rgba, w, h) = seg.overlay_rgba(); + assert_eq!((w, h), (4, 2)); + assert_eq!(rgba.len(), 4 * 2 * 4); + assert!(rgba.chunks_exact(4).all(|p| p[3] == 255), "fully opaque"); + + // The vertical seam between the halves must be drawn. + let at = |x: usize, y: usize| { + let p = (y * 4 + x) * 4; + [rgba[p], rgba[p + 1], rgba[p + 2]] + }; + assert_eq!(at(1, 0), [0, 0, 0], "the boundary column is black"); + assert_ne!(at(0, 0), [0, 0, 0], "the interior is not"); + } + + #[test] + fn the_two_halves_get_different_colours() { + let seg = segmentation(2); + let (rgba, _, _) = seg.overlay_rgba(); + let left = [rgba[0], rgba[1], rgba[2]]; + let right = { + let p = 3 * 4; + [rgba[p], rgba[p + 1], rgba[p + 2]] + }; + assert_ne!(left, right); + } + + #[test] + fn colours_are_stable_and_distinct() { + assert_eq!(group_colour(7), group_colour(7), "same id, same colour"); + assert_ne!(group_colour(0), group_colour(1)); + assert_ne!(group_colour(1), group_colour(2)); + } + + #[test] + fn every_colour_is_visible_against_the_image() { + // The overlay is composited over a photograph, so a colour that is + // nearly black is a region the user cannot see they have selected. + for id in 0..64u32 { + let [r, g, b] = group_colour(id); + let brightest = r.max(g).max(b); + assert!(brightest >= 200, "group {id} is too dark: {r},{g},{b}"); + } + } + + #[test] + fn tuning_changes_the_signature_but_the_prior_does_not() { + let base = Options::default(); + + let mut coarser = base; + coarser.segment.max_edge = 800; + assert_ne!(tuning_hash(&base), tuning_hash(&coarser), "proxy size"); + + let mut no_model = base; + no_model.semantic = false; + assert_ne!(tuning_hash(&base), tuning_hash(&no_model), "arm B on/off"); + + // The prior reweights the merge order over the *same* regions, so a + // mask's ids still mean what they meant. Invalidating them here would + // discard work for no reason. + let mut stronger = base; + stronger.prior.strength = 0.1; + assert_eq!(tuning_hash(&base), tuning_hash(&stronger)); + } +}