| |
| from typing import Optional |
|
|
| import torch |
| from mmengine.structures import InstanceData |
|
|
| from mmdet.models.task_modules.assigners.assign_result import AssignResult |
| from mmdet.models.task_modules.assigners.max_iou_assigner import MaxIoUAssigner |
| from mmdet.registry import TASK_UTILS |
|
|
|
|
| @TASK_UTILS.register_module() |
| class TransMaxIoUAssigner(MaxIoUAssigner): |
|
|
| def assign(self, |
| pred_instances: InstanceData, |
| gt_instances: InstanceData, |
| gt_instances_ignore: Optional[InstanceData] = None, |
| **kwargs) -> AssignResult: |
| """Assign gt to bboxes. |
| |
| This method assign a gt bbox to every bbox (proposal/anchor), each bbox |
| will be assigned with -1, or a semi-positive number. -1 means negative |
| sample, semi-positive number is the index (0-based) of assigned gt. |
| The assignment is done in following steps, the order matters. |
| |
| 1. assign every bbox to the background |
| 2. assign proposals whose iou with all gts < neg_iou_thr to 0 |
| 3. for each bbox, if the iou with its nearest gt >= pos_iou_thr, |
| assign it to that bbox |
| 4. for each gt bbox, assign its nearest proposals (may be more than |
| one) to itself |
| |
| Args: |
| pred_instances (:obj:`InstanceData`): Instances of model |
| predictions. It includes ``priors``, and the priors can |
| be anchors or points, or the bboxes predicted by the |
| previous stage, has shape (n, 4). The bboxes predicted by |
| the current model or stage will be named ``bboxes``, |
| ``labels``, and ``scores``, the same as the ``InstanceData`` |
| in other places. |
| gt_instances (:obj:`InstanceData`): Ground truth of instance |
| annotations. It usually includes ``bboxes``, with shape (k, 4), |
| and ``labels``, with shape (k, ). |
| gt_instances_ignore (:obj:`InstanceData`, optional): Instances |
| to be ignored during training. It includes ``bboxes`` |
| attribute data that is ignored during training and testing. |
| Defaults to None. |
| |
| Returns: |
| :obj:`AssignResult`: The assign result. |
| |
| Example: |
| >>> from mmengine.structures import InstanceData |
| >>> self = MaxIoUAssigner(0.5, 0.5) |
| >>> pred_instances = InstanceData() |
| >>> pred_instances.priors = torch.Tensor([[0, 0, 10, 10], |
| ... [10, 10, 20, 20]]) |
| >>> gt_instances = InstanceData() |
| >>> gt_instances.bboxes = torch.Tensor([[0, 0, 10, 9]]) |
| >>> gt_instances.labels = torch.Tensor([0]) |
| >>> assign_result = self.assign(pred_instances, gt_instances) |
| >>> expected_gt_inds = torch.LongTensor([1, 0]) |
| >>> assert torch.all(assign_result.gt_inds == expected_gt_inds) |
| """ |
| gt_bboxes = gt_instances.bboxes |
| priors = pred_instances.priors |
| gt_labels = gt_instances.labels |
| if gt_instances_ignore is not None: |
| gt_bboxes_ignore = gt_instances_ignore.bboxes |
| else: |
| gt_bboxes_ignore = None |
|
|
| assign_on_cpu = True if (self.gpu_assign_thr > 0) and ( |
| gt_bboxes.shape[0] > self.gpu_assign_thr) else False |
| |
| if assign_on_cpu: |
| device = priors.device |
| priors = priors.cpu() |
| gt_bboxes = gt_bboxes.cpu() |
| gt_labels = gt_labels.cpu() |
| if gt_bboxes_ignore is not None: |
| gt_bboxes_ignore = gt_bboxes_ignore.cpu() |
|
|
| trans_priors = torch.cat([ |
| priors[..., 1].view(-1, 1), priors[..., 0].view(-1, 1), |
| priors[..., 3].view(-1, 1), priors[..., 2].view(-1, 1) |
| ], |
| dim=-1) |
| overlaps = self.iou_calculator(gt_bboxes, trans_priors) |
|
|
| if (self.ignore_iof_thr > 0 and gt_bboxes_ignore is not None |
| and gt_bboxes_ignore.numel() > 0 and trans_priors.numel() > 0): |
| if self.ignore_wrt_candidates: |
| ignore_overlaps = self.iou_calculator( |
| trans_priors, gt_bboxes_ignore, mode='iof') |
| ignore_max_overlaps, _ = ignore_overlaps.max(dim=1) |
| else: |
| ignore_overlaps = self.iou_calculator( |
| gt_bboxes_ignore, trans_priors, mode='iof') |
| ignore_max_overlaps, _ = ignore_overlaps.max(dim=0) |
| overlaps[:, ignore_max_overlaps > self.ignore_iof_thr] = -1 |
|
|
| assign_result = self.assign_wrt_overlaps(overlaps, gt_labels) |
| if assign_on_cpu: |
| assign_result.gt_inds = assign_result.gt_inds.to(device) |
| assign_result.max_overlaps = assign_result.max_overlaps.to(device) |
| if assign_result.labels is not None: |
| assign_result.labels = assign_result.labels.to(device) |
| return assign_result |
|
|
