feat: added some post processing for retinaface

bounding-box/src/draw.rs

@@ -0,0 +1,39 @@
+use crate::*;
+pub use color::Rgba8;
+
+use ndarray::{Array3, ArrayViewMut3};
+
+pub trait Draw<T> {
+    fn draw(&mut self, item: T, color: color::Rgba8, thickness: usize);
+}
+
+// impl<T: Drawable<Self>> Draw<T> for Array3<u8> {
+//     fn draw(&self, item: T, color: color::Rgba8, thickness: usize) {
+//         item.draw(&self, color, thickness);
+//     }
+// }
+
+pub trait Drawable<Canvas, T> {
+    fn draw(&self, canvas: &mut Canvas, color: color::Rgba8, thickness: T);
+}
+
+/// Implementing Drawable for Aabb2 with Array3<u8> as the canvas type
+/// Assuming Array3<u8> is a 3D array representing an image with RGB/RGBA channels
+impl<T> Drawable<ArrayViewMut3<'_, u8>, T> for Aabb2<T>
+where
+    T: Num + core::ops::SubAssign + core::ops::AddAssign + core::ops::DivAssign,
+    T: PartialOrd,
+{
+    fn draw(&self, canvas: &mut ArrayViewMut3<u8>, color: color::Rgba8, thickness: T) {
+        use itertools::Itertools;
+        let (height, width, channels) = canvas.dim();
+        self.corners()
+            .iter()
+            .zip(self.padding(thickness).corners())
+            .tuple_windows()
+            .for_each(|((a, b), (c, d))| {
+                let bbox = Aabb2::from_vertices([*a, b, *c, d]);
+                todo!();
+            });
+    }
+}

bounding-box/src/lib.rs

@@ -1,26 +1,53 @@
-use nalgebra::{Point, Point2, SVector};
+pub mod draw;
+pub mod nms;
+
+use nalgebra::{Point, Point2, Point3, SVector};
 pub trait Num: num::Num + Copy + core::fmt::Debug + 'static {}
 impl<T: num::Num + Copy + core::fmt::Debug + 'static> Num for T {}
 
+/// An axis aligned bounding box in `D` dimensions, defined by the minimum vertex and a size vector.
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
-pub struct BoundingBox<T: Num, const D: usize> {
+pub struct AxisAlignedBoundingBox<T: Num, const D: usize> {
     /// The point of the bounding box closest to the origin
     point: Point<T, D>,
     /// The size of the bounding box in each dimension
     size: SVector<T, D>,
 }
 
-impl<T: Num, const D: usize> BoundingBox<T, D> {
+pub type Aabb<T, const D: usize> = AxisAlignedBoundingBox<T, D>;
+pub type Aabb2<T> = AxisAlignedBoundingBox<T, 2>;
+pub type Aabb3<T> = AxisAlignedBoundingBox<T, 3>;
+
+impl<T: Num, const D: usize> AxisAlignedBoundingBox<T, D> {
     pub fn new(point: Point<T, D>, size: SVector<T, D>) -> Self {
         Self { point, size }
     }
 
-    pub fn point(&self) -> &Point<T, D> {
-        &self.point
+    pub fn from_min_max_vertices(point1: Point<T, D>, point2: Point<T, D>) -> Self
+    where
+        T: core::ops::SubAssign,
+    {
+        let size = point2 - point1;
+        Self::new(point1, SVector::from(size))
     }
 
-    pub fn size(&self) -> &SVector<T, D> {
-        &self.size
+    pub fn from_vertices(points: [Point<T, D>; 4]) -> Option<Self>
+    where
+        T: core::ops::SubAssign,
+        T: PartialOrd,
+    {
+        // find the closest and farthest points from the origin
+        let min = points
+            .iter()
+            .reduce(|acc, p| (acc > p).then_some(p).unwrap_or(acc))?;
+        let max = points
+            .iter()
+            .reduce(|acc, p| (acc < p).then_some(p).unwrap_or(acc))?;
+        Some(Self::from_min_max_vertices(*min, *max))
+    }
+
+    pub fn size(&self) -> SVector<T, D> {
+        self.size
     }
 
     pub fn center(&self) -> Point<T, D>
@@ -55,23 +82,109 @@ impl<T: Num, const D: usize> BoundingBox<T, D> {
         self.point += translation;
     }
 
-    pub fn contains(&self, point: &Point<T, D>) -> bool
+    pub fn min_vertex(&self) -> Point<T, D>
+    where
+        T: core::ops::SubAssign,
+    {
+        self.point
+    }
+
+    pub fn max_vertex(&self) -> Point<T, D>
+    where
+        T: core::ops::AddAssign,
+    {
+        self.point + self.size
+    }
+
+    pub fn contains_point(&self, point: &Point<T, D>) -> bool
     where
         T: core::ops::AddAssign,
         T: core::ops::SubAssign,
         T: PartialOrd,
     {
-        let p1 = self.point.coords;
-        let p2 = self.point.coords + self.size;
+        let min = self.min_vertex();
+        let max = self.max_vertex();
 
-        point.coords > p1 && point.coords < p2
+        *point > min && *point < max
+    }
+
+    pub fn scale(self, vector: SVector<T, D>) -> Self
+    where
+        T: core::ops::MulAssign,
+        T: core::ops::DivAssign,
+        T: core::ops::SubAssign,
+    {
+        let two = T::one() + T::one();
+        let new_size = self.size.component_mul(&vector);
+        let new_point = self.point.coords - new_size / two;
+        Self {
+            point: Point::from(new_point),
+            size: new_size,
+        }
+    }
+
+    pub fn contains_bbox(&self, other: &Self) -> bool
+    where
+        T: core::ops::AddAssign,
+        T: core::ops::SubAssign,
+        T: PartialOrd,
+    {
+        let self_min = self.min_vertex();
+        let self_max = self.max_vertex();
+        let other_min = other.min_vertex();
+        let other_max = other.max_vertex();
+
+        other_min >= self_min && other_max <= self_max
+    }
+
+    pub fn union(&self, other: &Self) -> Self
+    where
+        T: core::ops::AddAssign,
+        T: core::ops::SubAssign,
+        T: PartialOrd,
+    {
+        let self_min = self.min_vertex();
+        let self_max = self.max_vertex();
+        let other_min = other.min_vertex();
+        let other_max = other.max_vertex();
+        let max_of_min = (self_min.coords < other_min.coords)
+            .then_some(self_min.coords)
+            .unwrap_or(other_min.coords);
+        let min_of_max = (self_max.coords > other_max.coords)
+            .then_some(self_max.coords)
+            .unwrap_or(other_max.coords);
+        Self::from_min_max_vertices(Point::from(max_of_min), Point::from(min_of_max))
+    }
+
+    pub fn intersection(&self, other: &Self) -> Option<Self>
+    where
+        T: core::ops::AddAssign,
+        T: core::ops::SubAssign,
+        T: PartialOrd,
+    {
+        let self_min = self.min_vertex();
+        let self_max = self.max_vertex();
+        let other_min = other.min_vertex();
+        let other_max = other.max_vertex();
+
+        if self_max < other_min || other_max < self_min {
+            return None; // No intersection
+        }
+
+        let min = (self_min.coords > other_min.coords)
+            .then_some(self_min.coords)
+            .unwrap_or(other_min.coords);
+        let max = (self_max.coords < other_max.coords)
+            .then_some(self_max.coords)
+            .unwrap_or(other_max.coords);
+        Some(Self::from_min_max_vertices(
+            Point::from(min),
+            Point::from(max),
+        ))
     }
 }
 
-pub type BoundingBox2D<T> = BoundingBox<T, 2>;
-pub type BoundingBox3D<T> = BoundingBox<T, 3>;
-
-impl<T: Num> BoundingBox2D<T> {
+impl<T: Num> Aabb2<T> {
     pub fn new_2d(point1: Point2<T>, point2: Point2<T>) -> Self
     where
         T: core::ops::SubAssign,
@@ -110,6 +223,41 @@ impl<T: Num> BoundingBox2D<T> {
     {
         [self.x1y1(), self.x2y1(), self.x2y2(), self.x1y2()]
     }
+
+    pub fn area(&self) -> T
+    where
+        T: core::ops::Mul<Output = T>,
+    {
+        self.size.x * self.size.y
+    }
+
+    pub fn iou(&self, other: &Self) -> Option<T>
+    where
+        T: core::ops::AddAssign,
+        T: core::ops::SubAssign,
+        T: PartialOrd,
+    {
+        let intersection = self.intersection(other)?;
+        let union = self.union(other);
+        Some(intersection.area() / union.area())
+    }
+}
+
+impl<T: Num> Aabb3<T> {
+    pub fn new_3d(point1: Point3<T>, point2: Point3<T>) -> Self
+    where
+        T: core::ops::SubAssign,
+    {
+        let size = point2.coords - point1.coords;
+        Self::new(point1, SVector::from(size))
+    }
+
+    pub fn volume(&self) -> T
+    where
+        T: core::ops::Mul<Output = T>,
+    {
+        self.size.x * self.size.y * self.size.z
+    }
 }
 
 #[test]
@@ -118,10 +266,10 @@ fn test_bbox_new() {
 
     let point1 = Point2::new(1.0, 2.0);
     let point2 = Point2::new(4.0, 6.0);
-    let bbox = BoundingBox::new_2d(point1, point2);
+    let bbox = AxisAlignedBoundingBox::new_2d(point1, point2);
 
-    assert_eq!(bbox.point(), &point1);
-    assert_eq!(bbox.size(), &Vector2::new(3.0, 4.0));
+    assert_eq!(bbox.min_vertex(), point1);
+    assert_eq!(bbox.size(), Vector2::new(3.0, 4.0));
     assert_eq!(bbox.center(), Point2::new(2.5, 4.0));
 }
 
@@ -131,10 +279,10 @@ fn test_bounding_box_center_2d() {
 
     let point = Point2::new(1.0, 2.0);
     let size = Vector2::new(3.0, 4.0);
-    let bbox = BoundingBox::new(point, size);
+    let bbox = AxisAlignedBoundingBox::new(point, size);
 
-    assert_eq!(bbox.point(), &point);
-    assert_eq!(bbox.size(), &size);
+    assert_eq!(bbox.min_vertex(), point);
+    assert_eq!(bbox.size(), size);
     assert_eq!(bbox.center(), Point2::new(2.5, 4.0));
 }
 
@@ -144,10 +292,10 @@ fn test_bounding_box_center_3d() {
 
     let point = Point3::new(1.0, 2.0, 3.0);
     let size = Vector3::new(4.0, 5.0, 6.0);
-    let bbox = BoundingBox::new(point, size);
+    let bbox = AxisAlignedBoundingBox::new(point, size);
 
-    assert_eq!(bbox.point(), &point);
-    assert_eq!(bbox.size(), &size);
+    assert_eq!(bbox.min_vertex(), point);
+    assert_eq!(bbox.size(), size);
     assert_eq!(bbox.center(), Point3::new(3.0, 4.5, 6.0));
 }
 
@@ -157,9 +305,68 @@ fn test_bounding_box_padding_2d() {
 
     let point = Point2::new(1.0, 2.0);
     let size = Vector2::new(3.0, 4.0);
-    let bbox = BoundingBox::new(point, size);
+    let bbox = AxisAlignedBoundingBox::new(point, size);
 
     let padded_bbox = bbox.padding(1.0);
-    assert_eq!(padded_bbox.point(), &Point2::new(0.5, 1.5));
-    assert_eq!(padded_bbox.size(), &Vector2::new(4.0, 5.0));
+    assert_eq!(padded_bbox.min_vertex(), Point2::new(0.5, 1.5));
+    assert_eq!(padded_bbox.size(), Vector2::new(4.0, 5.0));
+}
+
+#[test]
+fn test_bounding_box_scaling_2d() {
+    use nalgebra::{Point2, Vector2};
+
+    let point = Point2::new(1.0, 1.0);
+    let size = Vector2::new(3.0, 4.0);
+    let bbox = AxisAlignedBoundingBox::new(point, size);
+
+    let padded_bbox = bbox.scale(Vector2::new(2.0, 2.0));
+    assert_eq!(padded_bbox.min_vertex(), Point2::new(-2.0, -3.0));
+    assert_eq!(padded_bbox.size(), Vector2::new(6.0, 8.0));
+}
+
+#[test]
+fn test_bounding_box_contains_2d() {
+    use nalgebra::Point2;
+
+    let point1 = Point2::new(1.0, 2.0);
+    let point2 = Point2::new(4.0, 6.0);
+    let bbox = AxisAlignedBoundingBox::new_2d(point1, point2);
+
+    assert!(bbox.contains_point(&Point2::new(2.0, 3.0)));
+    assert!(!bbox.contains_point(&Point2::new(5.0, 7.0)));
+}
+
+#[test]
+fn test_bounding_box_union_2d() {
+    use nalgebra::{Point2, Vector2};
+
+    let point1 = Point2::new(1.0, 2.0);
+    let point2 = Point2::new(4.0, 6.0);
+    let bbox1 = AxisAlignedBoundingBox::new_2d(point1, point2);
+
+    let point3 = Point2::new(3.0, 5.0);
+    let point4 = Point2::new(7.0, 8.0);
+    let bbox2 = AxisAlignedBoundingBox::new_2d(point3, point4);
+
+    let union_bbox = bbox1.union(&bbox2);
+    assert_eq!(union_bbox.min_vertex(), Point2::new(1.0, 2.0));
+    assert_eq!(union_bbox.size(), Vector2::new(6.0, 6.0));
+}
+
+#[test]
+fn test_bounding_box_intersection_2d() {
+    use nalgebra::{Point2, Vector2};
+
+    let point1 = Point2::new(1.0, 2.0);
+    let point2 = Point2::new(4.0, 6.0);
+    let bbox1 = AxisAlignedBoundingBox::new_2d(point1, point2);
+
+    let point3 = Point2::new(3.0, 5.0);
+    let point4 = Point2::new(5.0, 7.0);
+    let bbox2 = AxisAlignedBoundingBox::new_2d(point3, point4);
+
+    let intersection_bbox = bbox1.intersection(&bbox2).unwrap();
+    assert_eq!(intersection_bbox.min_vertex(), Point2::new(3.0, 5.0));
+    assert_eq!(intersection_bbox.size(), Vector2::new(1.0, 1.0));
 }

bounding-box/src/nms.rs

@@ -0,0 +1,68 @@
+use std::collections::HashSet;
+
+use crate::*;
+/// Apply Non-Maximum Suppression to a set of bounding boxes.
+///
+/// # Arguments
+///
+/// * `boxes` - A slice of bounding boxes to apply NMS on.
+/// * `threshold` - The IoU threshold for suppression.
+///
+/// # Returns
+///
+/// A vector of indices of the bounding boxes that are kept after applying NMS.
+pub fn nms<T>(
+    boxes: &[Aabb2<T>],
+    scores: &[T],
+    score_threshold: T,
+    nms_threshold: T,
+) -> Vec<Aabb2<T>>
+where
+    T: Num + num::Float + core::iter::Product<T> + core::ops::AddAssign + core::ops::SubAssign,
+{
+    use itertools::Itertools;
+    let bboxes: Vec<_> = boxes
+        .iter()
+        .zip(scores.iter())
+        .filter_map(|(bbox, score)| (score >= &score_threshold).then_some((bbox, score)))
+        .sorted_by(|(_, score_a), (_, score_b)| {
+            score_b
+                .partial_cmp(score_a)
+                .unwrap_or(std::cmp::Ordering::Equal)
+        })
+        .map(|(bbox, _)| bbox)
+        .collect();
+
+    let outputs = bboxes
+        .iter()
+        .enumerate()
+        .scan(
+            HashSet::with_capacity(bboxes.len()),
+            |state, (index, bbox)| {
+                if state.is_empty() {
+                    state.insert(index);
+                    return Some(Some(bbox));
+                } else {
+                    if state.contains(&index) {
+                        return Some(None);
+                    }
+                    let to_remove = bboxes
+                        .iter()
+                        .enumerate()
+                        .skip(index + 1)
+                        .filter_map(|(index, bbox_b)| {
+                            (!state.contains(&index)).then_some(index)?;
+                            let iou = bbox.iou(bbox_b)?;
+                            (iou >= nms_threshold).then_some(index)
+                        })
+                        .collect_vec();
+                    state.extend(to_remove);
+                    Some(Some(bbox))
+                }
+            },
+        )
+        .flatten()
+        .map(|bbox| **bbox)
+        .collect_vec();
+    outputs
+}