Update bbox_nms.py

HikariTJU · Feb 3, 2022 · 66674c5 · 66674c5
1 parent 971d0cc
commit 66674c5
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/mmdet/core/post_processing/bbox_nms.py b/mmdet/core/post_processing/bbox_nms.py
@@ -173,16 +173,21 @@ def multiclass_nms(multi_bboxes,
             scores = scores[keep][:max_num]
             labels = labels[keep][:max_num]
             bboxes = bboxes[keep][:max_num]
+            keep = keep[:max_num]
         dets = torch.cat([bboxes, scores[:, None]], dim=1)
+        if return_inds:
+            return dets, labels, keep
+        else:
+            return dets, labels    
     elif nms_cfg['type']=='nms':    # Original NMS
         dets, keep = batched_nms(bboxes, scores, labels, nms_cfg)
 
         if max_num > 0:
             dets = dets[:max_num]
-            labels = labels[keep][:max_num]
+            keep = keep[:max_num]
 
     if return_inds:
-        return dets, labels[keep], keep
+        return dets, labels[keep], inds[keep]
     else:
         return dets, labels[keep]