[Feature] NMS update

mmdet/core/post_processing/bbox_nms.py

@@ -31,15 +31,15 @@ def multiclass_nms(multi_bboxes,
         tuple: (dets, labels, indices (optional)), tensors of shape (k, 5),
             (k), and (k). Dets are boxes with scores. Labels are 0-based.
     """
-    num_classes = multi_scores.size(1) - 1
+    num_classes = int(multi_scores.size(1) - 1)
     # exclude background category
     if multi_bboxes.shape[1] > 4:
         bboxes = multi_bboxes.view(multi_scores.size(0), -1, 4)
     else:
         bboxes = multi_bboxes[:, None].expand(
             multi_scores.size(0), num_classes, 4)
 
-    scores = multi_scores[:, :-1]
+    scores = multi_scores[:, :num_classes]
 
     labels = torch.arange(num_classes, dtype=torch.long)
     labels = labels.view(1, -1).expand_as(scores)
@@ -82,11 +82,11 @@ def multiclass_nms(multi_bboxes,
         else:
             return dets, labels
 
-    dets, keep = batched_nms(bboxes, scores, labels, nms_cfg)
+    if nms_cfg.type == 'nms':
+        nms_cfg.score_threshold = score_thr if score_factors is None else 0
+        nms_cfg.max_num = max_num
 
-    if max_num > 0:
-        dets = dets[:max_num]
-        keep = keep[:max_num]
+    dets, keep = batched_nms(bboxes, scores, labels, nms_cfg)
 
     if return_inds:
         return dets, labels[keep], keep

mmdet/models/dense_heads/atss_head.py

@@ -449,7 +449,7 @@ def _get_bboxes(self,
                 anchors = anchors.expand_as(bbox_pred)
 
             bboxes = self.bbox_coder.decode(
-                anchors, bbox_pred, max_shape=img_shapes)
+                anchors, bbox_pred, max_shape=torch.tensor(img_shapes[0][:2]))
             mlvl_bboxes.append(bboxes)
             mlvl_scores.append(scores)
             mlvl_centerness.append(centerness)

mmdet/models/dense_heads/cascade_rpn_head.py

@@ -659,6 +659,8 @@ def _get_bboxes_single(self,
                     'Please delete max_num which will be deprecated.'
             else:
                 cfg.max_per_img = cfg.max_num
+        if 'max_per_img' in cfg:
+            cfg.nms.max_num = cfg.max_per_img
         if 'nms_thr' in cfg:
             assert cfg.nms.iou_threshold == cfg.nms_thr, f'You set' \
                 f' iou_threshold in nms and ' \
@@ -668,7 +670,7 @@ def _get_bboxes_single(self,
                 f'which will be deprecated.'
 
         dets, keep = batched_nms(proposals, scores, ids, cfg.nms)
-        return dets[:cfg.max_per_img]
+        return dets
 
 
 @HEADS.register_module()

mmdet/models/dense_heads/fovea_head.py

@@ -288,11 +288,11 @@ def get_bboxes(self,
             ]
             img_shape = img_metas[img_id]['img_shape']
             scale_factor = img_metas[img_id]['scale_factor']
-            det_bboxes = self._get_bboxes_single(cls_score_list,
+            det_result = self._get_bboxes_single(cls_score_list,
                                                  bbox_pred_list, featmap_sizes,
                                                  points, img_shape,
                                                  scale_factor, cfg, rescale)
-            result_list.append(det_bboxes)
+            result_list.append(det_result)
         return result_list
 
     def _get_bboxes_single(self,

mmdet/models/dense_heads/rpn_head.py

@@ -216,12 +216,14 @@ def _get_bboxes_single(self,
                 proposals = proposals[valid_mask]
                 scores = scores[valid_mask]
                 ids = ids[valid_mask]
+        if 'max_per_img' in cfg:
+            cfg.nms.max_num = cfg.max_per_img
         if proposals.numel() > 0:
             dets, keep = batched_nms(proposals, scores, ids, cfg.nms)
         else:
             return proposals.new_zeros(0, 5)
 
-        return dets[:cfg.max_per_img]
+        return dets
 
     def onnx_export(self, x, img_metas):
         """Test without augmentation.

mmdet/models/roi_heads/trident_roi_head.py

@@ -33,12 +33,10 @@ def merge_trident_bboxes(self, trident_det_bboxes, trident_det_labels):
             nms_scores = trident_det_bboxes[:, 4].contiguous()
             nms_inds = trident_det_labels
             nms_cfg = self.test_cfg['nms']
+            nms_cfg['max_num'] = self.test_cfg['max_per_img']
             det_bboxes, keep = batched_nms(nms_bboxes, nms_scores, nms_inds,
                                            nms_cfg)
             det_labels = trident_det_labels[keep]
-            if self.test_cfg['max_per_img'] > 0:
-                det_labels = det_labels[:self.test_cfg['max_per_img']]
-                det_bboxes = det_bboxes[:self.test_cfg['max_per_img']]
 
         return det_bboxes, det_labels