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