keep all preds

pyronear · Jul 18, 2023 · 68ccd14 · 68ccd14
1 parent bd85bd0
commit 68ccd14
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/pyroengine/vision.py b/pyroengine/vision.py
@@ -62,11 +62,13 @@ def __call__(self, pil_img: Image.Image) -> np.ndarray:
 
         # ONNX inference
         y = self.ort_session.run(["output0"], {"images": np_img})[0][0]
-        y = y[:, y[-1, :] > 0.1]
+        # Post processing
         y = np.transpose(y)
         y = xywh2xyxy(y)
+        # Sort by confidence
         y = y[y[:, 4].argsort()]
         y = NMS(y)
+        # Normalize preds
         if len(y) > 0:
             y[:, :4:2] /= self.img_size[1]
             y[:, 1:4:2] /= self.img_size[0]