diff --git a/ppdet/modeling/assigners/atss_assigner.py b/ppdet/modeling/assigners/atss_assigner.py index 8d8e555ae024d75f4bdfe26cbf39e7f954f1a5b0..4835dd1d563b735a070836b0d18980483dfe0814 100644 --- a/ppdet/modeling/assigners/atss_assigner.py +++ b/ppdet/modeling/assigners/atss_assigner.py @@ -114,6 +114,14 @@ class ATSSAssigner(nn.Layer): num_anchors, _ = anchor_bboxes.shape batch_size, num_max_boxes, _ = gt_bboxes.shape + # negative batch + if num_max_boxes == 0: + assigned_labels = paddle.full([batch_size, num_anchors], bg_index) + assigned_bboxes = paddle.zeros([batch_size, num_anchors, 4]) + assigned_scores = paddle.zeros( + [batch_size, num_anchors, self.num_classes]) + return assigned_labels, assigned_bboxes, assigned_scores + # 1. compute iou between gt and anchor bbox, [B, n, L] ious = iou_similarity(gt_bboxes.reshape([-1, 4]), anchor_bboxes) ious = ious.reshape([batch_size, -1, num_anchors]) diff --git a/ppdet/modeling/assigners/task_aligned_assigner.py b/ppdet/modeling/assigners/task_aligned_assigner.py index ed43c4936bbf3f3ea8b5ebad0c283d7d46238b0e..bae00cc7cac72289e7c2cb2d7bab65b13737e679 100644 --- a/ppdet/modeling/assigners/task_aligned_assigner.py +++ b/ppdet/modeling/assigners/task_aligned_assigner.py @@ -78,6 +78,14 @@ class TaskAlignedAssigner(nn.Layer): batch_size, num_anchors, num_classes = pred_scores.shape _, num_max_boxes, _ = gt_bboxes.shape + # negative batch + if num_max_boxes == 0: + assigned_labels = paddle.full([batch_size, num_anchors], bg_index) + assigned_bboxes = paddle.zeros([batch_size, num_anchors, 4]) + assigned_scores = paddle.zeros( + [batch_size, num_anchors, num_classes]) + return assigned_labels, assigned_bboxes, assigned_scores + # compute iou between gt and pred bbox, [B, n, L] ious = iou_similarity(gt_bboxes, pred_bboxes) # gather pred bboxes class score