feat: Added facenet

src/main.rs

@@ -1,9 +1,13 @@
 mod cli;
 mod errors;
-use detector::facedet::retinaface::FaceDetectionConfig;
+use bounding_box::roi::MultiRoi;
+use detector::{facedet::retinaface::FaceDetectionConfig, faceembed};
 use errors::*;
+use fast_image_resize::ResizeOptions;
+use nalgebra::zero;
 use ndarray_image::*;
 const RETINAFACE_MODEL: &[u8] = include_bytes!("../models/retinaface.mnn");
+const FACENET_MODEL: &[u8] = include_bytes!("../models/facenet.mnn");
 pub fn main() -> Result<()> {
     tracing_subscriber::fmt()
         .with_env_filter("trace")
@@ -15,29 +19,84 @@ pub fn main() -> Result<()> {
     match args.cmd {
         cli::SubCommand::Detect(detect) => {
             use detector::facedet;
-            let model = facedet::retinaface::FaceDetection::new_from_bytes(RETINAFACE_MODEL)
+            let retinaface = facedet::retinaface::FaceDetection::new_from_bytes(RETINAFACE_MODEL)
                 .change_context(errors::Error)
                 .attach_printable("Failed to create face detection model")?;
+            let facenet = faceembed::facenet::EmbeddingGenerator::new_from_bytes(FACENET_MODEL)
+                .change_context(errors::Error)
+                .attach_printable("Failed to create face embedding model")?;
             let image = image::open(detect.image).change_context(Error)?;
             let image = image.into_rgb8();
             let mut array = image
                 .into_ndarray()
                 .change_context(errors::Error)
                 .attach_printable("Failed to convert image to ndarray")?;
-            let output = model
+            let output = retinaface
                 .detect_faces(
-                    array.clone(),
+                    array.view(),
                     FaceDetectionConfig::default()
                         .with_threshold(detect.threshold)
                         .with_nms_threshold(detect.nms_threshold),
                 )
                 .change_context(errors::Error)
                 .attach_printable("Failed to detect faces")?;
-            for bbox in output.bbox {
+            for bbox in &output.bbox {
                 tracing::info!("Detected face: {:?}", bbox);
                 use bounding_box::draw::*;
                 array.draw(bbox, color::palette::css::GREEN_YELLOW.to_rgba8(), 1);
             }
+            use ndarray::{Array2, Array3, Array4, Axis};
+            use ndarray_resize::NdFir;
+            let face_rois = array
+                .view()
+                .multi_roi(&output.bbox)
+                .change_context(Error)?
+                .into_iter()
+                // .inspect(|f| {
+                //     tracing::info!("Face ROI shape before resize: {:?}", f.dim());
+                // })
+                .map(|roi| {
+                    roi.as_standard_layout()
+                        .fast_resize(512, 512, &ResizeOptions::default())
+                        .change_context(Error)
+                })
+                // .inspect(|f| {
+                //     f.as_ref().inspect(|f| {
+                //         tracing::info!("Face ROI shape after resize: {:?}", f.dim());
+                //     });
+                // })
+                .collect::<Result<Vec<_>>>()?;
+            let face_roi_views = face_rois.iter().map(|roi| roi.view()).collect::<Vec<_>>();
+
+            let embeddings = face_roi_views
+                .chunks(8)
+                .map(|chunk| {
+                    tracing::info!("Processing chunk of size: {}", chunk.len());
+
+                    if chunk.len() < 8 {
+                        tracing::warn!("Chunk size is less than 8, padding with zeros");
+                        let zeros = Array3::zeros((512, 512, 3));
+                        let padded: Vec<ndarray::ArrayView3<'_, u8>> = chunk
+                            .iter()
+                            .cloned()
+                            .chain(core::iter::repeat(zeros.view()))
+                            .take(8)
+                            .collect();
+                        let face_rois: Array4<u8> = ndarray::stack(Axis(0), padded.as_slice())
+                            .change_context(errors::Error)
+                            .attach_printable("Failed to stack rois together")?;
+                        let output = facenet.run_models(face_rois.view()).change_context(Error)?;
+                        Ok(output)
+                    } else {
+                        let face_rois: Array4<u8> = ndarray::stack(Axis(0), chunk)
+                            .change_context(errors::Error)
+                            .attach_printable("Failed to stack rois together")?;
+                        let output = facenet.run_models(face_rois.view()).change_context(Error)?;
+                        Ok(output)
+                    }
+                })
+                .collect::<Result<Vec<Array2<f32>>>>();
+
             let v = array.view();
             if let Some(output) = detect.output {
                 let image: image::RgbImage = v