feat: Added postprocessing for retinaface

bounding-box/src/draw.rs

@@ -1,39 +1,53 @@
 use crate::*;
 pub use color::Rgba8;
 
-use ndarray::{Array3, ArrayViewMut3};
+use ndarray::{Array1, 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);
-//     }
-// }
+impl Draw<Aabb2<usize>> for Array3<u8> {
+    fn draw(&mut self, item: Aabb2<usize>, 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);
+pub trait Drawable<Canvas> {
+    fn draw(&self, canvas: &mut Canvas, color: color::Rgba8, thickness: usize);
 }
 
 /// 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) {
+impl Drawable<ArrayViewMut3<'_, u8>> for Aabb2<usize> {
+    fn draw(&self, canvas: &mut ArrayViewMut3<u8>, color: color::Rgba8, thickness: usize) {
         use itertools::Itertools;
-        let (height, width, channels) = canvas.dim();
+        // let (height, width, channels) = canvas.dim();
+        let color = Array1::from_vec(vec![color.r, color.g, color.b, color.a]);
         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!();
+                let bbox = Aabb2::from_vertices([*a, b, *c, d]).expect("Invalid bounding box");
+                use crate::roi::RoiMut;
+                let mut out = canvas.roi_mut(bbox).expect("Failed to get ROI");
+                out.lanes_mut(ndarray::Axis(2))
+                    .into_iter()
+                    .for_each(|mut pixel| {
+                        pixel.assign(&color);
+                    });
             });
     }
 }
+
+impl Drawable<Array3<u8>> for Aabb2<usize> {
+    fn draw(&self, canvas: &mut Array3<u8>, color: color::Rgba8, thickness: usize) {
+        use itertools::Itertools;
+        // let (height, width, channels) = canvas.dim();
+        let color = Array1::from_vec(vec![color.r, color.g, color.b, color.a]);
+        let pixel_size = canvas.dim().2;
+        let color = color.slice(ndarray::s![..pixel_size]);
+        let [x1y1, x2y1, x1y2, x2y2] = self.corners();
+    }
+}

bounding-box/src/lib.rs

@@ -1,5 +1,6 @@
 pub mod draw;
 pub mod nms;
+pub mod roi;
 
 use nalgebra::{Point, Point2, Point3, SVector};
 pub trait Num: num::Num + Copy + core::fmt::Debug + 'static {}
@@ -31,6 +32,8 @@ impl<T: Num, const D: usize> AxisAlignedBoundingBox<T, D> {
         Self::new(point1, SVector::from(size))
     }
 
+    /// Only considers the points closest and furthest from origin
+    /// Points which are rotated along in the z axis (in 2d) are not considered
     pub fn from_vertices(points: [Point<T, D>; 4]) -> Option<Self>
     where
         T: core::ops::SubAssign,
@@ -182,9 +185,51 @@ impl<T: Num, const D: usize> AxisAlignedBoundingBox<T, D> {
             Point::from(max),
         ))
     }
+
+    pub fn denormalize(&self, factor: nalgebra::SVector<T, D>) -> Self
+    where
+        T: core::ops::MulAssign,
+        T: core::ops::AddAssign,
+        // nalgebra::constraint::ShapeConstraint:
+        //     nalgebra::constraint::DimEq<nalgebra::Const<D>, nalgebra::Const<D>>,
+    {
+        Self {
+            point: (self.point.coords.component_mul(&factor)).into(),
+            size: self.size.component_mul(&factor),
+        }
+    }
+
+    pub fn cast<T2>(&self) -> Option<Aabb<T2, D>>
+    where
+        // T: num::NumCast,
+        T2: Num + simba::scalar::SubsetOf<T>,
+    {
+        Some(Aabb {
+            point: Point::from(self.point.coords.try_cast::<T2>()?),
+            size: self.size.try_cast::<T2>()?,
+        })
+    }
+
+    // pub fn as_<T2>(&self) -> Option<Aabb<T2, D>>
+    // where
+    //     T2: Num + simba::scalar::SubsetOf<T>,
+    // {
+    //     Some(Aabb {
+    //         point: Point::from(self.point.coords.as_()),
+    //         size: self.size.as_(),
+    //     })
+    // }
 }
 
 impl<T: Num> Aabb2<T> {
+    pub fn from_x1y1x2y2(x1: T, x2: T, y1: T, y2: T) -> Self
+    where
+        T: core::ops::SubAssign,
+    {
+        let point1 = Point2::new(x1, y1);
+        let point2 = Point2::new(x2, y2);
+        Self::from_min_max_vertices(point1, point2)
+    }
     pub fn new_2d(point1: Point2<T>, point2: Point2<T>) -> Self
     where
         T: core::ops::SubAssign,
@@ -217,6 +262,28 @@ impl<T: Num> Aabb2<T> {
         Point2::new(self.point.x, self.point.y + self.size.y)
     }
 
+    pub fn x1(&self) -> T {
+        self.point.x
+    }
+
+    pub fn y1(&self) -> T {
+        self.point.y
+    }
+
+    pub fn x2(&self) -> T
+    where
+        T: core::ops::AddAssign,
+    {
+        self.point.x + self.size.x
+    }
+
+    pub fn y2(&self) -> T
+    where
+        T: core::ops::AddAssign,
+    {
+        self.point.y + self.size.y
+    }
+
     pub fn corners(&self) -> [Point2<T>; 4]
     where
         T: core::ops::AddAssign,

bounding-box/src/roi.rs

@@ -0,0 +1,97 @@
+use crate::*;
+use ndarray::{Array3, ArrayView3, ArrayViewMut3};
+/// A trait that extracts a region of interest from an image
+pub trait Roi<'a, Output> {
+    type Error;
+    fn roi(&'a self, aabb: Aabb2<usize>) -> Result<Output, Self::Error>;
+}
+pub trait RoiMut<'a, Output> {
+    type Error;
+    fn roi_mut(&'a mut self, aabb: Aabb2<usize>) -> Result<Output, Self::Error>;
+}
+#[derive(thiserror::Error, Debug, Copy, Clone)]
+pub enum RoiError {
+    #[error("Region of intereset is out of bounds")]
+    RoiOutOfBounds,
+}
+
+impl<'a, T: Num> Roi<'a, ArrayView3<'a, T>> for Array3<T> {
+    type Error = RoiError;
+    fn roi(&'a self, aabb: Aabb2<usize>) -> Result<ArrayView3<'a, T>, Self::Error> {
+        let x1 = aabb.x1();
+        let x2 = aabb.x2();
+        let y1 = aabb.y1();
+        let y2 = aabb.y2();
+        if x1 >= x2 || y1 >= y2 || x2 > self.shape()[1] || y2 > self.shape()[0] {
+            return Err(RoiError::RoiOutOfBounds);
+        }
+        Ok(self.slice(ndarray::s![y1..y2, x1..x2, ..]))
+    }
+}
+
+impl<'a, T: Num> RoiMut<'a, ArrayViewMut3<'a, T>> for Array3<T> {
+    type Error = RoiError;
+    fn roi_mut(&'a mut self, aabb: Aabb2<usize>) -> Result<ArrayViewMut3<'a, T>, Self::Error> {
+        let x1 = aabb.x1();
+        let x2 = aabb.x2();
+        let y1 = aabb.y1();
+        let y2 = aabb.y2();
+        if x1 >= x2 || y1 >= y2 || x2 > self.shape()[1] || y2 > self.shape()[0] {
+            return Err(RoiError::RoiOutOfBounds);
+        }
+        Ok(self.slice_mut(ndarray::s![y1..y2, x1..x2, ..]))
+    }
+}
+
+impl<'a, 'b, T: Num> Roi<'a, ArrayView3<'b, T>> for ArrayView3<'b, T> {
+    type Error = RoiError;
+    fn roi(&'a self, aabb: Aabb2<usize>) -> Result<ArrayView3<'b, T>, Self::Error> {
+        let x1 = aabb.x1();
+        let x2 = aabb.x2();
+        let y1 = aabb.y1();
+        let y2 = aabb.y2();
+        if x1 >= x2 || y1 >= y2 || x2 > self.shape()[1] || y2 > self.shape()[0] {
+            return Err(RoiError::RoiOutOfBounds);
+        }
+        Ok(self.slice_move(ndarray::s![y1..y2, x1..x2, ..]))
+    }
+}
+// impl<'a, 'b, T: Num> Roi<'a, ArrayViewMut3<'b, T>> for ArrayViewMut3<'b, T> {
+//     type Error = RoiError;
+//     fn roi(&'a self, aabb: Aabb2<usize>) -> Result<ArrayViewMut3<'b, T>, Self::Error> {
+//         let x1 = aabb.x1();
+//         let x2 = aabb.x2();
+//         let y1 = aabb.y1();
+//         let y2 = aabb.y2();
+//         if x1 >= x2 || y1 >= y2 || x2 > self.shape()[1] || y2 > self.shape()[0] {
+//             return Err(RoiError::RoiOutOfBounds);
+//         }
+//         Ok(self.slice(ndarray::s![y1..y2, x1..x2, ..]))
+//     }
+// }
+
+impl<'a, 'b: 'a, T: Num> RoiMut<'a, ArrayViewMut3<'a, T>> for ArrayViewMut3<'b, T> {
+    type Error = RoiError;
+    fn roi_mut(&'a mut self, aabb: Aabb2<usize>) -> Result<ArrayViewMut3<'a, T>, Self::Error> {
+        let x1 = aabb.x1();
+        let x2 = aabb.x2();
+        let y1 = aabb.y1();
+        let y2 = aabb.y2();
+        if x1 >= x2 || y1 >= y2 || x2 > self.shape()[1] || y2 > self.shape()[0] {
+            return Err(RoiError::RoiOutOfBounds);
+        }
+        let out: ArrayViewMut3<'a, T> = self.slice_mut(ndarray::s![y1..y2, x1..x2, ..]);
+        Ok(out)
+    }
+}
+
+#[test]
+pub fn reborrow_test() {
+    let ndarray = ndarray::Array::from_shape_vec((5, 5, 5), vec![33; 5 * 5 * 5]).unwrap();
+    let aabb = Aabb2::from_x1y1x2y2(2, 3, 4, 5);
+    let y = {
+        let view = ndarray.view();
+        view.roi(aabb).unwrap()
+    };
+    dbg!(y);
+}