Spaces:
Runtime error
Runtime error
File size: 6,188 Bytes
f549064 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 |
# Copyright (c) OpenMMLab. All rights reserved.
from typing import Optional
import torch
from mmengine.structures import InstanceData
from mmdet.registry import TASK_UTILS
from .assign_result import AssignResult
from .max_iou_assigner import MaxIoUAssigner
@TASK_UTILS.register_module()
class MultiInstanceAssigner(MaxIoUAssigner):
"""Assign a corresponding gt bbox or background to each proposal bbox. If
we need to use a proposal box to generate multiple predict boxes,
`MultiInstanceAssigner` can assign multiple gt to each proposal box.
Args:
num_instance (int): How many bboxes are predicted by each proposal box.
"""
def __init__(self, num_instance: int = 2, **kwargs):
super().__init__(**kwargs)
self.num_instance = num_instance
def assign(self,
pred_instances: InstanceData,
gt_instances: InstanceData,
gt_instances_ignore: Optional[InstanceData] = None,
**kwargs) -> AssignResult:
"""Assign gt to bboxes.
This method assign gt bboxes to every bbox (proposal/anchor), each bbox
is assigned a set of gts, and the number of gts in this set is defined
by `self.num_instance`.
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.
"""
gt_bboxes = gt_instances.bboxes
priors = pred_instances.priors
# Set the FG label to 1 and add ignored annotations
gt_labels = gt_instances.labels + 1
if gt_instances_ignore is not None:
gt_bboxes_ignore = gt_instances_ignore.bboxes
if hasattr(gt_instances_ignore, 'labels'):
gt_labels_ignore = gt_instances_ignore.labels
else:
gt_labels_ignore = torch.ones_like(gt_bboxes_ignore)[:, 0] * -1
else:
gt_bboxes_ignore = None
gt_labels_ignore = None
assign_on_cpu = True if (self.gpu_assign_thr > 0) and (
gt_bboxes.shape[0] > self.gpu_assign_thr) else False
# compute overlap and assign gt on CPU when number of GT is large
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()
gt_labels_ignore = gt_labels_ignore.cpu()
if gt_bboxes_ignore is not None:
all_bboxes = torch.cat([gt_bboxes, gt_bboxes_ignore], dim=0)
all_labels = torch.cat([gt_labels, gt_labels_ignore], dim=0)
else:
all_bboxes = gt_bboxes
all_labels = gt_labels
all_priors = torch.cat([priors, all_bboxes], dim=0)
overlaps_normal = self.iou_calculator(
all_priors, all_bboxes, mode='iou')
overlaps_ignore = self.iou_calculator(
all_priors, all_bboxes, mode='iof')
gt_ignore_mask = all_labels.eq(-1).repeat(all_priors.shape[0], 1)
overlaps_normal = overlaps_normal * ~gt_ignore_mask
overlaps_ignore = overlaps_ignore * gt_ignore_mask
overlaps_normal, overlaps_normal_indices = overlaps_normal.sort(
descending=True, dim=1)
overlaps_ignore, overlaps_ignore_indices = overlaps_ignore.sort(
descending=True, dim=1)
# select the roi with the higher score
max_overlaps_normal = overlaps_normal[:, :self.num_instance].flatten()
gt_assignment_normal = overlaps_normal_indices[:, :self.
num_instance].flatten()
max_overlaps_ignore = overlaps_ignore[:, :self.num_instance].flatten()
gt_assignment_ignore = overlaps_ignore_indices[:, :self.
num_instance].flatten()
# ignore or not
ignore_assign_mask = (max_overlaps_normal < self.pos_iou_thr) * (
max_overlaps_ignore > max_overlaps_normal)
overlaps = (max_overlaps_normal * ~ignore_assign_mask) + (
max_overlaps_ignore * ignore_assign_mask)
gt_assignment = (gt_assignment_normal * ~ignore_assign_mask) + (
gt_assignment_ignore * ignore_assign_mask)
assigned_labels = all_labels[gt_assignment]
fg_mask = (overlaps >= self.pos_iou_thr) * (assigned_labels != -1)
bg_mask = (overlaps < self.neg_iou_thr) * (overlaps >= 0)
assigned_labels[fg_mask] = 1
assigned_labels[bg_mask] = 0
overlaps = overlaps.reshape(-1, self.num_instance)
gt_assignment = gt_assignment.reshape(-1, self.num_instance)
assigned_labels = assigned_labels.reshape(-1, self.num_instance)
assign_result = AssignResult(
num_gts=all_bboxes.size(0),
gt_inds=gt_assignment,
max_overlaps=overlaps,
labels=assigned_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
|