DarkRoom/ui/dr-ui/src/segmentation.rs

//! Finding what can be selected in the photograph on screen.
//!
//! One model pass, and what it recognised. That is the whole subsystem now.
//!
//! # What used to be here
//!
//! A watershed over-segmented the image, a merge tree turned that into a
//! granularity ladder, and a click walked up it (docs/segmentation.md, arms A
//! and C). It is gone from this path, and the reason is measured rather than
//! aesthetic: on a real photograph the saddles are near zero almost
//! everywhere, so the merge order joins everything meaningful before it joins
//! anything spurious. Cutting a 45,808-basin field of a 22 MP frame to 2,000
//! regions left **one** region covering nearly the whole picture plus specks
//! (§15). A ladder that collapses is not a ladder.
//!
//! The passes and the hierarchy still exist in `dr-gpu` and `dr-segment`,
//! tested and documented, because it is the *merge criterion* that fails and
//! that is one function. What does not exist any more is a product path
//! through them, or a control offering a choice that does nothing.
//!
//! # This is a precompute, and it is slow
//!
//! ~700 ms on a 22 MP frame: a proxy render, a readback, and the model. It
//! runs **once per image, when the user asks**, and never on the frame path.
//! Every interaction it enables — click a subject, grow a mask, change a
//! falloff — reads its cached output.
//!
//! All of it is on a worker. [`compute`] takes an owned buffer and a
//! `GpuContext`, which is what makes that possible — nothing here touches the
//! develop session, and [`crate::develop::SegmentationJob`] is the piece that
//! carries the proxy render across with it.

use std::sync::Arc;

use dr_gpu::GpuContext;
use dr_pipeline::mask::segmentation_signature;

/// One recognised object.
#[derive(Debug, Clone)]
pub struct InstanceSummary {
    pub class_name: Arc<str>,
    pub score: f32,
    /// Coverage at proxy resolution, quantised to a byte.
    ///
    /// A byte rather than the `f32` the model produces: 256 levels is far
    /// finer than an edge anyone can see, and at four bytes a pixel a handful
    /// of objects would be most of a hundred megabytes for one photograph.
    ///
    /// This is the *source* a distance field is built from, not the mask
    /// itself — `dr_segment::Shaped` turns it into one.
    pub mask: Vec<u8>,
    /// `(x0, y0, x1, y1)` in [`Segmentation::proxy_size`] pixels.
    ///
    /// Carried through from `dr_segment::Instance` rather than re-derived
    /// from the mask, so a refine pass knows what region to crop without
    /// scanning a proxy-sized buffer for its own extent.
    pub bbox: (f32, f32, f32, f32),
}

/// One image's recognised objects, ready to mask.
pub struct Segmentation {
    instances: Vec<InstanceSummary>,
    /// Identifies this run, so a stored layer can tell whether the index it
    /// holds still means what it meant.
    signature: u64,
    /// The space instance masks are defined in, in **source** proxy pixels.
    ///
    /// Everything a mask is built from lives here, which is what lets the
    /// render sample it *after* the framing map rather than before — so a mask
    /// stays on the photograph through a zoom, a pan and a crop.
    proxy: (usize, usize),
}

impl Segmentation {
    pub fn signature(&self) -> u64 {
        self.signature
    }

    pub fn proxy_size(&self) -> (usize, usize) {
        self.proxy
    }

    pub fn instances(&self) -> &[InstanceSummary] {
        &self.instances
    }

    /// One instance's coverage, at [`Self::proxy_size`].
    pub fn instance_mask(&self, index: usize) -> Option<&[u8]> {
        self.instances.get(index).map(|i| i.mask.as_slice())
    }

    /// Replace one instance in place, keeping every other index and the
    /// signature unchanged.
    ///
    /// What a refine pass calls once it has a sharper mask for the subject at
    /// `index`: the layers pointing at this run by index still mean what they
    /// meant, they just read better pixels now.
    pub fn replace_instance(&mut self, index: usize, instance: InstanceSummary) {
        if let Some(slot) = self.instances.get_mut(index) {
            *slot = instance;
        }
    }

    /// The strongest instance covering a point in normalised image
    /// coordinates.
    ///
    /// Strongest rather than smallest: detections are score-ordered and
    /// overlapping ones are usually the same object found twice, so the more
    /// confident is the better guess. A person in front of a bus wins over the
    /// bus, because the person's mask is the one under the cursor at all.
    pub fn instance_at(&self, x: f32, y: f32) -> Option<usize> {
        if !(0.0..1.0).contains(&x) || !(0.0..1.0).contains(&y) {
            return None;
        }
        let (w, h) = self.proxy;
        if w == 0 || h == 0 {
            return None;
        }
        let px = ((x * w as f32) as usize).min(w - 1);
        let py = ((y * h as f32) as usize).min(h - 1);
        let p = py * w + px;

        self.instances
            .iter()
            .enumerate()
            .filter(|(_, i)| i.mask.get(p).is_some_and(|&v| v >= 128))
            .max_by(|(_, a), (_, b)| a.score.total_cmp(&b.score))
            .map(|(i, _)| i)
    }

    /// A false-coloured picture of what a click can select, in source space.
    ///
    /// **Transparent where nothing is selectable.** The region map this
    /// replaced covered every pixel and so hid the photograph it was drawn
    /// over; the question an overlay exists to answer is whether an outline
    /// follows the subject, and that can only be answered by seeing both.
    pub fn overlay_rgba(&self) -> (Vec<u8>, u32, u32) {
        let (w, h) = self.proxy;
        let mut out = vec![0u8; w * h * 4];

        // Weakest first, so where two detections overlap the more confident
        // one is the colour on top — matching which a click would select.
        let mut order: Vec<usize> = (0..self.instances.len()).collect();
        order.sort_by(|&a, &b| self.instances[a].score.total_cmp(&self.instances[b].score));

        for &i in &order {
            let [r, g, b] = instance_colour(i as u32);
            for (p, &cov) in self.instances[i].mask.iter().enumerate() {
                if cov < 128 || p * 4 + 3 >= out.len() {
                    continue;
                }
                out[p * 4] = r;
                out[p * 4 + 1] = g;
                out[p * 4 + 2] = b;
                out[p * 4 + 3] = 255;
            }
        }

        // The outline drawn opaque white over the fill. It is the part being
        // judged — a fill can look right while its edge sits several pixels
        // off the subject — and it survives the low opacity the fill is
        // composited at.
        let solid = |p: usize| out.get(p * 4 + 3).is_some_and(|&a| a > 0);
        let mut edges = Vec::new();
        for y in 0..h {
            for x in 0..w {
                let p = y * w + x;
                if !solid(p) {
                    continue;
                }
                let boundary = (x + 1 == w || !solid(p + 1))
                    || (x == 0 || !solid(p - 1))
                    || (y + 1 == h || !solid(p + w))
                    || (y == 0 || !solid(p - w));
                if boundary {
                    edges.push(p);
                }
            }
        }
        for p in edges {
            out[p * 4..p * 4 + 4].copy_from_slice(&[255, 255, 255, 255]);
        }

        (out, w as u32, h as u32)
    }
}

/// A distinct colour per instance.
///
/// Golden-angle hue stepping, deterministic rather than random: the same
/// object is the same colour every time the overlay is drawn, so the eye can
/// track it while a mask is being shaped.
fn instance_colour(index: u32) -> [u8; 3] {
    let h = (index as f32 * 137.508) % 360.0;
    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 {
    /// Detections below this are dropped.
    ///
    /// Deliberately low. A weak detection costs a spurious entry in a list the
    /// user is choosing from, and its score is shown beside it; a missed one
    /// costs a subject that cannot be selected at all, which is the worse
    /// failure for a selection tool.
    pub confidence: f32,
    /// TRACES: FR-DEV-3
    /// Run the model over overlapping tiles instead of the whole frame once.
    ///
    /// Off by default and deliberately so. The graph's input is fixed at
    /// 640x640 (docs/segmentation.md F6), so every image is letterboxed into
    /// it and a subject 200px across in a 1600px proxy reaches the model at
    /// 80px — which is where a coarse outline comes from. Tiling is the only
    /// route to more resolution with a fixed window, and it costs one
    /// inference per tile: about 2.8s for a 3x2 grid against 470ms whole-frame.
    ///
    /// Six times the wait is the wrong default for the common case, where the
    /// subject is large in frame and whole-frame inference is already the best
    /// answer. It is the right answer for a bird against sky, so it is offered
    /// per-image rather than chosen once for all of them.
    pub fine: bool,
}

impl Default for Options {
    fn default() -> Self {
        Self {
            confidence: 0.30,
            fine: false,
        }
    }
}

/// Find what can be selected in this photograph.
///
/// `rgb` is the proxy the model reads — tightly packed RGB floats at
/// `(width, height)`. Passed in rather than derived here because the caller
/// already has the rendered proxy, and re-deriving it would mean a second
/// readback of something the CPU is holding.
pub fn compute(
    _ctx: &GpuContext,
    rgb: &[f32],
    width: usize,
    height: usize,
    options: &Options,
) -> Result<Segmentation, String> {
    let found = detect(rgb, width, height, options.fine)?;

    let instances: Vec<InstanceSummary> = found
        .iter()
        .filter(|i| i.score >= options.confidence)
        .map(|i| InstanceSummary {
            class_name: i.class_name.clone(),
            score: i.score,
            mask: quantise(&i.mask),
            bbox: i.bbox,
        })
        .collect();

    let signature = segmentation_signature(
        width as u32,
        height as u32,
        instances.len() as u32,
        // The tiling choice belongs in the signature as much as the
        // confidence does. A mask stores the signature of the segmentation its
        // region ids index into (`MaskSource::Regions`), and a tiled run finds
        // different instances in a different order — so if the two runs shared
        // a signature, a layer built against the coarse pass would be silently
        // reinterpreted against the fine one. That is a *wrong* mask, which is
        // far worse than a stale one, because nothing announces it.
        options.confidence.to_bits() as u64
            ^ if options.fine {
                0x9E37_79B9_7F4A_7C15
            } else {
                0
            },
    );

    Ok(Segmentation {
        instances,
        signature,
        proxy: (width, height),
    })
}

/// Load the model and run it.
///
/// Loading is ~24 ms against the ~470 ms of inference that follows, and this
/// runs once per image — so caching the session would keep 11 MB of weights
/// resident for the life of the app to save five percent of a background task.
fn detect(
    rgb: &[f32],
    width: usize,
    height: usize,
    fine: bool,
) -> Result<Vec<dr_segment::Instance>, String> {
    let mut model = dr_segment::SemanticModel::embedded().map_err(|e| e.to_string())?;
    let options = dr_segment::SemanticOptions {
        // A quarter shared with each neighbour. It has to exceed zero at all,
        // or a subject sitting on a seam is cut in half by both tiles and
        // recognised by neither; a quarter is enough to carry a whole subject
        // inside one tile at the sizes tiling is reached for.
        tiling: if fine {
            dr_segment::Tiling::Grid { overlap: 0.25 }
        } else {
            dr_segment::Tiling::Whole
        },
        ..dr_segment::SemanticOptions::default()
    };
    model
        .detect(rgb, width, height, &options)
        .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
/// threshold the selection tests against instead of one below it.
fn quantise(mask: &[f32]) -> Vec<u8> {
    mask.iter()
        .map(|&v| (v.clamp(0.0, 1.0) * 255.0).round() as u8)
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Two objects: a big weak one on the left, a small strong one that
    /// overlaps it.
    fn overlapping() -> Segmentation {
        let (w, h) = (8usize, 4usize);
        let mut big = vec![0u8; w * h];
        let mut small = vec![0u8; w * h];
        for y in 0..h {
            for x in 0..6 {
                big[y * w + x] = 255;
            }
            for x in 4..8 {
                small[y * w + x] = 255;
            }
        }

        Segmentation {
            instances: vec![
                InstanceSummary {
                    class_name: "bus".into(),
                    score: 0.5,
                    mask: big,
                    bbox: (0.0, 0.0, 6.0, h as f32),
                },
                InstanceSummary {
                    class_name: "person".into(),
                    score: 0.9,
                    mask: small,
                    bbox: (4.0, 0.0, 8.0, h as f32),
                },
            ],
            signature: 1,
            proxy: (w, h),
        }
    }

    #[test]
    fn a_click_finds_the_object_under_it() {
        let seg = overlapping();
        assert_eq!(seg.instance_at(0.1, 0.5), Some(0), "only the bus here");
        assert_eq!(seg.instance_at(0.95, 0.5), Some(1), "only the person here");
    }

    /// The overlap rule, and the one that decides what a click means where two
    /// detections cover the same pixel.
    #[test]
    fn overlapping_objects_resolve_to_the_more_confident() {
        let seg = overlapping();
        assert_eq!(
            seg.instance_at(0.6, 0.5),
            Some(1),
            "the person at 0.9 beats the bus at 0.5"
        );
    }

    #[test]
    fn a_click_outside_the_frame_selects_nothing() {
        let seg = overlapping();
        assert_eq!(seg.instance_at(-0.1, 0.5), None);
        assert_eq!(seg.instance_at(1.5, 0.5), None);
        assert_eq!(seg.instance_at(0.5, 1.0), None, "the far edge is exclusive");
    }

    #[test]
    fn a_click_on_nothing_selects_nothing() {
        let seg = Segmentation {
            instances: Vec::new(),
            signature: 1,
            proxy: (4, 4),
        };
        assert_eq!(seg.instance_at(0.5, 0.5), None);
    }

    #[test]
    fn colours_are_stable_and_distinct() {
        assert_eq!(instance_colour(7), instance_colour(7));
        assert_ne!(instance_colour(0), instance_colour(1));
        assert_ne!(instance_colour(1), instance_colour(2));
    }

    #[test]
    fn every_colour_is_visible_against_a_photograph() {
        for i in 0..64u32 {
            let [r, g, b] = instance_colour(i);
            assert!(r.max(g).max(b) >= 200, "instance {i} is too dark");
        }
    }

    /// The overlay must not cover the picture: that is the difference between
    /// this and the region map it replaced.
    #[test]
    fn the_overlay_is_transparent_where_nothing_was_found() {
        let seg = Segmentation {
            instances: Vec::new(),
            signature: 1,
            proxy: (4, 4),
        };
        let (px, w, h) = seg.overlay_rgba();
        assert_eq!((w, h), (4, 4));
        assert!(px.chunks_exact(4).all(|p| p[3] == 0), "nothing to draw");
    }

    #[test]
    fn the_overlay_outlines_what_it_fills() {
        let seg = overlapping();
        let (px, w, _) = seg.overlay_rgba();
        let at = |x: usize, y: usize| {
            let p = (y * w as usize + x) * 4;
            [px[p], px[p + 1], px[p + 2], px[p + 3]]
        };
        // The rightmost column of the person is against the frame edge, so it
        // is an outline pixel.
        assert_eq!(at(7, 1), [255, 255, 255, 255]);
        // And an interior pixel keeps its fill.
        assert_ne!(at(2, 1)[3], 0, "the bus is filled");
        assert_ne!(at(2, 1), [255, 255, 255, 255], "and not all outline");
    }

    #[test]
    fn quantising_rounds_rather_than_truncates() {
        // Exactly half must reach the threshold the selection tests against.
        assert_eq!(quantise(&[0.5]), vec![128]);
        assert_eq!(quantise(&[0.0, 1.0]), vec![0, 255]);
        // And values outside the range cannot wrap.
        assert_eq!(quantise(&[-1.0, 2.0]), vec![0, 255]);
    }

    #[test]
    fn confidence_changes_the_signature() {
        // A different threshold is a different instance list, so the indices a
        // stored layer holds mean something else.
        let a = segmentation_signature(100, 100, 3, Options::default().confidence.to_bits() as u64);
        let b = segmentation_signature(100, 100, 3, 0.9f32.to_bits() as u64);
        assert_ne!(a, b);
    }
}