Fix nms error of CornerNet when there is no prediction box (open-mmla…

…b#4409)
Clearlangw · Jan 10, 2021 · 6cdb0a3 · 6cdb0a3
1 parent 4689a73
commit 6cdb0a3
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/mmdet/models/dense_heads/corner_head.py b/mmdet/models/dense_heads/corner_head.py
@@ -762,6 +762,8 @@ def _get_bboxes_single(self,
         return detections, labels
 
     def _bboxes_nms(self, bboxes, labels, cfg):
+        if labels.numel() == 0:
+            return bboxes, labels
         out_bboxes, keep = batched_nms(bboxes[:, :4], bboxes[:, -1], labels,
                                        cfg.nms_cfg)
         out_labels = labels[keep]