task_aligned_assigner.py 6.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158
  1. # Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from __future__ import absolute_import
  15. from __future__ import division
  16. from __future__ import print_function
  17. import paddle
  18. import paddle.nn as nn
  19. import paddle.nn.functional as F
  20. from paddlex.ppdet.core.workspace import register
  21. from ..bbox_utils import iou_similarity
  22. from .utils import (pad_gt, gather_topk_anchors, check_points_inside_bboxes,
  23. compute_max_iou_anchor)
  24. @register
  25. class TaskAlignedAssigner(nn.Layer):
  26. """TOOD: Task-aligned One-stage Object Detection
  27. """
  28. def __init__(self, topk=13, alpha=1.0, beta=6.0, eps=1e-9):
  29. super(TaskAlignedAssigner, self).__init__()
  30. self.topk = topk
  31. self.alpha = alpha
  32. self.beta = beta
  33. self.eps = eps
  34. @paddle.no_grad()
  35. def forward(self,
  36. pred_scores,
  37. pred_bboxes,
  38. anchor_points,
  39. gt_labels,
  40. gt_bboxes,
  41. bg_index,
  42. gt_scores=None):
  43. r"""This code is based on
  44. https://github.com/fcjian/TOOD/blob/master/mmdet/core/bbox/assigners/task_aligned_assigner.py
  45. The assignment is done in following steps
  46. 1. compute alignment metric between all bbox (bbox of all pyramid levels) and gt
  47. 2. select top-k bbox as candidates for each gt
  48. 3. limit the positive sample's center in gt (because the anchor-free detector
  49. only can predict positive distance)
  50. 4. if an anchor box is assigned to multiple gts, the one with the
  51. highest iou will be selected.
  52. Args:
  53. pred_scores (Tensor, float32): predicted class probability, shape(B, L, C)
  54. pred_bboxes (Tensor, float32): predicted bounding boxes, shape(B, L, 4)
  55. anchor_points (Tensor, float32): pre-defined anchors, shape(L, 2), "cxcy" format
  56. gt_labels (Tensor|List[Tensor], int64): Label of gt_bboxes, shape(B, n, 1)
  57. gt_bboxes (Tensor|List[Tensor], float32): Ground truth bboxes, shape(B, n, 4)
  58. bg_index (int): background index
  59. gt_scores (Tensor|List[Tensor]|None, float32) Score of gt_bboxes,
  60. shape(B, n, 1), if None, then it will initialize with one_hot label
  61. Returns:
  62. assigned_labels (Tensor): (B, L)
  63. assigned_bboxes (Tensor): (B, L, 4)
  64. assigned_scores (Tensor): (B, L, C)
  65. """
  66. assert pred_scores.ndim == pred_bboxes.ndim
  67. gt_labels, gt_bboxes, pad_gt_scores, pad_gt_mask = pad_gt(
  68. gt_labels, gt_bboxes, gt_scores)
  69. assert gt_labels.ndim == gt_bboxes.ndim and \
  70. gt_bboxes.ndim == 3
  71. batch_size, num_anchors, num_classes = pred_scores.shape
  72. _, num_max_boxes, _ = gt_bboxes.shape
  73. # negative batch
  74. if num_max_boxes == 0:
  75. assigned_labels = paddle.full([batch_size, num_anchors], bg_index)
  76. assigned_bboxes = paddle.zeros([batch_size, num_anchors, 4])
  77. assigned_scores = paddle.zeros(
  78. [batch_size, num_anchors, num_classes])
  79. return assigned_labels, assigned_bboxes, assigned_scores
  80. # compute iou between gt and pred bbox, [B, n, L]
  81. ious = iou_similarity(gt_bboxes, pred_bboxes)
  82. # gather pred bboxes class score
  83. pred_scores = pred_scores.transpose([0, 2, 1])
  84. batch_ind = paddle.arange(
  85. end=batch_size, dtype=gt_labels.dtype).unsqueeze(-1)
  86. gt_labels_ind = paddle.stack(
  87. [batch_ind.tile([1, num_max_boxes]), gt_labels.squeeze(-1)],
  88. axis=-1)
  89. bbox_cls_scores = paddle.gather_nd(pred_scores, gt_labels_ind)
  90. # compute alignment metrics, [B, n, L]
  91. alignment_metrics = bbox_cls_scores.pow(self.alpha) * ious.pow(
  92. self.beta)
  93. # check the positive sample's center in gt, [B, n, L]
  94. is_in_gts = check_points_inside_bboxes(anchor_points, gt_bboxes)
  95. # select topk largest alignment metrics pred bbox as candidates
  96. # for each gt, [B, n, L]
  97. is_in_topk = gather_topk_anchors(
  98. alignment_metrics * is_in_gts,
  99. self.topk,
  100. topk_mask=pad_gt_mask.tile([1, 1, self.topk]).astype(paddle.bool))
  101. # select positive sample, [B, n, L]
  102. mask_positive = is_in_topk * is_in_gts * pad_gt_mask
  103. # if an anchor box is assigned to multiple gts,
  104. # the one with the highest iou will be selected, [B, n, L]
  105. mask_positive_sum = mask_positive.sum(axis=-2)
  106. if mask_positive_sum.max() > 1:
  107. mask_multiple_gts = (mask_positive_sum.unsqueeze(1) > 1).tile(
  108. [1, num_max_boxes, 1])
  109. is_max_iou = compute_max_iou_anchor(ious)
  110. mask_positive = paddle.where(mask_multiple_gts, is_max_iou,
  111. mask_positive)
  112. mask_positive_sum = mask_positive.sum(axis=-2)
  113. assigned_gt_index = mask_positive.argmax(axis=-2)
  114. assert mask_positive_sum.max() == 1, \
  115. ("one anchor just assign one gt, but received not equals 1. "
  116. "Received: %f" % mask_positive_sum.max().item())
  117. # assigned target
  118. assigned_gt_index = assigned_gt_index + batch_ind * num_max_boxes
  119. assigned_labels = paddle.gather(
  120. gt_labels.flatten(), assigned_gt_index.flatten(), axis=0)
  121. assigned_labels = assigned_labels.reshape([batch_size, num_anchors])
  122. assigned_labels = paddle.where(
  123. mask_positive_sum > 0, assigned_labels,
  124. paddle.full_like(assigned_labels, bg_index))
  125. assigned_bboxes = paddle.gather(
  126. gt_bboxes.reshape([-1, 4]), assigned_gt_index.flatten(), axis=0)
  127. assigned_bboxes = assigned_bboxes.reshape([batch_size, num_anchors, 4])
  128. assigned_scores = F.one_hot(assigned_labels, num_classes)
  129. # rescale alignment metrics
  130. alignment_metrics *= mask_positive
  131. max_metrics_per_instance = alignment_metrics.max(axis=-1, keepdim=True)
  132. max_ious_per_instance = (ious * mask_positive).max(axis=-1,
  133. keepdim=True)
  134. alignment_metrics = alignment_metrics / (
  135. max_metrics_per_instance + self.eps) * max_ious_per_instance
  136. alignment_metrics = alignment_metrics.max(-2).unsqueeze(-1)
  137. assigned_scores = assigned_scores * alignment_metrics
  138. return assigned_labels, assigned_bboxes, assigned_scores