reppoint bug fix when negative image training

open-mmlab · Jul 5, 2022 · aedb8ae · aedb8ae
1 parent 01439d1
commit aedb8ae
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/mmrotate/core/bbox/assigners/max_convex_iou_assigner.py b/mmrotate/core/bbox/assigners/max_convex_iou_assigner.py
@@ -212,6 +212,8 @@ def convex_overlaps(self, gt_rbboxes, points):
             overlaps (torch.Tensor): Overlaps between k gt_bboxes and n \
                 bboxes, shape(k, n).
         """
+        if gt_rbboxes.size(0) == 0:
+            return gt_rbboxes.new_zeros((0, points.size(0)))
         overlaps = convex_iou(points, gt_rbboxes)
         overlaps = overlaps.transpose(1, 0)
         return overlaps
diff --git a/mmrotate/models/dense_heads/rotated_reppoints_head.py b/mmrotate/models/dense_heads/rotated_reppoints_head.py
@@ -822,6 +822,9 @@ def get_pos_loss(self, cls_score, pts_pred, label, bbox_gt, label_weight,
         Returns:
             Tensor: Losses of all positive samples in single image.
         """
+        if pos_inds.size(0) == 0:
+            pos_loss = bbox_gt.new_zeros(0)
+            return pos_loss,
         pos_scores = cls_score[pos_inds]
         pos_pts_pred = pts_pred[pos_inds]
         pos_bbox_gt = bbox_gt[pos_inds]