diff --git a/apps/darkroom-desktop/Cargo.toml b/apps/darkroom-desktop/Cargo.toml index 238b7f6..a97b2ae 100644 --- a/apps/darkroom-desktop/Cargo.toml +++ b/apps/darkroom-desktop/Cargo.toml @@ -6,7 +6,7 @@ rust-version.workspace = true license.workspace = true [dependencies] -dr-ui.workspace = true +dr-ui = { workspace = true, features = ["scene-model"] } anyhow.workspace = true env_logger.workspace = true log.workspace = true diff --git a/core/dr-gpu/src/mask.rs b/core/dr-gpu/src/mask.rs index 3c4aab0..a5c6d46 100644 --- a/core/dr-gpu/src/mask.rs +++ b/core/dr-gpu/src/mask.rs @@ -654,7 +654,13 @@ impl MaskPass { // 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 { .. } => { + // A category shares the subject's mode, and that is not a + // shortcut: both arrive as a soft coverage buffer at proxy + // resolution and both are turned into a distance field before they + // reach here. The shader has no way to tell them apart and no + // reason to want one — what differs is only which model produced + // the coverage. + MaskSource::Subject { .. } | MaskSource::Category { .. } => { let short = field_short_edge(width, height); MaskParams { mode: MODE_SUBJECT, diff --git a/ui/dr-ui/Cargo.toml b/ui/dr-ui/Cargo.toml index 03db558..8d2aded 100644 --- a/ui/dr-ui/Cargo.toml +++ b/ui/dr-ui/Cargo.toml @@ -114,6 +114,15 @@ slint-build.workspace = true serde_norway.workspace = true [features] +# Compile the scene model into the binary. +# +# On by default for the desktop app and *off* for Android, which unpacks the +# same graph from APK assets instead — 24 MB of constant is worth avoiding in a +# mobile install and not worth the plumbing to avoid on a desktop one. Without +# it `segmentation::scene_categories` falls back to the installed-file lookup, +# which is what a packaged desktop build uses too. +scene-model = ["dr-segment/embedded-scene-model"] + default = [] # Debug convenience: re-read style.yaml at startup so a palette can be tuned # without rebuilding. Costs the constant-folding of every token, so it stays diff --git a/ui/dr-ui/src/develop.rs b/ui/dr-ui/src/develop.rs index cd4bf25..7fa898c 100644 --- a/ui/dr-ui/src/develop.rs +++ b/ui/dr-ui/src/develop.rs @@ -1795,6 +1795,23 @@ impl DevelopSession { MaskSource::Subject { index, .. } => { mix(1); mix(*index as u64); + if layer.morphology.is_compound() { + mix(match layer.morphology { + Morphology::Close => 2, + Morphology::Open => 3, + _ => 0, + }); + mix(layer.morph_radius.to_bits() as u64); + } + } + // Hashed by *name*, and `4` rather than `1` so a category + // named the same as an instance index could never collide with + // it. The field has to be rebuilt when either changes. + MaskSource::Category { name, .. } => { + mix(4); + for b in name.as_bytes() { + mix(*b as u64); + } // Only the compound operations change the field itself. if layer.morphology.is_compound() { mix(match layer.morphology { @@ -1830,6 +1847,20 @@ impl DevelopSession { let mut fields: Vec> = Vec::new(); for layer in self.graph.masks().active() { let field = match &layer.source { + MaskSource::Category { name, .. } => seg + .category_mask(name) + .map(|coverage| { + dr_segment::Shaped::build( + coverage, + pw, + ph, + 128, + morphology_for(layer.morphology), + layer.morph_radius * pw.min(ph) as f32, + ) + .distance + }) + .unwrap_or_default(), MaskSource::Subject { index, .. } => seg .instance_mask(*index as usize) .map(|coverage| { @@ -2594,6 +2625,54 @@ impl DevelopSession { self.add_subject_mask(index) } + /// Add a layer covering one photographic category. + /// + /// The counterpart to [`Self::add_subject_mask`], and it takes a *name* + /// rather than an index for the reason `MaskSource::Category` stores one: + /// the descriptor grouping ADE20K's classes is editable, so an index would + /// silently repoint every stored layer the first time a category was + /// added to it. + /// + /// The layer is named after the category, because "sky" is a better name + /// for a layer than "Mask 3" and the user can rename it anyway. + pub fn add_category_mask(&mut self, name: &str) -> Option { + use dr_pipeline::mask::MaskSource; + + let seg = self.segmentation.as_ref()?; + let signature = seg.signature(); + // Refuse a category this run did not produce rather than creating a + // layer that renders empty: an empty mask looks like a broken + // adjustment, where a button that does nothing at least says so. + if seg.category_mask(name).is_none() { + return None; + } + + let id = self.graph.masks().next_id(); + let mut layer = MaskLayer::new( + id.clone(), + MaskSource::Category { + signature, + name: name.to_string(), + }, + ); + layer.name = name.to_string(); + if !self.graph.masks_mut().push(layer) { + return None; + } + self.active_masks = vec![id.clone()]; + self.history + .record(&self.graph, Edit::Action(labels::step::MASK_ADDED)); + Some(id) + } + + /// What the scene model found in this frame, largest category first. + /// + /// Empty when no scene model was available, which is an ordinary state — + /// see `segmentation::scene_categories`. + pub fn categories(&self) -> &[segmentation::CategorySummary] { + self.segmentation.as_ref().map_or(&[], |s| s.categories()) + } + /// Add a layer selecting one detected subject. /// /// The instance's own coverage is the mask, rather than the watershed @@ -2768,7 +2847,9 @@ impl DevelopSession { self.graph.masks().get(id).is_some_and(|l| { matches!( l.source, - MaskSource::Subject { .. } | MaskSource::Regions { .. } + MaskSource::Subject { .. } + | MaskSource::Category { .. } + | MaskSource::Regions { .. } ) }) } diff --git a/ui/dr-ui/src/segmentation.rs b/ui/dr-ui/src/segmentation.rs index cb19c1d..013238d 100644 --- a/ui/dr-ui/src/segmentation.rs +++ b/ui/dr-ui/src/segmentation.rs @@ -36,6 +36,23 @@ use dr_gpu::GpuContext; use dr_pipeline::mask::segmentation_signature; use dr_types::Orientation; +/// One photographic category, over the whole frame. +/// +/// The counterpart to [`InstanceSummary`] and deliberately thinner: a category +/// has no box, because it is not one object in one place — sky is wherever the +/// sky is, in as many disconnected pieces as the frame has windows. +#[derive(Debug, Clone)] +pub struct CategorySummary { + pub name: std::sync::Arc, + /// Fraction of the frame this category covers, for ordering the list and + /// for hiding a category that would give the user a control that does + /// nothing. + pub coverage: f32, + /// Coverage at proxy resolution, quantised to a byte — the same + /// representation, and for the same reasons, as `InstanceSummary::mask`. + pub mask: Vec, +} + /// One recognised object. #[derive(Debug, Clone)] pub struct InstanceSummary { @@ -61,6 +78,14 @@ pub struct InstanceSummary { /// One image's recognised objects, ready to mask. pub struct Segmentation { instances: Vec, + /// What the scene model made of the same frame, empty when no scene model + /// could be found. + /// + /// Empty is an ordinary state, not a failure: a build without the weights + /// compiled in and without them installed simply offers no categories, the + /// same way a missing face model turns face indexing off rather than + /// stopping the app. + categories: Vec, /// Identifies this run, so a stored layer can tell whether the index it /// holds still means what it meant. signature: u64, @@ -90,6 +115,21 @@ impl Segmentation { self.instances.get(index).map(|i| i.mask.as_slice()) } + pub fn categories(&self) -> &[CategorySummary] { + &self.categories + } + + /// One category's coverage, at [`Self::proxy_size`]. + /// + /// By name, matching `MaskSource::Category`. A linear scan because there + /// are eight of them and a map would be more machinery than lookup. + pub fn category_mask(&self, name: &str) -> Option<&[u8]> { + self.categories + .iter() + .find(|c| &*c.name == name) + .map(|c| c.mask.as_slice()) + } + /// Replace one instance in place, keeping every other index and the /// signature unchanged. /// @@ -309,8 +349,21 @@ pub fn compute( ^ orientation_key(orientation), ); + // The scene pass, on the same upright frame and laid back down the same + // way. Failures here are logged and dropped rather than propagated: no + // scene model is an ordinary state, and a photograph that can be masked by + // subject should not become unopenable because the categories are absent. + let categories = match scene_categories(&stood_up, uw, uh, orientation) { + Ok(c) => c, + Err(e) => { + log::info!("no scene categories for this frame: {e}"); + Vec::new() + } + }; + Ok(Segmentation { instances, + categories, signature, // **Sensor space, not the model's.** `lay_down` put every mask back, // so the grid a stored layer indexes into is the one it always was — @@ -475,6 +528,86 @@ fn detect( .map_err(|e| e.to_string()) } +/// Weigh the photographic categories, if this build can find a scene model. +/// +/// Separate from [`detect`] rather than folded into it because the two are +/// independent: a build with no scene model still segments subjects, and a +/// frame with no recognisable subject still has sky. Neither failure should +/// take the other down. +fn scene_categories( + rgb: &[f32], + width: usize, + height: usize, + orientation: Orientation, +) -> Result, String> { + let mut model = load_scene_model()?; + let scene = model + .analyse(rgb, width, height) + .map_err(|e| e.to_string())?; + + let mut out = Vec::new(); + for (index, name) in scene.categories().iter().enumerate() { + let coverage = scene.coverage(index); + // A category the model barely saw is not worth a mask buffer the size + // of the proxy, and offering it in the list would be offering a + // control that does nothing when moved. The threshold is the one the + // example prints against. + if coverage < 0.005 { + continue; + } + let Some(weights) = scene.rasterise(index, width, height) else { + continue; + }; + // Back into sensor space, exactly as an instance mask is: the grid a + // stored layer indexes into has to be the sensor's whatever the model + // was shown. `lay_down` wants a box too, so it gets the whole frame — + // a category has no meaningful extent. + let (mask, _) = lay_down( + &weights, + (0.0, 0.0, width as f32, height as f32), + width, + height, + orientation, + ); + out.push(CategorySummary { + name: name.clone(), + coverage, + mask: quantise(&mask), + }); + } + + // Largest first, which is the order the list is worth reading in. + out.sort_by(|a, b| b.coverage.total_cmp(&a.coverage)); + Ok(out) +} + +/// Find a scene model: compiled in if this build has one, installed otherwise. +/// +/// The order matters. A build with the weights compiled in should not be +/// silently overridden by a stale file in a data directory, and a build +/// without them has nothing to fall back *from* — so "embedded, then +/// installed" is the only ordering that is not surprising either way. +fn load_scene_model() -> Result { + #[cfg(feature = "scene-model")] + { + return dr_segment::SceneModel::embedded().map_err(|e| e.to_string()); + } + #[cfg(not(feature = "scene-model"))] + { + // Account-independent, like `shared_face_models_dir` and for the same + // reason: this runs on a worker with no session in hand. Android + // unpacks the APK's copy to exactly this directory before any store + // opens. + let dir = crate::library::shared_face_models_dir(); + dr_segment::SceneModel::from_path( + dir.join("yolo26s-sem-ade20k.onnx"), + dir.join("yolo26s-sem-ade20k.classes.json"), + dir.join("categories.txt"), + ) + .map_err(|e| e.to_string()) + } +} + /// The model's soft coverage, to a byte per pixel. /// /// Rounded rather than truncated, so a coverage of exactly 0.5 lands on the @@ -519,6 +652,7 @@ mod tests { bbox: (4.0, 0.0, 8.0, h as f32), }, ], + categories: Vec::new(), signature: 1, proxy: (w, h), } @@ -555,6 +689,7 @@ mod tests { fn a_click_on_nothing_selects_nothing() { let seg = Segmentation { instances: Vec::new(), + categories: Vec::new(), signature: 1, proxy: (4, 4), }; @@ -582,6 +717,7 @@ mod tests { fn the_overlay_is_transparent_where_nothing_was_found() { let seg = Segmentation { instances: Vec::new(), + categories: Vec::new(), signature: 1, proxy: (4, 4), };