feat: Added postprocessing for retinaface

Cargo.lock

@@ -255,6 +255,8 @@ dependencies = [
  "nalgebra",
  "ndarray 0.16.1",
  "num",
+ "simba",
+ "thiserror 2.0.12",
 ]
 
 [[package]]
@@ -498,6 +500,7 @@ dependencies = [
  "bounding-box",
  "clap",
  "clap_complete",
+ "color",
  "error-stack",
  "fast_image_resize",
  "image",

Cargo.toml

@@ -50,6 +50,7 @@ mnn = { workspace = true }
 mnn-bridge = { workspace = true }
 mnn-sync = { workspace = true }
 bounding-box = { version = "0.1.0", path = "bounding-box" }
+color = "0.3.1"
 
 [profile.release]
 debug = true

bounding-box/Cargo.toml

@@ -9,6 +9,8 @@ itertools = "0.14.0"
 nalgebra = "0.33.2"
 ndarray = { version = "0.16.1", optional = true }
 num = "0.4.3"
+simba = "0.9.0"
+thiserror = "2.0.12"
 
 [features]
 ndarray = ["dep:ndarray"]

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> {
+impl Draw<Aabb2<usize>> for Array3<u8> {
-//     fn draw(&self, item: T, color: color::Rgba8, thickness: usize) {
+    fn draw(&mut self, item: Aabb2<usize>, color: color::Rgba8, thickness: usize) {
-//         item.draw(&self, color, thickness);
+        item.draw(self, color, thickness)
-//     }
+    }
-// }
+}
 
-pub trait Drawable<Canvas, T> {
+pub trait Drawable<Canvas> {
-    fn draw(&self, canvas: &mut Canvas, color: color::Rgba8, thickness: T);
+    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>
+impl Drawable<ArrayViewMut3<'_, u8>> for Aabb2<usize> {
-where
+    fn draw(&self, canvas: &mut ArrayViewMut3<u8>, color: color::Rgba8, thickness: usize) {
-    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();
+        // 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]);
+                let bbox = Aabb2::from_vertices([*a, b, *c, d]).expect("Invalid bounding box");
-                todo!();
+                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);
+}

src/cli.rs

@@ -45,6 +45,8 @@ pub struct Detect {
     pub model: Option<PathBuf>,
     #[clap(short = 'M', long, default_value = "retina-face")]
     pub model_type: Models,
+    #[clap(short, long)]
+    pub output: Option<PathBuf>,
     pub image: PathBuf,
 }
 

src/main.rs

@@ -19,17 +19,45 @@ pub fn main() -> Result<()> {
                 .attach_printable("Failed to create face detection model")?;
             let image = image::open(detect.image).change_context(Error)?;
             let image = image.into_rgb8();
-            let array = image
+            let mut array = image
                 .into_ndarray()
                 .change_context(errors::Error)
                 .attach_printable("Failed to convert image to ndarray")?;
             let output = model
-                .detect_faces(array)
+                .detect_faces(array.clone())
                 .change_context(errors::Error)
                 .attach_printable("Failed to detect faces")?;
             // output.print(20);
-            let aabbs = output.postprocess(Default::default());
+            let aabbs = output
-            dbg!(aabbs);
+                .postprocess(Default::default())
+                .change_context(errors::Error)
+                .attach_printable("Failed to attach context")?;
+            for bbox in aabbs {
+                println!("Detected face: {:?}", bbox);
+                use bounding_box::draw::*;
+                let bbox = bbox
+                    .denormalize(nalgebra::SVector::<f32, 2>::new(
+                        array.shape()[1] as f32,
+                        array.shape()[0] as f32,
+                    ))
+                    .cast()
+                    .ok_or(errors::Error)
+                    .attach_printable("Failed to cast f32 to usize")?;
+                dbg!(bbox);
+                array.draw(bbox, color::palette::css::GREEN_YELLOW.to_rgba8(), 20);
+                break;
+            }
+            let v = array.view();
+            if let Some(output) = detect.output {
+                let image: image::RgbImage = v
+                    .to_image()
+                    .change_context(errors::Error)
+                    .attach_printable("Failed to convert ndarray to image")?;
+                image
+                    .save(output)
+                    .change_context(errors::Error)
+                    .attach_printable("Failed to save output image")?;
+            }
         }
         cli::SubCommand::List(list) => {
             println!("List: {:?}", list);