feat: added some post processing for retinaface

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));
 }