From 15ce8f83e4ed3ffaffe4bf65f569c71da9b1c77c Mon Sep 17 00:00:00 2001 From: aahmetdurmaz Date: Fri, 4 Sep 2026 12:27:23 +0300 Subject: [PATCH] fix: add support for customizable score and NMS thresholds --- frigate/detectors/plugins/deepx.py | 8 +++++- frigate/test/test_post_processing.py | 31 +++++++++++++++++++-- frigate/util/model.py | 41 +++++++++++++++++++++------- 3 files changed, 66 insertions(+), 14 deletions(-) diff --git a/frigate/detectors/plugins/deepx.py b/frigate/detectors/plugins/deepx.py index 0bb4efc543..4da5e973ff 100644 --- a/frigate/detectors/plugins/deepx.py +++ b/frigate/detectors/plugins/deepx.py @@ -1136,7 +1136,13 @@ class DeepxDetector(DetectionApi): self.nms_threshold, ) - return post_process_yolo(outputs, self.width, self.height) + return post_process_yolo( + outputs, + self.width, + self.height, + self.score_threshold, + self.nms_threshold, + ) def detect_raw(self, tensor_input): """Run inference and decode the DEEPX model output.""" diff --git a/frigate/test/test_post_processing.py b/frigate/test/test_post_processing.py index 9290c0c152..6fd33783a7 100644 --- a/frigate/test/test_post_processing.py +++ b/frigate/test/test_post_processing.py @@ -81,9 +81,21 @@ class TestYoloNmsPostProcess(unittest.TestCase): self.assertEqual(len(detections), 1) np.testing.assert_allclose(detections[0], [0, 0.90, *A_ROW], atol=2e-3) + def test_honors_caller_supplied_thresholds(self): + """A detector with its own thresholds must not be held to the 0.4 + default, which would drop this 0.3 confidence box.""" + output = self._single_output([[A_CX, A_CY, A_W, A_H, 0.30, 0.0]]) + + self.assertEqual(len(kept(post_process_yolo(output, WIDTH, HEIGHT))), 0) + + detections = kept(post_process_yolo(output, WIDTH, HEIGHT, 0.25, 0.45)) + + self.assertEqual(len(detections), 1) + np.testing.assert_allclose(detections[0], [0, 0.30, *A_ROW], atol=2e-3) + class TestMultipartYoloPostProcess(unittest.TestCase): - def _multipart_output(self) -> list[np.ndarray]: + def _multipart_output(self, a_conf=0.95, b_conf=0.90) -> list[np.ndarray]: """Build a 3-scale anchor-based YOLO output containing boxes A and B, both decoded through anchor 0 of the stride-32 scale.""" outputs = [ @@ -94,8 +106,8 @@ class TestMultipartYoloPostProcess(unittest.TestCase): stride, (anchor_w, anchor_h) = 32, (142, 110) for cx, cy, w, h, conf, class_channel in [ - (A_CX, A_CY, A_W, A_H, 0.95, 5), # class 0 - (B_CX, B_CY, B_W, B_H, 0.90, 6), # class 1 + (A_CX, A_CY, A_W, A_H, a_conf, 5), # class 0 + (B_CX, B_CY, B_W, B_H, b_conf, 6), # class 1 ]: cell_x, cell_y = int(cx // stride), int(cy // stride) dx = (cx / stride - cell_x + 0.5) / 2 @@ -127,6 +139,19 @@ class TestMultipartYoloPostProcess(unittest.TestCase): self.assertEqual(len(detections), 0) + def test_honors_caller_supplied_thresholds(self): + """Same as the single-tensor case: the multi-part branch must filter + at the caller's threshold, not the 0.4 default.""" + outputs = self._multipart_output(a_conf=0.35, b_conf=0.30) + + self.assertEqual(len(kept(post_process_yolo(outputs, WIDTH, HEIGHT))), 0) + + detections = kept(post_process_yolo(outputs, WIDTH, HEIGHT, 0.25, 0.45)) + + self.assertEqual(len(detections), 2) + np.testing.assert_allclose(detections[0], [0, 0.35, *A_ROW], atol=2e-3) + np.testing.assert_allclose(detections[1], [1, 0.30, *B_ROW], atol=2e-3) + class TestYoloxPostProcess(unittest.TestCase): def test_keeps_separated_objects_far_from_origin(self): diff --git a/frigate/util/model.py b/frigate/util/model.py index 80f0d29a04..e0d1d2f961 100644 --- a/frigate/util/model.py +++ b/frigate/util/model.py @@ -134,6 +134,8 @@ def __post_process_multipart_yolo( output_list, width, height, + score_threshold=0.4, + nms_threshold=0.4, ): anchors = [ [(12, 16), (19, 36), (40, 28)], @@ -166,7 +168,7 @@ def __post_process_multipart_yolo( class_conf = class_probs[class_id] conf = class_conf * pred[4] - if conf < 0.4: + if conf < score_threshold: continue dx = pred[0] @@ -191,8 +193,8 @@ def __post_process_multipart_yolo( indices = cv2.dnn.NMSBoxes( bboxes=xyxy_to_xywh_for_nms(all_boxes), scores=all_scores, - score_threshold=0.4, - nms_threshold=0.4, + score_threshold=score_threshold, + nms_threshold=nms_threshold, ) results = np.zeros((20, 6), np.float32) @@ -214,7 +216,13 @@ def __post_process_multipart_yolo( return results -def __post_process_nms_yolo(predictions: np.ndarray, width, height) -> np.ndarray: +def __post_process_nms_yolo( + predictions: np.ndarray, + width, + height, + score_threshold=0.4, + nms_threshold=0.4, +) -> np.ndarray: predictions = np.squeeze(predictions) # transpose the output so it has order (inferences, class_ids) @@ -222,8 +230,8 @@ def __post_process_nms_yolo(predictions: np.ndarray, width, height) -> np.ndarra predictions = predictions.T scores = np.max(predictions[:, 4:], axis=1) - predictions = predictions[scores > 0.4, :] - scores = scores[scores > 0.4] + predictions = predictions[scores > score_threshold, :] + scores = scores[scores > score_threshold] class_ids = np.argmax(predictions[:, 4:], axis=1) # Rescale box @@ -237,7 +245,10 @@ def __post_process_nms_yolo(predictions: np.ndarray, width, height) -> np.ndarra # run NMS indices = cv2.dnn.NMSBoxes( - xyxy_to_xywh_for_nms(boxes), scores, score_threshold=0.4, nms_threshold=0.4 + xyxy_to_xywh_for_nms(boxes), + scores, + score_threshold=score_threshold, + nms_threshold=nms_threshold, ) detections = np.zeros((20, 6), np.float32) for i, (bbox, confidence, class_id) in enumerate( @@ -258,11 +269,21 @@ def __post_process_nms_yolo(predictions: np.ndarray, width, height) -> np.ndarra return detections -def post_process_yolo(output: list[np.ndarray], width: int, height: int) -> np.ndarray: +def post_process_yolo( + output: list[np.ndarray], + width: int, + height: int, + score_threshold: float = 0.4, + nms_threshold: float = 0.4, +) -> np.ndarray: if len(output) > 1: - return __post_process_multipart_yolo(output, width, height) + return __post_process_multipart_yolo( + output, width, height, score_threshold, nms_threshold + ) else: - return __post_process_nms_yolo(output[0], width, height) + return __post_process_nms_yolo( + output[0], width, height, score_threshold, nms_threshold + ) def post_process_yolox(