|
|
|
import numpy as np |
|
import torch |
|
|
|
from mmdet.registry import TASK_UTILS |
|
from .random_sampler import RandomSampler |
|
|
|
|
|
@TASK_UTILS.register_module() |
|
class InstanceBalancedPosSampler(RandomSampler): |
|
"""Instance balanced sampler that samples equal number of positive samples |
|
for each instance.""" |
|
|
|
def _sample_pos(self, assign_result, num_expected, **kwargs): |
|
"""Sample positive boxes. |
|
|
|
Args: |
|
assign_result (:obj:`AssignResult`): The assigned results of boxes. |
|
num_expected (int): The number of expected positive samples |
|
|
|
Returns: |
|
Tensor or ndarray: sampled indices. |
|
""" |
|
pos_inds = torch.nonzero(assign_result.gt_inds > 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: |
|
unique_gt_inds = assign_result.gt_inds[pos_inds].unique() |
|
num_gts = len(unique_gt_inds) |
|
num_per_gt = int(round(num_expected / float(num_gts)) + 1) |
|
sampled_inds = [] |
|
for i in unique_gt_inds: |
|
inds = torch.nonzero( |
|
assign_result.gt_inds == i.item(), as_tuple=False) |
|
if inds.numel() != 0: |
|
inds = inds.squeeze(1) |
|
else: |
|
continue |
|
if len(inds) > num_per_gt: |
|
inds = self.random_choice(inds, num_per_gt) |
|
sampled_inds.append(inds) |
|
sampled_inds = torch.cat(sampled_inds) |
|
if len(sampled_inds) < num_expected: |
|
num_extra = num_expected - len(sampled_inds) |
|
extra_inds = np.array( |
|
list(set(pos_inds.cpu()) - set(sampled_inds.cpu()))) |
|
if len(extra_inds) > num_extra: |
|
extra_inds = self.random_choice(extra_inds, num_extra) |
|
extra_inds = torch.from_numpy(extra_inds).to( |
|
assign_result.gt_inds.device).long() |
|
sampled_inds = torch.cat([sampled_inds, extra_inds]) |
|
elif len(sampled_inds) > num_expected: |
|
sampled_inds = self.random_choice(sampled_inds, num_expected) |
|
return sampled_inds |
|
|
