DarkRoom/core/dr-segment/src/prior.rs

//! Arm C — semantic instances as a prior over the watershed merge order.
//!
//! docs/segmentation.md §5. The spec calls this the expected winner and it is
//! what ships, for a reason that survives the model turning out to be narrower
//! than §4 assumed: the two arms fail in *opposite* directions, so each one
//! covers the other's failure.
//!
//! - The watershed knows where every edge is and nothing about what it
//!   separates. Its coarse levels are geometric accidents — level 7 is *a*
//!   coarser partition, not *the* object.
//! - The model knows a dog is a dog and puts the dog's outline roughly where
//!   the dog is, at quarter resolution, with a soft edge.
//!
//! Weighting the merge by semantic agreement takes the outline from the
//! watershed and the grouping from the model. A region pair the model believes
//! belongs to one object merges early; a pair straddling an object's edge
//! merges late. **No boundary moves** — only the order in which boundaries
//! dissolve — which is why the result is pixel-accurate at every level while
//! its coarse levels are named things.
//!
//! # Why this is not "just use the model's mask"
//!
//! Because a mask edge is judged at 100% zoom, and the model's edge is a
//! quarter-resolution sigmoid. Using the instance mask directly gives a
//! selection that is semantically right and visibly soft — acceptable for
//! biasing, not acceptable as the mask itself. Snapping to watershed regions
//! ([`regions_for_instance`]) gives the same selection with the sensor's own
//! edges.

use std::collections::HashMap;

use crate::hierarchy::{Edge, RegionField};

/// How strongly the model is allowed to reorder the merge.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct PriorOptions {
    /// How far a confident semantic judgement may scale a saddle, `0.0..1.0`.
    ///
    /// At `0.0` this is arm A exactly. At `0.9` a pair the model is certain
    /// shares an object merges at a tenth of its true boundary strength.
    ///
    /// Not `1.0`, and the ceiling is the point: at `1.0` an agreeing pair gets
    /// a saddle of zero and merges *before* genuinely identical neighbours,
    /// which lets a confident-but-wrong detection flatten real structure it
    /// happens to cover. Leaving headroom keeps the image's own evidence able
    /// to outvote the model.
    pub strength: f32,

    /// Coverage above which a region counts as belonging to an instance.
    ///
    /// Applied to the *mean* of the instance's soft mask over the region, so
    /// this is "most of this region is inside the dog", not "some pixel is".
    pub membership: f32,

    /// Instances scoring below this contribute no prior at all.
    ///
    /// Higher than the detection threshold on purpose. A weak detection is
    /// still worth *offering* in a list a person picks from, where the cost of
    /// being wrong is an ignored entry — but not worth silently reshaping the
    /// hierarchy every other interaction depends on.
    pub confidence: f32,
}

impl Default for PriorOptions {
    fn default() -> Self {
        Self {
            strength: 0.75,
            membership: 0.5,
            confidence: 0.4,
        }
    }
}

/// Which instance, if any, each region belongs to.
///
/// One dominant instance per region rather than a vector of memberships:
/// regions are small — a proxy watershed makes thousands of them — and a
/// region spanning two objects means the watershed already failed there, which
/// is a case to leave to the image gradient rather than to average over.
#[derive(Debug, Clone, PartialEq)]
pub struct Membership {
    /// Per region: the instance index it mostly belongs to, and how much.
    pub of: Vec<Option<(usize, f32)>>,
}

impl Membership {
    /// Compute per-region membership from soft instance masks.
    ///
    /// `masks` is one full-image coverage buffer per instance, each
    /// `width * height` — [`crate::semantic::Instance::mask`] is exactly this,
    /// passed as slices so that this function needs no `semantic` feature and
    /// stays testable with hand-written masks.
    pub fn compute(
        field: &RegionField,
        masks: &[&[f32]],
        options: &PriorOptions,
    ) -> Result<Self, MembershipError> {
        let pixels = field.width * field.height;
        for (i, mask) in masks.iter().enumerate() {
            if mask.len() != pixels {
                return Err(MembershipError::MaskSize {
                    instance: i,
                    expected: pixels,
                    got: mask.len(),
                });
            }
        }

        // Sum coverage per (region, instance), then divide by region size.
        let mut totals = vec![0.0f32; field.region_count * masks.len()];
        let mut sizes = vec![0u32; field.region_count];

        for (p, &label) in field.labels.iter().enumerate() {
            let r = label as usize;
            sizes[r] += 1;
            for (i, mask) in masks.iter().enumerate() {
                totals[r * masks.len() + i] += mask[p];
            }
        }

        let of = (0..field.region_count)
            .map(|r| {
                let size = sizes[r].max(1) as f32;
                let row = &totals[r * masks.len()..(r + 1) * masks.len()];

                // `total_cmp` and an index tiebreak: two instances covering a
                // region equally must resolve the same way on every machine,
                // because the merge order below is derived from this and an
                // unstable merge order is an unstable label field (M5).
                let best = row
                    .iter()
                    .enumerate()
                    .max_by(|(ia, a), (ib, b)| a.total_cmp(b).then(ib.cmp(ia)))?;

                let coverage = best.1 / size;
                (coverage >= options.membership).then_some((best.0, coverage))
            })
            .collect();

        Ok(Self { of })
    }

    /// The instance a region belongs to, if any.
    pub fn instance_of(&self, region: u32) -> Option<usize> {
        self.of.get(region as usize).copied().flatten().map(|(i, _)| i)
    }

    /// How two regions relate semantically, in `-1.0..=1.0`.
    ///
    /// `+c` when both sit in the same instance with confidence `c`, `-c` when
    /// they sit in different ones or one is inside an object and the other is
    /// background, and `0.0` when neither belongs to anything — two patches of
    /// hillside get no opinion from a model that has no word for hillside, and
    /// fall back to arm A untouched.
    pub fn affinity(&self, a: u32, b: u32) -> f32 {
        let (a, b) = (
            self.of.get(a as usize).copied().flatten(),
            self.of.get(b as usize).copied().flatten(),
        );
        match (a, b) {
            (Some((ia, ca)), Some((ib, cb))) if ia == ib => ca.min(cb),
            (Some((_, ca)), Some((_, cb))) => -ca.min(cb),
            (Some((_, c)), None) | (None, Some((_, c))) => -c,
            (None, None) => 0.0,
        }
    }
}

#[derive(Debug, thiserror::Error, PartialEq)]
pub enum MembershipError {
    #[error("instance {instance} mask is {got} pixels, expected {expected}")]
    MaskSize {
        instance: usize,
        expected: usize,
        got: usize,
    },
}

/// Re-weight a region field's boundaries by semantic agreement.
///
/// Returns a field whose labels are untouched and whose adjacency saddles have
/// been scaled — so [`crate::MergeTree::build`] over the result yields a
/// hierarchy that climbs toward objects instead of toward whatever happened to
/// be smooth.
///
/// The scaling is `saddle * (1 - strength * affinity)`, which has the three
/// properties that matter: agreement shrinks a saddle toward zero without ever
/// reaching it, disagreement grows one without bound, and an affinity of zero
/// is exactly arm A. A pair the model has no opinion about is left alone
/// rather than nudged.
pub fn apply_semantic_prior(
    field: &RegionField,
    membership: &Membership,
    options: &PriorOptions,
) -> RegionField {
    let strength = options.strength.clamp(0.0, 0.99);

    let mut adjacency: Vec<Edge> = field
        .adjacency
        .iter()
        .map(|e| Edge {
            a: e.a,
            b: e.b,
            saddle: e.saddle * (1.0 - strength * membership.affinity(e.a, e.b)),
        })
        .collect();

    // Re-sorted because `MergeTree::build` consumes this in order and trusts
    // it to be sorted; the same total order as `RegionField::from_roots` uses,
    // for the same determinism reason.
    adjacency.sort_by(|x, y| {
        x.saddle
            .total_cmp(&y.saddle)
            .then(x.a.cmp(&y.a))
            .then(x.b.cmp(&y.b))
    });

    RegionField {
        width: field.width,
        height: field.height,
        labels: field.labels.clone(),
        region_count: field.region_count,
        adjacency,
    }
}

/// The regions making up one instance — click-to-select, snapped to edges.
///
/// This is the interaction the whole spike exists to enable, and the reason it
/// returns *region ids* rather than a raster: a mask that is a set of integers
/// is diffable, mergeable at node level under FR-NC-9, and cheap in a sidecar
/// (docs/segmentation.md §1). A raster is none of those.
///
/// The returned ids are sorted, so the same click always produces the same
/// mask — which is what lets it be a cache key.
pub fn regions_for_instance(
    field: &RegionField,
    mask: &[f32],
    options: &PriorOptions,
) -> Vec<u32> {
    let mut coverage = vec![0.0f32; field.region_count];
    let mut sizes = vec![0u32; field.region_count];

    for (p, &label) in field.labels.iter().enumerate() {
        coverage[label as usize] += mask.get(p).copied().unwrap_or(0.0);
        sizes[label as usize] += 1;
    }

    (0..field.region_count as u32)
        .filter(|&r| coverage[r as usize] / sizes[r as usize].max(1) as f32 >= options.membership)
        .collect()
}

/// Every pixel covered by a set of region ids, as a binary mask.
///
/// The other direction: region ids are what gets *stored*, and a rasteriser
/// needs pixels. On the shipping path this happens in a shader (ARCH §5.4);
/// this exists for export, for tests, and for the example.
pub fn rasterise(field: &RegionField, regions: &[u32]) -> Vec<f32> {
    let selected: HashMap<u32, ()> = regions.iter().map(|&r| (r, ())).collect();
    field
        .labels
        .iter()
        .map(|l| if selected.contains_key(l) { 1.0 } else { 0.0 })
        .collect()
}

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

    /// A 4x2 field: regions 0 and 1 on the left, 2 and 3 on the right.
    fn field() -> RegionField {
        RegionField {
            width: 4,
            height: 2,
            labels: vec![0, 0, 2, 2, 1, 1, 3, 3],
            region_count: 4,
            adjacency: vec![
                Edge { a: 0, b: 1, saddle: 1.0 },
                Edge { a: 0, b: 2, saddle: 1.0 },
                Edge { a: 1, b: 3, saddle: 1.0 },
                Edge { a: 2, b: 3, saddle: 1.0 },
            ],
        }
    }

    /// An instance covering the left half — regions 0 and 1.
    fn left_half() -> Vec<f32> {
        vec![1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0]
    }

    #[test]
    fn membership_finds_the_covered_regions() {
        let f = field();
        let m = Membership::compute(&f, &[&left_half()], &PriorOptions::default()).unwrap();

        assert_eq!(m.instance_of(0), Some(0));
        assert_eq!(m.instance_of(1), Some(0));
        assert_eq!(m.instance_of(2), None, "right half is outside the instance");
        assert_eq!(m.instance_of(3), None);
    }

    #[test]
    fn affinity_is_signed_by_agreement() {
        let f = field();
        let m = Membership::compute(&f, &[&left_half()], &PriorOptions::default()).unwrap();

        assert!(m.affinity(0, 1) > 0.0, "both inside the instance");
        assert!(m.affinity(0, 2) < 0.0, "across the instance boundary");
        assert_eq!(m.affinity(2, 3), 0.0, "model has no opinion on either");
    }

    /// The property arm C exists for: with equal image evidence everywhere,
    /// the semantic pair merges first.
    #[test]
    fn the_prior_reorders_the_merge_toward_the_object() {
        let f = field();
        let m = Membership::compute(&f, &[&left_half()], &PriorOptions::default()).unwrap();

        // Arm A alone: every saddle is 1.0, so the merge order is arbitrary
        // and 0-1 has no reason to come first.
        let plain = MergeTree::build(&f);
        assert_eq!(plain.merges.len(), 3);

        let biased = MergeTree::build(&apply_semantic_prior(&f, &m, &PriorOptions::default()));
        let first = biased.merges[0];
        assert_eq!(
            (first.a, first.b),
            (0, 1),
            "the two regions inside the instance should merge first"
        );
        assert!(
            first.saddle < 1.0,
            "agreement should lower the saddle, got {}",
            first.saddle
        );

        // And the boundary the model believes in should now be the last to go.
        let last = biased.merges.last().unwrap();
        assert!(
            last.saddle > 1.0,
            "a semantic boundary should outlast the others, got {}",
            last.saddle
        );
    }

    #[test]
    fn a_cut_at_two_groups_splits_along_the_instance() {
        let f = field();
        let m = Membership::compute(&f, &[&left_half()], &PriorOptions::default()).unwrap();
        let tree = MergeTree::build(&apply_semantic_prior(&f, &m, &PriorOptions::default()));

        let grouping = tree.cut_to(2);
        assert_eq!(grouping[0], grouping[1], "instance regions share a group");
        assert_eq!(grouping[2], grouping[3], "background regions share a group");
        assert_ne!(grouping[0], grouping[2], "and the two groups differ");
    }

    #[test]
    fn zero_strength_is_arm_a_exactly() {
        let f = field();
        let opts = PriorOptions { strength: 0.0, ..PriorOptions::default() };
        let m = Membership::compute(&f, &[&left_half()], &opts).unwrap();
        assert_eq!(apply_semantic_prior(&f, &m, &opts).adjacency, f.adjacency);
    }

    #[test]
    fn selection_snaps_to_whole_regions() {
        let f = field();

        // A mask that is ragged at the pixel level — as a quarter-resolution
        // sigmoid would be — still selects clean whole regions.
        let ragged = vec![1.0, 0.9, 0.1, 0.0, 0.8, 1.0, 0.0, 0.2];
        let regions = regions_for_instance(&f, &ragged, &PriorOptions::default());
        assert_eq!(regions, vec![0, 1]);

        let pixels = rasterise(&f, &regions);
        assert_eq!(pixels, vec![1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0]);
    }

    #[test]
    fn regions_are_sorted_so_a_click_is_a_cache_key() {
        let f = field();
        let all = vec![1.0; 8];
        let regions = regions_for_instance(&f, &all, &PriorOptions::default());
        assert!(regions.windows(2).all(|w| w[0] < w[1]));
    }

    #[test]
    fn a_wrong_sized_mask_is_an_error_not_a_panic() {
        let f = field();
        let err = Membership::compute(&f, &[&vec![0.0; 3]], &PriorOptions::default()).unwrap_err();
        assert_eq!(
            err,
            MembershipError::MaskSize { instance: 0, expected: 8, got: 3 }
        );
    }

    /// Two instances, so the "different objects repel" branch is covered.
    #[test]
    fn different_instances_repel() {
        let f = field();
        let right_half = vec![0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 1.0];
        let m = Membership::compute(
            &f,
            &[&left_half(), &right_half],
            &PriorOptions::default(),
        )
        .unwrap();

        assert_eq!(m.instance_of(0), Some(0));
        assert_eq!(m.instance_of(2), Some(1));
        assert!(m.affinity(0, 2) < 0.0, "two different objects should repel");
        assert!(m.affinity(2, 3) > 0.0, "same object should attract");
    }
}