mecha10-core 0.6.2

//! Shared 2D axis-aligned bounding-box geometry and greedy Non-Maximum
//! Suppression (NMS).
//!
//! `mecha10-core`'s own [`crate::messages::perception::BoundingBox`] and
//! `mecha10-vision`'s `BoundingBox` are separate wire/domain types (they
//! carry different derives and doc semantics), but both need the same
//! axis-aligned-rectangle math — area, center, and Intersection-over-Union.
//! The [`Rect`] trait factors that math out once so both types (and any
//! future box type) implement it via a handful of field accessors instead of
//! carrying their own copy of the IoU formula.
//!
//! [`non_max_suppression`] factors out the matching greedy-NMS loop that
//! `mecha10-nodes-object-detector` and `mecha10-vision`'s YOLO detector both
//! run over their own `Detection` types after scoring boxes with [`Rect`].

/// Axis-aligned rectangle accessors, with IoU/area/center derived once from
/// them so callers don't hand-roll the same formula per box type.
pub trait Rect {
    fn x(&self) -> f32;
    fn y(&self) -> f32;
    fn width(&self) -> f32;
    fn height(&self) -> f32;

    /// Area of the rectangle.
    fn area(&self) -> f32 {
        self.width() * self.height()
    }

    /// Center point of the rectangle.
    fn center(&self) -> (f32, f32) {
        (self.x() + self.width() / 2.0, self.y() + self.height() / 2.0)
    }

    /// Intersection-over-Union with another rectangle (0.0 if they don't overlap).
    fn iou<R: Rect>(&self, other: &R) -> f32 {
        let x1 = self.x().max(other.x());
        let y1 = self.y().max(other.y());
        let x2 = (self.x() + self.width()).min(other.x() + other.width());
        let y2 = (self.y() + self.height()).min(other.y() + other.height());

        if x2 < x1 || y2 < y1 {
            return 0.0;
        }

        let intersection = (x2 - x1) * (y2 - y1);
        let union = self.area() + other.area() - intersection;

        intersection / union
    }
}

/// Greedy Non-Maximum Suppression: sorts `items` by descending score, then
/// repeatedly keeps the highest-scoring remaining item and discards the rest
/// of the same class whose IoU with it is at or above `iou_threshold`.
///
/// Generic over the item type via the three accessor closures so callers
/// don't need a shared `Detection` type — only a shared score/class-key/IoU
/// shape, which is all the duplicated logic actually needed.
pub fn non_max_suppression<T: Clone, K: PartialEq>(
    mut items: Vec<T>,
    iou_threshold: f32,
    mut score: impl FnMut(&T) -> f32,
    mut class_key: impl FnMut(&T) -> K,
    mut iou: impl FnMut(&T, &T) -> f32,
) -> Vec<T> {
    items.sort_by(|a, b| score(b).partial_cmp(&score(a)).unwrap());

    let mut keep = Vec::new();
    while !items.is_empty() {
        let current = items.remove(0);
        keep.push(current.clone());

        items.retain(|item| class_key(item) != class_key(&current) || iou(item, &current) < iou_threshold);
    }

    keep
}

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

    #[derive(Clone, Copy, Debug, PartialEq)]
    struct TestRect {
        x: f32,
        y: f32,
        width: f32,
        height: f32,
    }

    impl Rect for TestRect {
        fn x(&self) -> f32 {
            self.x
        }
        fn y(&self) -> f32 {
            self.y
        }
        fn width(&self) -> f32 {
            self.width
        }
        fn height(&self) -> f32 {
            self.height
        }
    }

    fn rect(x: f32, y: f32, width: f32, height: f32) -> TestRect {
        TestRect { x, y, width, height }
    }

    #[test]
    fn area_and_center() {
        let r = rect(0.0, 0.0, 4.0, 2.0);
        assert_eq!(r.area(), 8.0);
        assert_eq!(r.center(), (2.0, 1.0));
    }

    #[test]
    fn iou_of_identical_boxes_is_one() {
        let a = rect(0.0, 0.0, 2.0, 2.0);
        let b = rect(0.0, 0.0, 2.0, 2.0);
        assert!((a.iou(&b) - 1.0).abs() < 1e-6);
    }

    #[test]
    fn iou_of_disjoint_boxes_is_zero() {
        let a = rect(0.0, 0.0, 1.0, 1.0);
        let b = rect(5.0, 5.0, 1.0, 1.0);
        assert_eq!(a.iou(&b), 0.0);
    }

    #[test]
    fn iou_of_half_overlapping_boxes() {
        let a = rect(0.0, 0.0, 2.0, 2.0);
        let b = rect(1.0, 0.0, 2.0, 2.0);
        // intersection = 1x2 = 2, union = 4 + 4 - 2 = 6
        assert!((a.iou(&b) - (2.0 / 6.0)).abs() < 1e-6);
    }

    #[derive(Clone, Debug, PartialEq)]
    struct Scored {
        class: u32,
        score: f32,
        rect: TestRect,
    }

    #[test]
    fn nms_keeps_highest_scoring_of_overlapping_same_class_items() {
        let items = vec![
            Scored {
                class: 0,
                score: 0.9,
                rect: rect(0.0, 0.0, 2.0, 2.0),
            },
            Scored {
                class: 0,
                score: 0.5,
                rect: rect(0.2, 0.0, 2.0, 2.0), // heavily overlaps the 0.9 box
            },
            Scored {
                class: 1,
                score: 0.4,
                rect: rect(0.2, 0.0, 2.0, 2.0), // same location, different class: kept
            },
        ];

        let kept = non_max_suppression(items, 0.5, |i| i.score, |i| i.class, |a, b| a.rect.iou(&b.rect));

        assert_eq!(kept.len(), 2);
        assert!(kept.iter().any(|i| i.class == 0 && i.score == 0.9));
        assert!(kept.iter().any(|i| i.class == 1));
    }

    #[test]
    fn nms_keeps_non_overlapping_items_of_same_class() {
        let items = vec![
            Scored {
                class: 0,
                score: 0.9,
                rect: rect(0.0, 0.0, 1.0, 1.0),
            },
            Scored {
                class: 0,
                score: 0.8,
                rect: rect(10.0, 10.0, 1.0, 1.0),
            },
        ];

        let kept = non_max_suppression(items, 0.5, |i| i.score, |i| i.class, |a, b| a.rect.iou(&b.rect));

        assert_eq!(kept.len(), 2);
    }
}