feat: Added threshold for scores and nms

src/cli.rs

@@ -47,6 +47,8 @@ pub struct Detect {
     pub model_type: Models,
     #[clap(short, long)]
     pub output: Option<PathBuf>,
+    #[clap(short, long, default_value_t = 0.8)]
+    pub threshold: f32,
     pub image: PathBuf,
 }
 

src/facedet/retinaface.rs

@@ -10,6 +10,31 @@ pub struct FaceDetectionConfig {
     min_sizes: Vec<Vector2<usize>>,
     steps: Vec<usize>,
     variance: Vec<f32>,
+    threshold: f32,
+    nms_threshold: f32,
+}
+
+impl FaceDetectionConfig {
+    pub fn with_min_sizes(mut self, min_sizes: Vec<Vector2<usize>>) -> Self {
+        self.min_sizes = min_sizes;
+        self
+    }
+    pub fn with_steps(mut self, steps: Vec<usize>) -> Self {
+        self.steps = steps;
+        self
+    }
+    pub fn with_variance(mut self, variance: Vec<f32>) -> Self {
+        self.variance = variance;
+        self
+    }
+    pub fn with_threshold(mut self, threshold: f32) -> Self {
+        self.threshold = threshold;
+        self
+    }
+    pub fn with_nms_threshold(mut self, nms_threshold: f32) -> Self {
+        self.nms_threshold = nms_threshold;
+        self
+    }
 }
 
 impl Default for FaceDetectionConfig {
@@ -22,6 +47,8 @@ impl Default for FaceDetectionConfig {
             ],
             steps: vec![8, 16, 32],
             variance: vec![0.1, 0.2],
+            threshold: 0.8,
+            nms_threshold: 0.6,
         }
     }
 }
@@ -35,8 +62,13 @@ pub struct FaceDetectionModelOutput {
     pub landmark: ndarray::Array3<f32>,
 }
 
+pub struct FaceDetectionProcessedOutput {
+    pub bbox: Vec<Aabb2<f32>>,
+    pub confidence: Vec<f32>,
+}
+
 impl FaceDetectionModelOutput {
-    pub fn postprocess(self, config: FaceDetectionConfig) -> Result<Vec<Aabb2<f32>>> {
+    pub fn postprocess(self, config: FaceDetectionConfig) -> Result<FaceDetectionProcessedOutput> {
         let mut anchors = Vec::new();
         for (k, &step) in config.steps.iter().enumerate() {
             let feature_size = 640 / step;
@@ -54,6 +86,7 @@ impl FaceDetectionModelOutput {
             }
         }
         let mut boxes = Vec::new();
+        let mut scores = Vec::new();
         let var0 = config.variance[0];
         let var1 = config.variance[1];
         let bbox_data = self.bbox;
@@ -74,14 +107,15 @@ impl FaceDetectionModelOutput {
             let x_max = pred_cx + pred_w / 2.0;
             let y_max = pred_cy + pred_h / 2.0;
             let score = conf_data[[0, idx, 1]];
-            if score > 0.6 {
-                boxes.push(Aabb2::from_min_max_vertices(
-                    Point2::new(x_min, y_min),
-                    Point2::new(x_max, y_max),
-                ));
+            if score > config.threshold {
+                boxes.push(Aabb2::from_x1y1x2y2(x_min, y_min, x_max, y_max));
+                scores.push(score);
             }
         }
-        Ok(boxes)
+        Ok(FaceDetectionProcessedOutput {
+            bbox: boxes,
+            confidence: scores,
+        })
     }
 }
 

src/main.rs

@@ -1,5 +1,6 @@
 mod cli;
 mod errors;
+use detector::facedet::retinaface::FaceDetectionConfig;
 use errors::*;
 use ndarray_image::*;
 const RETINAFACE_MODEL: &[u8] = include_bytes!("../models/retinaface.mnn");
@@ -29,11 +30,11 @@ pub fn main() -> Result<()> {
                 .attach_printable("Failed to detect faces")?;
             // output.print(20);
             let aabbs = output
-                .postprocess(Default::default())
+                .postprocess(FaceDetectionConfig::default().with_threshold(detect.threshold))
                 .change_context(errors::Error)
                 .attach_printable("Failed to attach context")?;
             for bbox in aabbs {
-                println!("Detected face: {:?}", bbox);
+                tracing::info!("Detected face: {:?}", bbox);
                 use bounding_box::draw::*;
                 let bbox = bbox
                     .denormalize(nalgebra::SVector::<f32, 2>::new(