fix OOB bugs in yolov5/7 and rfdetr

Author: leeclemnet

inference_models/inference_models/models/common/roboflow/post_processing.py

@@ -25,14 +25,11 @@ def run_nms_for_object_detection(
     boxes = output[:, :4, :]
     scores = output[:, 4:, :]
     results = []
-    per_class_thresh = (
-        conf_thresh.to(output.device) if isinstance(conf_thresh, torch.Tensor) else None
-    )
     for b in range(bs):
         class_scores = scores[b]
         class_conf, class_ids = class_scores.max(0)
-        if per_class_thresh is not None:
-            mask = class_conf > per_class_thresh[class_ids]
+        if isinstance(conf_thresh, torch.Tensor):
+            mask = class_conf > conf_thresh.to(output.device)[class_ids]
         else:
             mask = class_conf > conf_thresh
         if not torch.any(mask):
@@ -74,15 +71,12 @@ def post_process_nms_fused_model_output(
 ) -> List[torch.Tensor]:
     bs = output.shape[0]
     nms_results = []
-    per_class_thresh = (
-        conf_thresh.to(output.device) if isinstance(conf_thresh, torch.Tensor) else None
-    )
     for batch_element_id in range(bs):
         batch_element_result = output[batch_element_id]
-        if per_class_thresh is not None:
+        if isinstance(conf_thresh, torch.Tensor):
             class_ids = batch_element_result[:, 5].long()
             batch_element_result = batch_element_result[
-                batch_element_result[:, 4] >= per_class_thresh[class_ids]
+                batch_element_result[:, 4] >= conf_thresh.to(output.device)[class_ids]
             ]
         else:
             batch_element_result = batch_element_result[
@@ -105,17 +99,14 @@ def run_nms_for_instance_segmentation(
     scores = output[:, 4:-32, :]  # (N, 80, 8400)
     masks = output[:, -32:, :]
     results = []
-    per_class_thresh = (
-        conf_thresh.to(output.device) if isinstance(conf_thresh, torch.Tensor) else None
-    )
 
     for b in range(bs):
         bboxes = boxes[b].T  # (8400, 4)
         class_scores = scores[b].T  # (8400, 80)
         box_masks = masks[b].T
         class_conf, class_ids = class_scores.max(1)  # (8400,), (8400,)
-        if per_class_thresh is not None:
-            mask = class_conf > per_class_thresh[class_ids]
+        if isinstance(conf_thresh, torch.Tensor):
+            mask = class_conf > conf_thresh.to(output.device)[class_ids]
         else:
             mask = class_conf > conf_thresh
         if mask.sum() == 0:
@@ -164,14 +155,11 @@ def run_nms_for_key_points_detection(
     scores = output[:, 4 : 4 + num_classes, :]
     key_points = output[:, 4 + num_classes :, :]
     results = []
-    per_class_thresh = (
-        conf_thresh.to(output.device) if isinstance(conf_thresh, torch.Tensor) else None
-    )
     for b in range(bs):
         class_scores = scores[b]
         class_conf, class_ids = class_scores.max(0)
-        if per_class_thresh is not None:
-            mask = class_conf > per_class_thresh[class_ids]
+        if isinstance(conf_thresh, torch.Tensor):
+            mask = class_conf > conf_thresh.to(output.device)[class_ids]
         else:
             mask = class_conf > conf_thresh
         if not torch.any(mask):

inference_models/inference_models/models/rfdetr/common.py

@@ -66,6 +66,13 @@ def post_process_instance_segmentation_results(
             confidence = confidence[remapping_mask]
             image_bboxes = image_bboxes[remapping_mask]
             image_masks = image_masks[remapping_mask]
+        else:
+            # drop DETR no-object rows
+            named = top_classes < threshold.shape[0]
+            confidence = confidence[named]
+            top_classes = top_classes[named]
+            image_bboxes = image_bboxes[named]
+            image_masks = image_masks[named]
         confidence_mask = confidence > threshold[top_classes.long()]
         confidence = confidence[confidence_mask]
         top_classes = top_classes[confidence_mask]

inference_models/inference_models/models/rfdetr/rfdetr_object_detection_onnx.py

@@ -238,6 +238,12 @@ def post_process(
                 ]
                 predicted_confidence = predicted_confidence[remapping_mask]
                 image_bboxes = image_bboxes[remapping_mask]
+            else:
+                # drop DETR no-object rows
+                named = top_classes < len(self.class_names)
+                predicted_confidence = predicted_confidence[named]
+                top_classes = top_classes[named]
+                image_bboxes = image_bboxes[named]
             confidence_mask = predicted_confidence > thresholds[top_classes.long()]
             predicted_confidence = predicted_confidence[confidence_mask]
             top_classes = top_classes[confidence_mask]

inference_models/inference_models/models/rfdetr/rfdetr_object_detection_pytorch.py

@@ -493,6 +493,12 @@ def post_process(
                 labels = self._classes_re_mapping.class_mapping[labels[remapping_mask]]
                 boxes = boxes[remapping_mask]
             score_thresholds = thresholds.to(dtype=scores.dtype)
+            if self._classes_re_mapping is None:
+                # drop DETR no-object rows
+                named = labels < score_thresholds.shape[0]
+                scores = scores[named]
+                labels = labels[named]
+                boxes = boxes[named]
             keep = scores > score_thresholds[labels.long()]
             scores = scores[keep]
             labels = labels[keep]

inference_models/inference_models/models/rfdetr/rfdetr_object_detection_trt.py

@@ -295,6 +295,12 @@ def post_process(
                     ]
                     predicted_confidence = predicted_confidence[remapping_mask]
                     image_bboxes = image_bboxes[remapping_mask]
+                else:
+                    # drop DETR no-object rows
+                    named = top_classes < len(self.class_names)
+                    predicted_confidence = predicted_confidence[named]
+                    top_classes = top_classes[named]
+                    image_bboxes = image_bboxes[named]
                 confidence_mask = predicted_confidence > thresholds[top_classes.long()]
                 predicted_confidence = predicted_confidence[confidence_mask]
                 top_classes = top_classes[confidence_mask]

inference_models/inference_models/models/yolact/yolact_instance_segmentation_onnx.py

@@ -338,16 +338,13 @@ def run_nms_for_instance_segmentation(
     scores = output[:, :, 4:-32]  # (N, 19248, num_classes)
     masks = output[:, :, -32:]
     results = []
-    per_class_thresh = (
-        conf_thresh.to(output.device) if isinstance(conf_thresh, torch.Tensor) else None
-    )
     for b in range(bs):
         bboxes = boxes[b]  # (19248, 4)
         class_scores = scores[b]  # (19248, 80)
         box_masks = masks[b]
         class_conf, class_ids = class_scores.max(1)  # (8400,), (8400,)
-        if per_class_thresh is not None:
-            mask = class_conf > per_class_thresh[class_ids]
+        if isinstance(conf_thresh, torch.Tensor):
+            mask = class_conf > conf_thresh.to(output.device)[class_ids]
         else:
             mask = class_conf > conf_thresh
         if mask.sum() == 0:

inference_models/inference_models/models/yolact/yolact_instance_segmentation_trt.py

@@ -422,16 +422,13 @@ def run_nms_for_instance_segmentation(
     scores = output[:, :, 4:-32]  # (N, 19248, num_classes)
     masks = output[:, :, -32:]
     results = []
-    per_class_thresh = (
-        conf_thresh.to(output.device) if isinstance(conf_thresh, torch.Tensor) else None
-    )
     for b in range(bs):
         bboxes = boxes[b]  # (19248, 4)
         class_scores = scores[b]  # (19248, 80)
         box_masks = masks[b]
         class_conf, class_ids = class_scores.max(1)  # (8400,), (8400,)
-        if per_class_thresh is not None:
-            mask = class_conf > per_class_thresh[class_ids]
+        if isinstance(conf_thresh, torch.Tensor):
+            mask = class_conf > conf_thresh.to(output.device)[class_ids]
         else:
             mask = class_conf > conf_thresh
         if mask.sum() == 0:

inference_models/inference_models/models/yolonas/nms.py

@@ -15,14 +15,11 @@ def run_yolonas_nms_for_object_detection(
     boxes = output[:, :, :4]
     scores = output[:, :, 4:]
     results = []
-    per_class_thresh = (
-        conf_thresh.to(output.device) if isinstance(conf_thresh, torch.Tensor) else None
-    )
     for b in range(bs):
         class_scores = scores[b]  # (8400, cls_num)
         class_conf, class_ids = torch.max(class_scores, dim=-1)
-        if per_class_thresh is not None:
-            mask = class_conf > per_class_thresh[class_ids]
+        if isinstance(conf_thresh, torch.Tensor):
+            mask = class_conf > conf_thresh.to(output.device)[class_ids]
         else:
             mask = class_conf > conf_thresh
         if not torch.any(mask):

inference_models/inference_models/models/yolov5/nms.py

@@ -16,14 +16,11 @@ def run_nms_yolov5(
     top_classes_conf = output[:, 4, :]
     scores = output[:, 5:, :]
     results = []
-    per_class_thresh = (
-        conf_thresh.to(output.device) if isinstance(conf_thresh, torch.Tensor) else None
-    )
     for b in range(bs):
         class_scores = scores[b]
         class_conf, class_ids = class_scores.max(0)
-        if per_class_thresh is not None:
-            mask = top_classes_conf[b] > per_class_thresh[class_ids]
+        if isinstance(conf_thresh, torch.Tensor):
+            mask = top_classes_conf[b] > conf_thresh.to(output.device)[class_ids]
         else:
             mask = top_classes_conf[b] > conf_thresh
         if not torch.any(mask):
@@ -68,17 +65,20 @@ def run_yolov5_nms_for_instance_segmentation(
     scores = output[:, 4:-32, :]
     masks = output[:, -32:, :]
     results = []
-    per_class_thresh = (
-        conf_thresh.to(output.device) if isinstance(conf_thresh, torch.Tensor) else None
-    )
 
     for b in range(bs):
         bboxes = boxes[b].T
         class_scores = scores[b].T
         box_masks = masks[b].T
         class_conf, class_ids = class_scores.max(1)
-        if per_class_thresh is not None:
-            mask = top_classes_conf[b] > per_class_thresh[class_ids]
+        if isinstance(conf_thresh, torch.Tensor):
+            # class_ids are slice-indexed: 0 is objectness, k>=1 is class_{k-1}.
+            # Keep objectness-dominated rows (no real class to threshold) and
+            # filter class-dominated rows against the real class's threshold.
+            thresh = conf_thresh.to(output.device)
+            real_class = (class_ids - 1).clamp(min=0, max=thresh.shape[0] - 1)
+            is_obj = class_ids == 0
+            mask = is_obj | (top_classes_conf[b] > thresh[real_class])
         else:
             mask = top_classes_conf[b] > conf_thresh
         if mask.sum() == 0:

inference_models/inference_models/models/yolov7/yolov7_instance_segmentation_onnx.py

@@ -38,7 +38,9 @@
     align_instance_segmentation_results,
     crop_masks_to_boxes,
     preprocess_segmentation_masks,
-    run_nms_for_instance_segmentation,
+)
+from inference_models.models.yolov5.nms import (
+    run_yolov5_nms_for_instance_segmentation,
 )
 from inference_models.models.common.roboflow.pre_processing import (
     pre_process_network_input,
@@ -214,7 +216,7 @@ def post_process(
         )
         confidence = confidence_filter.per_class_thresholds(self.class_names)
         instances, protos = model_results
-        nms_results = run_nms_for_instance_segmentation(
+        nms_results = run_yolov5_nms_for_instance_segmentation(
             output=instances.permute(0, 2, 1),
             conf_thresh=confidence,
             iou_thresh=iou_threshold,