filter boxes with low score (#177)

pyronear · Oct 31, 2023 · 4051cd1 · 4051cd1
1 parent 32f6ffb
commit 4051cd1
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/pyroengine/engine.py b/pyroengine/engine.py
@@ -254,6 +254,9 @@ def _update_states(self, frame: Image.Image, preds: np.ndarray, cam_key: str) ->
             # Limit bbox size in api
             output_predictions = np.round(output_predictions, 3)  # max 3 digit
             output_predictions = output_predictions[:5, :]  # max 5 bbox
+            output_predictions = output_predictions[
+                output_predictions[:, 4] > self.conf_thresh
+            ]  # send only boxes above conf
 
         self._states[cam_key]["last_predictions"].append(
             (frame, preds, output_predictions.tolist(), datetime.utcnow().isoformat(), False)