feat: Added threshold for scores and nms

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,
+        })
     }
 }