[Bug-fix]Add support for max_detection_per_image hyperparam (#543)

* add support for max_detection_per_image hyperparam

* format with black

* update doc strings

* changing default to -1 so that it bypasses by default and filter only if max_detections is less than no of outputs

Author: RihaanBH-1810

app.py

@@ -27,7 +27,7 @@ def browse_dataset_path():
 st.session_state.setdefault("config_option", "Manual Configuration")
 st.session_state.setdefault("confidence_threshold", 0.5)
 st.session_state.setdefault("nms_threshold", 0.5)
-st.session_state.setdefault("max_detections", 100)
+st.session_state.setdefault("max_detections", -1)
 st.session_state.setdefault("device", best_device)
 st.session_state.setdefault("batch_size", 1)
 st.session_state.setdefault("evaluation_step", 5)
@@ -122,7 +122,7 @@ def browse_dataset_path():
                 )
                 st.number_input(
                     "Max Detections/Image",
-                    min_value=1,
+                    min_value=-1,
                     max_value=1000,
                     step=1,
                     key="max_detections",
@@ -307,7 +307,7 @@ def browse_dataset_path():
                         st.session_state.get("confidence_threshold", 0.5)
                     )
                     nms_threshold = float(st.session_state.get("nms_threshold", 0.5))
-                    max_detections = int(st.session_state.get("max_detections", 100))
+                    max_detections = int(st.session_state.get("max_detections", -1))
                     device = st.session_state.get("device", "cpu")
                     batch_size = int(st.session_state.get("batch_size", 1))
                     evaluation_step = int(st.session_state.get("evaluation_step", 5))

perceptionmetrics/models/torch_detection.py

@@ -298,10 +298,14 @@ def forward(self, x):
         # Load confidence and NMS thresholds from config
         self.confidence_threshold = self.model_cfg.get("confidence_threshold", 0.5)
         self.nms_threshold = self.model_cfg.get("nms_threshold", 0.3)
+        self.max_detections_per_image = self.model_cfg.get(
+            "max_detections_per_image", -1
+        )
 
         self.postprocess_args = [self.confidence_threshold]
         if self.model_format == "yolo":
             self.postprocess_args.append(self.nms_threshold)
+        self.postprocess_args.append(self.max_detections_per_image)
 
         # Add reverse mapping for idx to class_name
         self.idx_to_class_name = {v["idx"]: k for k, v in self.ontology.items()}

perceptionmetrics/models/utils/torchvision.py

@@ -1,10 +1,14 @@
-def postprocess_detection(output: dict, confidence_threshold: float = 0.5):
+def postprocess_detection(
+    output: dict, confidence_threshold: float = 0.5, max_detections: int = -1
+):
     """Post-process torchvision model output.
 
     :param output: Dictionary with keys 'boxes', 'labels', and 'scores'.
     :type output: dict
     :param confidence_threshold: Confidence threshold to filter boxes.
     :type confidence_threshold: float
+    :param max_detections: Maximum number of best detections to keep per image after filtering.
+    :type max_detections: int
     :return: Dictionary with keys 'boxes', 'labels', and 'scores'.
     :rtype: dict
     """
@@ -15,4 +19,13 @@ def postprocess_detection(output: dict, confidence_threshold: float = 0.5):
             "labels": output["labels"][keep_mask],
             "scores": output["scores"][keep_mask],
         }
+
+    if max_detections < output["scores"].shape[0] and max_detections > 0:
+        limited_idx = output["scores"].argsort(descending=True)[:max_detections]
+        output = {
+            "boxes": output["boxes"][limited_idx],
+            "labels": output["labels"][limited_idx],
+            "scores": output["scores"][limited_idx],
+        }
+
     return output

perceptionmetrics/models/utils/yolo.py

@@ -1,14 +1,14 @@
 import torch
 from torchvision.ops import nms
 
-
 CLASS_NMS_OFFSET = 7680  # offset to apply to boxes for class-wise NMS
 
 
 def postprocess_detection(
     output: torch.Tensor,
     confidence_threshold: float = 0.25,
     nms_threshold: float = 0.45,
+    max_detections: int = -1,
 ):
     """Post-process YOLO model output.
 
@@ -18,6 +18,8 @@ def postprocess_detection(
     :type confidence_threshold: float
     :param nms_threshold: IoU threshold for Non-Maximum Suppression (NMS). Some models may not perform NMS (e.g. YOLOv26).
     :type nms_threshold: float
+    :param max_detections: Maximum number of best detections to keep per image after filtering.
+    :type max_detections: int
     :return: Dictionary with keys 'boxes', 'labels', and 'scores'.
     :rtype: dict
     """
@@ -57,4 +59,10 @@ def postprocess_detection(
         scores = scores[keep_idx]
         labels = labels[keep_idx]
 
+        if max_detections > 0 and max_detections < scores.shape[0]:
+            limited_idx = scores.argsort(descending=True)[:max_detections]
+            boxes_xyxy = boxes_xyxy[limited_idx]
+            scores = scores[limited_idx]
+            labels = labels[limited_idx]
+
     return {"boxes": boxes_xyxy, "labels": labels, "scores": scores}