locus_core/

segmentation.rs

//! Connected components labeling (CCL) using Union-Find.
//!
//! This module implements the second stage of the pipeline, identifying and grouping
//! dark pixels that could potentially form fiducial markers.
//!
//! It provides:
//! - **Standard CCL**: Efficient binary component labeling.
//! - **Threshold-Model CCL**: Advanced connectivity based on local adaptive thresholds.

#![allow(unsafe_code)]

use bumpalo::Bump;
use bumpalo::collections::Vec as BumpVec;
use rayon::prelude::*;

/// A disjoint-set forest (Union-Find) with path compression and rank optimization.
pub struct UnionFind<'a> {
    parent: &'a mut [u32],
    rank: &'a mut [u8],
}

impl<'a> UnionFind<'a> {
    /// Create a new UnionFind structure backed by the provided arena.
    pub fn new_in(arena: &'a Bump, size: usize) -> Self {
        let parent = arena.alloc_slice_fill_with(size, |i| i as u32);
        let rank = arena.alloc_slice_fill_copy(size, 0u8);
        Self { parent, rank }
    }

    /// Find the representative (root) of the set containing `i`.
    #[inline]
    pub fn find(&mut self, i: u32) -> u32 {
        let mut root = i;
        while self.parent[root as usize] != root {
            self.parent[root as usize] = self.parent[self.parent[root as usize] as usize];
            root = self.parent[root as usize];
        }
        root
    }

    /// Unite the sets containing `i` and `j`.
    #[inline]
    pub fn union(&mut self, i: u32, j: u32) {
        let root_i = self.find(i);
        let root_j = self.find(j);
        if root_i != root_j {
            match self.rank[root_i as usize].cmp(&self.rank[root_j as usize]) {
                std::cmp::Ordering::Less => self.parent[root_i as usize] = root_j,
                std::cmp::Ordering::Greater => self.parent[root_j as usize] = root_i,
                std::cmp::Ordering::Equal => {
                    self.parent[root_i as usize] = root_j;
                    self.rank[root_j as usize] += 1;
                },
            }
        }
    }
}

/// Bounding box and statistics for a connected component.
#[derive(Clone, Copy, Debug)]
pub struct ComponentStats {
    /// Minimum x coordinate.
    pub min_x: u16,
    /// Maximum x coordinate.
    pub max_x: u16,
    /// Minimum y coordinate.
    pub min_y: u16,
    /// Maximum y coordinate.
    pub max_y: u16,
    /// Total number of pixels in the component.
    pub pixel_count: u32,
    /// First encountered pixel X (for boundary start).
    pub first_pixel_x: u16,
    /// First encountered pixel Y (for boundary start).
    pub first_pixel_y: u16,
    /// Spatial moment: Σ x (raw, integer).
    pub m10: u64,
    /// Spatial moment: Σ y (raw, integer).
    pub m01: u64,
    /// Spatial moment: Σ x² (raw, integer).
    pub m20: u64,
    /// Spatial moment: Σ y² (raw, integer).
    pub m02: u64,
    /// Spatial moment: Σ xy (raw, integer).
    pub m11: u64,
}

impl Default for ComponentStats {
    fn default() -> Self {
        Self {
            min_x: u16::MAX,
            max_x: 0,
            min_y: u16::MAX,
            max_y: 0,
            pixel_count: 0,
            first_pixel_x: 0,
            first_pixel_y: 0,
            m10: 0,
            m01: 0,
            m20: 0,
            m02: 0,
            m11: 0,
        }
    }
}

/// Compute shape descriptors from a component's spatial moments.
///
/// Returns `Some((elongation, density))` where:
/// - `elongation` ε = λ_max / λ_min (ratio of principal-axis variances; ≥ 1.0).
///   A perfect square has ε ≈ 1.0; a thin line or needle has ε >> 1.0.
/// - `density` D = pixel_count / bbox_area (fraction of bounding box filled).
///
/// Returns `None` for degenerate components (λ_min < 1e-9, e.g. single-pixel rows).
#[must_use]
pub fn compute_moment_shape(stats: &ComponentStats) -> Option<(f64, f64)> {
    let m00 = f64::from(stats.pixel_count);
    if m00 < 1.0 {
        return None;
    }

    let x_bar = stats.m10 as f64 / m00;
    let y_bar = stats.m01 as f64 / m00;

    // Central moments (unnormalized: sum of squared deviations, not divided by m00)
    let mu20 = stats.m20 as f64 - x_bar * stats.m10 as f64;
    let mu02 = stats.m02 as f64 - y_bar * stats.m01 as f64;
    let mu11 = stats.m11 as f64 - x_bar * stats.m01 as f64;

    // Eigenvalues of the normalized 2x2 covariance matrix.
    // λ = ((μ20 + μ02) ± sqrt((μ20 - μ02)² + 4·μ11²)) / (2·m00)
    let sum = mu20 + mu02;
    let diff = mu20 - mu02;
    let disc = (diff * diff + 4.0 * mu11 * mu11).sqrt();
    let lambda_max = (sum + disc) / (2.0 * m00);
    let lambda_min = (sum - disc) / (2.0 * m00);

    if lambda_min < 1e-9 {
        return None;
    }

    let elongation = lambda_max / lambda_min;

    let bbox_w = u32::from(stats.max_x - stats.min_x) + 1;
    let bbox_h = u32::from(stats.max_y - stats.min_y) + 1;
    let bbox_area = f64::from(bbox_w * bbox_h);
    let density = m00 / bbox_area;

    Some((elongation, density))
}

/// Result of connected component labeling.
pub struct LabelResult<'a> {
    /// Flat array of pixel labels (row-major).
    pub labels: &'a [u32],
    /// Statistics for each component (indexed by label - 1).
    pub component_stats: Vec<ComponentStats>,
}

/// A detected run of background pixels in a row.
#[derive(Clone, Copy, Debug)]
struct Run {
    y: u32,
    x_start: u32,
    x_end: u32,
    id: u32,
}

/// Label connected components in a binary image.
pub fn label_components<'a>(
    arena: &'a Bump,
    binary: &[u8],
    width: usize,
    height: usize,
    use_8_connectivity: bool,
) -> &'a [u32] {
    label_components_with_stats(arena, binary, width, height, use_8_connectivity).labels
}

/// Label components and compute bounding box stats for each.
#[allow(clippy::too_many_lines)]
pub fn label_components_with_stats<'a>(
    arena: &'a Bump,
    binary: &[u8],
    width: usize,
    height: usize,
    use_8_connectivity: bool,
) -> LabelResult<'a> {
    // Pass 1: Extract runs - Optimized with Rayon parallel processing
    let all_runs: Vec<Vec<Run>> = binary
        .par_chunks(width)
        .enumerate()
        .map(|(y, row)| {
            let mut row_runs = Vec::with_capacity(width / 4 + 1);
            let mut x = 0;
            while x < width {
                // Find start of background run (0)
                if let Some(pos) = row[x..].iter().position(|&p| p == 0) {
                    let start = x + pos;
                    // Find end of run
                    let len = row[start..].iter().take_while(|&&p| p == 0).count();
                    row_runs.push(Run {
                        y: y as u32,
                        x_start: start as u32,
                        x_end: (start + len - 1) as u32,
                        id: 0,
                    });
                    x = start + len;
                } else {
                    break;
                }
            }
            row_runs
        })
        .collect();

    let total_runs: usize = all_runs.iter().map(std::vec::Vec::len).sum();
    let mut runs: BumpVec<Run> = BumpVec::with_capacity_in(total_runs, arena);
    for (id, mut run) in all_runs.into_iter().flatten().enumerate() {
        run.id = id as u32;
        runs.push(run);
    }

    if runs.is_empty() {
        return LabelResult {
            labels: arena.alloc_slice_fill_copy(width * height, 0u32),
            component_stats: Vec::new(),
        };
    }

    let mut uf = UnionFind::new_in(arena, runs.len());
    let mut curr_row_range = 0..0; // Initialize curr_row_range
    let mut i = 0;

    // Pass 2: Link runs between adjacent rows using two-pointer linear scan
    while i < runs.len() {
        let y = runs[i].y;

        // Identify the range of runs in the current row
        let start = i;
        while i < runs.len() && runs[i].y == y {
            i += 1;
        }
        let prev_row_range = curr_row_range; // Now correctly uses the previously assigned curr_row_range
        curr_row_range = start..i;

        if y > 0 && !prev_row_range.is_empty() && runs[prev_row_range.start].y == y - 1 {
            let mut p_idx = prev_row_range.start;
            for c_idx in curr_row_range.clone() {
                let curr = &runs[c_idx];

                if use_8_connectivity {
                    // 8-connectivity check: overlap if [xs1, xe1] and [xs2-1, xe2+1] intersect
                    while p_idx < prev_row_range.end && runs[p_idx].x_end + 1 < curr.x_start {
                        p_idx += 1;
                    }
                    let mut temp_p = p_idx;
                    while temp_p < prev_row_range.end && runs[temp_p].x_start <= curr.x_end + 1 {
                        uf.union(curr.id, runs[temp_p].id);
                        temp_p += 1;
                    }
                } else {
                    // 4-connectivity check: overlap if [xs1, xe1] and [xs2, xe2] intersect
                    while p_idx < prev_row_range.end && runs[p_idx].x_end < curr.x_start {
                        p_idx += 1;
                    }
                    let mut temp_p = p_idx;
                    while temp_p < prev_row_range.end && runs[temp_p].x_start <= curr.x_end {
                        uf.union(curr.id, runs[temp_p].id);
                        temp_p += 1;
                    }
                }
            }
        }
    }

    // Pass 3: Collect stats per root and assign labels
    let mut root_to_label: Vec<u32> = vec![0; runs.len()];
    let mut component_stats: Vec<ComponentStats> = Vec::new();
    let mut next_label = 1u32;

    // Pre-resolve roots to avoid repeated find() in Pass 4
    let mut run_roots = Vec::with_capacity(runs.len());

    for run in &runs {
        let root = uf.find(run.id) as usize;
        run_roots.push(root);
        if root_to_label[root] == 0 {
            root_to_label[root] = next_label;
            next_label += 1;
            let new_stat = ComponentStats {
                first_pixel_x: run.x_start as u16,
                first_pixel_y: run.y as u16,
                ..Default::default()
            };
            component_stats.push(new_stat);
        }
        let label_idx = (root_to_label[root] - 1) as usize;
        let stats = &mut component_stats[label_idx];
        stats.min_x = stats.min_x.min(run.x_start as u16);
        stats.max_x = stats.max_x.max(run.x_end as u16);
        stats.min_y = stats.min_y.min(run.y as u16);
        stats.max_y = stats.max_y.max(run.y as u16);
        stats.pixel_count += run.x_end - run.x_start + 1;
        // Accumulate spatial moments using closed-form per-run sums.
        // Run covers x in [a, b) exclusive, i.e. x = a, a+1, ..., b-1.
        let a = u64::from(run.x_start);
        let b = u64::from(run.x_end) + 1; // convert inclusive end to exclusive
        let yu = u64::from(run.y);
        // SAFETY: `a - 1` and `2*a - 1` are always multiplied by `a`, so their
        // value is irrelevant when a = 0. saturating_sub avoids u64 underflow in
        // debug builds (release wraps, but the `* a` factor zeros the term anyway).
        stats.m10 += b * (b - 1) / 2 - a * a.saturating_sub(1) / 2;
        stats.m01 += yu * (b - a);
        stats.m20 +=
            (b - 1) * b * (2 * b - 1) / 6 - a.saturating_sub(1) * a * (2 * a).saturating_sub(1) / 6;
        stats.m02 += yu * yu * (b - a);
        stats.m11 += yu * (b * (b - 1) / 2 - a * a.saturating_sub(1) / 2);
    }

    // Pass 4: Assign labels to pixels - Optimized with slice fill
    let labels = arena.alloc_slice_fill_copy(width * height, 0u32);
    for (run, root) in runs.iter().zip(run_roots) {
        let label = root_to_label[root];
        let row_off = run.y as usize * width;
        labels[row_off + run.x_start as usize..=row_off + run.x_end as usize].fill(label);
    }

    LabelResult {
        labels,
        component_stats,
    }
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::cast_sign_loss)]
mod tests {
    use super::*;
    use bumpalo::Bump;
    use proptest::prelude::prop;
    use proptest::proptest;

    #[test]
    fn test_union_find() {
        let arena = Bump::new();
        let mut uf = UnionFind::new_in(&arena, 10);

        uf.union(1, 2);
        uf.union(2, 3);
        uf.union(5, 6);

        assert_eq!(uf.find(1), uf.find(3));
        assert_eq!(uf.find(1), uf.find(2));
        assert_ne!(uf.find(1), uf.find(5));

        uf.union(3, 5);
        assert_eq!(uf.find(1), uf.find(6));
    }

    #[test]
    fn test_label_components_simple() {
        let arena = Bump::new();
        // 6x6 image with two separate 2x2 squares that are NOT 8-connected.
        // 0 = background (black), 255 = foreground (white)
        // Tag detector looks for black components (0)
        let binary = [
            0, 0, 255, 255, 255, 255, 0, 0, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
            255, 255, 0, 0, 255, 255, 255, 255, 0, 0, 255, 255, 255, 255, 255, 255, 255,
        ];
        let width = 6;
        let height = 6;

        let result = label_components_with_stats(&arena, &binary, width, height, false);

        assert_eq!(result.component_stats.len(), 2);

        // Component 1 (top-left)
        let s1 = result.component_stats[0];
        assert_eq!(s1.pixel_count, 4);
        assert_eq!(s1.min_x, 0);
        assert_eq!(s1.max_x, 1);
        assert_eq!(s1.min_y, 0);
        assert_eq!(s1.max_y, 1);

        // Component 2 (middle-rightish)
        let s2 = result.component_stats[1];
        assert_eq!(s2.pixel_count, 4);
        assert_eq!(s2.min_x, 3);
        assert_eq!(s2.max_x, 4);
        assert_eq!(s2.min_y, 3);
        assert_eq!(s2.max_y, 4);
    }

    #[test]
    fn test_segmentation_with_decimation() {
        let arena = Bump::new();
        let width = 32;
        let height = 32;
        let mut binary = vec![255u8; width * height];
        // Draw a 10x10 black square at (10,10)
        for y in 10..20 {
            for x in 10..20 {
                binary[y * width + x] = 0;
            }
        }

        // use statement moved to top of module or block
        let img =
            crate::image::ImageView::new(&binary, width, height, width).expect("valid creation");

        // Decimate by 2 -> 16x16
        let mut decimated_data = vec![0u8; 16 * 16];
        let decimated_img = img
            .decimate_to(2, &mut decimated_data)
            .expect("decimation failed");
        // In decimated image, square should be roughly at (5,5) with size 5x5
        let result = label_components_with_stats(&arena, decimated_img.data, 16, 16, true);

        assert_eq!(result.component_stats.len(), 1);
        let s = result.component_stats[0];
        assert_eq!(s.pixel_count, 25);
        assert_eq!(s.min_x, 5);
        assert_eq!(s.max_x, 9);
        assert_eq!(s.min_y, 5);
        assert_eq!(s.max_y, 9);
    }

    proptest! {
        #[test]
        fn prop_union_find_reflexivity(size in 1..1000usize) {
            let arena = Bump::new();
            let mut uf = UnionFind::new_in(&arena, size);
            for i in 0..size as u32 {
                assert_eq!(uf.find(i), i);
            }
        }

        #[test]
        fn prop_union_find_transitivity(size in 1..1000usize, pairs in prop::collection::vec((0..1000u32, 0..1000u32), 0..100)) {
            let arena = Bump::new();
            let real_size = size.max(1001); // Ensure indices are in range
            let mut uf = UnionFind::new_in(&arena, real_size);

            for (a, b) in pairs {
                let a = a % real_size as u32;
                let b = b % real_size as u32;
                uf.union(a, b);
                assert_eq!(uf.find(a), uf.find(b));
            }
        }

        #[test]
        fn prop_label_components_no_panic(
            width in 1..64usize,
            height in 1..64usize,
            data in prop::collection::vec(0..=1u8, 64 * 64)
        ) {
            let arena = Bump::new();
            let binary: Vec<u8> = data.iter().map(|&b| if b == 0 { 0 } else { 255 }).collect();
            let real_width = width.min(64);
            let real_height = height.min(64);
            let slice = &binary[..real_width * real_height];

            let result = label_components_with_stats(&arena, slice, real_width, real_height, true);

            for stat in result.component_stats {
                assert!(stat.pixel_count > 0);
                assert!(stat.max_x < real_width as u16);
                assert!(stat.max_y < real_height as u16);
                assert!(stat.min_x <= stat.max_x);
                assert!(stat.min_y <= stat.max_y);
            }
        }
    }

    // ========================================================================
    // SEGMENTATION ROBUSTNESS TESTS
    // ========================================================================

    use crate::config::TagFamily;
    use crate::image::ImageView;
    use crate::test_utils::{TestImageParams, generate_test_image_with_params};
    use crate::threshold::ThresholdEngine;

    /// Helper: Generate a binarized tag image at the given size.
    fn generate_binarized_tag(tag_size: usize, canvas_size: usize) -> (Vec<u8>, [[f64; 2]; 4]) {
        let params = TestImageParams {
            family: TagFamily::AprilTag36h11,
            id: 0,
            tag_size,
            canvas_size,
            ..Default::default()
        };

        let (data, corners) = generate_test_image_with_params(&params);
        let img = ImageView::new(&data, canvas_size, canvas_size, canvas_size).unwrap();

        let arena = Bump::new();
        let engine = ThresholdEngine::new();
        let stats = engine.compute_tile_stats(&arena, &img);
        let mut binary = vec![0u8; canvas_size * canvas_size];
        engine.apply_threshold(&arena, &img, &stats, &mut binary);

        (binary, corners)
    }

    /// Test segmentation at varying tag sizes.
    #[test]
    fn test_segmentation_at_varying_tag_sizes() {
        let canvas_size = 640;
        let tag_sizes = [32, 64, 100, 200, 300];

        for tag_size in tag_sizes {
            let arena = Bump::new();
            let (binary, corners) = generate_binarized_tag(tag_size, canvas_size);

            let result =
                label_components_with_stats(&arena, &binary, canvas_size, canvas_size, true);

            assert!(
                !result.component_stats.is_empty(),
                "Tag size {tag_size}: No components found"
            );

            let largest = result
                .component_stats
                .iter()
                .max_by_key(|s| s.pixel_count)
                .unwrap();

            let expected_min_x = corners[0][0] as u16;
            let expected_max_x = corners[1][0] as u16;
            let tolerance = 5;

            assert!(
                (i32::from(largest.min_x) - i32::from(expected_min_x)).abs() <= tolerance,
                "Tag size {tag_size}: min_x mismatch"
            );
            assert!(
                (i32::from(largest.max_x) - i32::from(expected_max_x)).abs() <= tolerance,
                "Tag size {tag_size}: max_x mismatch"
            );

            println!(
                "Tag size {:>3}px: {} components, largest has {} px",
                tag_size,
                result.component_stats.len(),
                largest.pixel_count
            );
        }
    }

    /// Test component pixel counts are reasonable for clean binarization.
    #[test]
    fn test_segmentation_component_accuracy() {
        let canvas_size = 320;
        let tag_size = 120;

        let arena = Bump::new();
        let (binary, corners) = generate_binarized_tag(tag_size, canvas_size);

        let result = label_components_with_stats(&arena, &binary, canvas_size, canvas_size, true);

        let largest = result
            .component_stats
            .iter()
            .max_by_key(|s| s.pixel_count)
            .unwrap();

        let expected_min = (tag_size * tag_size / 3) as u32;
        let expected_max = (tag_size * tag_size) as u32;

        assert!(largest.pixel_count >= expected_min);
        assert!(largest.pixel_count <= expected_max);

        let gt_width = (corners[1][0] - corners[0][0]).abs() as i32;
        let gt_height = (corners[2][1] - corners[0][1]).abs() as i32;
        let bbox_width = i32::from(largest.max_x - largest.min_x);
        let bbox_height = i32::from(largest.max_y - largest.min_y);

        assert!((bbox_width - gt_width).abs() <= 2);
        assert!((bbox_height - gt_height).abs() <= 2);

        println!(
            "Component accuracy: {} pixels, bbox={}x{} (GT: {}x{})",
            largest.pixel_count, bbox_width, bbox_height, gt_width, gt_height
        );
    }

    /// Test segmentation with noisy binary boundaries.
    #[test]
    fn test_segmentation_noisy_boundaries() {
        let canvas_size = 320;
        let tag_size = 120;

        let arena = Bump::new();
        let (mut binary, _corners) = generate_binarized_tag(tag_size, canvas_size);

        let noise_rate = 0.05;

        for y in 1..(canvas_size - 1) {
            for x in 1..(canvas_size - 1) {
                let idx = y * canvas_size + x;
                let current = binary[idx];
                let left = binary[idx - 1];
                let right = binary[idx + 1];
                let up = binary[idx - canvas_size];
                let down = binary[idx + canvas_size];

                let is_edge =
                    current != left || current != right || current != up || current != down;
                if is_edge && rand::random_range(0.0..1.0_f32) < noise_rate {
                    binary[idx] = if current == 0 { 255 } else { 0 };
                }
            }
        }

        let result = label_components_with_stats(&arena, &binary, canvas_size, canvas_size, true);

        assert!(!result.component_stats.is_empty());

        let largest = result
            .component_stats
            .iter()
            .max_by_key(|s| s.pixel_count)
            .unwrap();

        let min_expected = (tag_size * tag_size / 4) as u32;
        assert!(largest.pixel_count >= min_expected);

        println!(
            "Noisy segmentation: {} components, largest has {} px",
            result.component_stats.len(),
            largest.pixel_count
        );
    }

    #[test]
    fn test_segmentation_correctness_large_image() {
        let arena = Bump::new();
        let width = 3840; // 4K width
        let height = 2160; // 4K height
        let mut binary = vec![255u8; width * height];

        // Create multiple separate components
        // 1. A square at the top left
        for y in 100..200 {
            for x in 100..200 {
                binary[y * width + x] = 0;
            }
        }

        // 2. A long horizontal strip in the middle
        for x in 500..3000 {
            binary[1000 * width + x] = 0;
            binary[1001 * width + x] = 0;
        }

        // 3. Horizontal stripes at the bottom
        for y in 1800..2000 {
            if y % 4 == 0 {
                for x in 1800..2000 {
                    binary[y * width + x] = 0;
                }
            }
        }

        // 4. Noise (avoiding the square area)
        for y in 0..height {
            if y % 10 == 0 {
                for x in 0..width {
                    if (!(100..200).contains(&x) || !(100..200).contains(&y)) && (x + y) % 31 == 0 {
                        binary[y * width + x] = 0;
                    }
                }
            }
        }

        // Run segmentation
        let start = std::time::Instant::now();
        let result = label_components_with_stats(&arena, &binary, width, height, true);
        let duration = start.elapsed();

        // Basic verification of component counts
        assert!(result.component_stats.len() > 1000);

        println!(
            "Found {} components on 4K image in {:?}",
            result.component_stats.len(),
            duration
        );
    }
}