| |
| from typing import Union |
|
|
| import torch |
| from mmengine.structures import InstanceData |
| from numpy import ndarray |
| from torch import Tensor |
|
|
| from mmdet.registry import TASK_UTILS |
| from ..assigners import AssignResult |
| from .multi_instance_sampling_result import MultiInstanceSamplingResult |
| from .random_sampler import RandomSampler |
|
|
|
|
| @TASK_UTILS.register_module() |
| class MultiInsRandomSampler(RandomSampler): |
| """Random sampler for multi instance. |
| |
| Note: |
| Multi-instance means to predict multiple detection boxes with |
| one proposal box. `AssignResult` may assign multiple gt boxes |
| to each proposal box, in this case `RandomSampler` should be |
| replaced by `MultiInsRandomSampler` |
| """ |
|
|
| def _sample_pos(self, assign_result: AssignResult, num_expected: int, |
| **kwargs) -> Union[Tensor, ndarray]: |
| """Randomly sample some positive samples. |
| |
| Args: |
| assign_result (:obj:`AssignResult`): Bbox assigning results. |
| num_expected (int): The number of expected positive samples |
| |
| Returns: |
| Tensor or ndarray: sampled indices. |
| """ |
| pos_inds = torch.nonzero( |
| assign_result.labels[:, 0] > 0, as_tuple=False) |
| if pos_inds.numel() != 0: |
| pos_inds = pos_inds.squeeze(1) |
| if pos_inds.numel() <= num_expected: |
| return pos_inds |
| else: |
| return self.random_choice(pos_inds, num_expected) |
|
|
| def _sample_neg(self, assign_result: AssignResult, num_expected: int, |
| **kwargs) -> Union[Tensor, ndarray]: |
| """Randomly sample some negative samples. |
| |
| Args: |
| assign_result (:obj:`AssignResult`): Bbox assigning results. |
| num_expected (int): The number of expected positive samples |
| |
| Returns: |
| Tensor or ndarray: sampled indices. |
| """ |
| neg_inds = torch.nonzero( |
| assign_result.labels[:, 0] == 0, as_tuple=False) |
| if neg_inds.numel() != 0: |
| neg_inds = neg_inds.squeeze(1) |
| if len(neg_inds) <= num_expected: |
| return neg_inds |
| else: |
| return self.random_choice(neg_inds, num_expected) |
|
|
| def sample(self, assign_result: AssignResult, pred_instances: InstanceData, |
| gt_instances: InstanceData, |
| **kwargs) -> MultiInstanceSamplingResult: |
| """Sample positive and negative bboxes. |
| |
| Args: |
| assign_result (:obj:`AssignResult`): Assigning results from |
| MultiInstanceAssigner. |
| 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, ). |
| |
| Returns: |
| :obj:`MultiInstanceSamplingResult`: Sampling result. |
| """ |
|
|
| assert 'batch_gt_instances_ignore' in kwargs, \ |
| 'batch_gt_instances_ignore is necessary for MultiInsRandomSampler' |
|
|
| gt_bboxes = gt_instances.bboxes |
| ignore_bboxes = kwargs['batch_gt_instances_ignore'].bboxes |
| gt_and_ignore_bboxes = torch.cat([gt_bboxes, ignore_bboxes], dim=0) |
| priors = pred_instances.priors |
| if len(priors.shape) < 2: |
| priors = priors[None, :] |
| priors = priors[:, :4] |
|
|
| gt_flags = priors.new_zeros((priors.shape[0], ), dtype=torch.uint8) |
| priors = torch.cat([priors, gt_and_ignore_bboxes], dim=0) |
| gt_ones = priors.new_ones( |
| gt_and_ignore_bboxes.shape[0], dtype=torch.uint8) |
| gt_flags = torch.cat([gt_flags, gt_ones]) |
|
|
| num_expected_pos = int(self.num * self.pos_fraction) |
| pos_inds = self.pos_sampler._sample_pos(assign_result, |
| num_expected_pos) |
| |
| |
| pos_inds = pos_inds.unique() |
| num_sampled_pos = pos_inds.numel() |
| num_expected_neg = self.num - num_sampled_pos |
| if self.neg_pos_ub >= 0: |
| _pos = max(1, num_sampled_pos) |
| neg_upper_bound = int(self.neg_pos_ub * _pos) |
| if num_expected_neg > neg_upper_bound: |
| num_expected_neg = neg_upper_bound |
| neg_inds = self.neg_sampler._sample_neg(assign_result, |
| num_expected_neg) |
| neg_inds = neg_inds.unique() |
|
|
| sampling_result = MultiInstanceSamplingResult( |
| pos_inds=pos_inds, |
| neg_inds=neg_inds, |
| priors=priors, |
| gt_and_ignore_bboxes=gt_and_ignore_bboxes, |
| assign_result=assign_result, |
| gt_flags=gt_flags) |
| return sampling_result |
|
|
