def forward_test(self, imgs, img_metas, **kwargs):
"""
Args:
imgs (List[Tensor]): the outer list indicates test-time
augmentations and inner Tensor should have a shape NxCxHxW,
which contains all images in the batch.
img_metas (List[List[dict]]): the outer list indicates test-time
augs (multiscale, flip, etc.) and the inner list indicates
images in a batch.
Note that if `kwargs` contains either `forward_export=True` or
`dummy_forward=True` parameters, one of special branches of code is
enabled for ONNX export
(see the methods `forward_export` and `forward_dummy`).
"""
if kwargs.get('forward_export'):
logger = get_root_logger()
logger.info('Calling forward_export inside forward_test')
return self.forward_export(imgs)
if kwargs.get('dummy_forward'):
return self.forward_dummy(imgs[0])
for var, name in [(imgs, 'imgs'), (img_metas, 'img_metas')]:
if not isinstance(var, list):
raise TypeError(f'{name} must be a list, but got {type(var)}')
num_augs = len(imgs)
if num_augs != len(img_metas):
raise ValueError(f'num of augmentations ({len(imgs)}) '
f'!= num of image meta ({len(img_metas)})')
# TODO: remove the restriction of samples_per_gpu == 1 when prepared
samples_per_gpu = imgs[0].size(0)
assert samples_per_gpu == 1
if num_augs == 1:
"""
proposals (List[List[Tensor]]): the outer list indicates test-time
augs (multiscale, flip, etc.) and the inner list indicates
images in a batch. The Tensor should have a shape Px4, where
P is the number of proposals.
"""
if 'proposals' in kwargs:
kwargs['proposals'] = kwargs['proposals'][0]
with no_nncf_trace():
return self.simple_test(imgs[0], img_metas[0], **kwargs)
else:
# TODO: support test augmentation for predefined proposals
assert 'proposals' not in kwargs
return self.aug_test(imgs, img_metas, **kwargs)
def batched_nms(bboxes, scores, inds, nms_cfg, class_agnostic=False):
"""Performs non-maximum suppression in a batched fashion.
Modified from https://github.com/pytorch/vision/blob
/505cd6957711af790211896d32b40291bea1bc21/torchvision/ops/boxes.py#L39.
In order to perform NMS independently per class, we add an offset to all
the boxes. The offset is dependent only on the class idx, and is large
enough so that boxes from different classes do not overlap.
Arguments:
bboxes (torch.Tensor): bboxes in shape (N, 4).
scores (torch.Tensor): scores in shape (N, ).
inds (torch.Tensor): each index value correspond to a bbox cluster,
and NMS will not be applied between elements of different inds,
shape (N, ).
nms_cfg (dict): specify nms type and class_agnostic as well as other
parameters like iou_thr.
class_agnostic (bool): if true, nms is class agnostic,
i.e. IoU thresholding happens over all bboxes,
regardless of the predicted class
Returns:
tuple: kept bboxes and indice.
"""
nms_cfg_ = nms_cfg.copy()
class_agnostic = nms_cfg_.pop('class_agnostic', class_agnostic)
if class_agnostic:
bboxes_for_nms = bboxes
else:
max_coordinate = bboxes.max()
offsets = inds.to(bboxes) * (max_coordinate + 1)
with no_nncf_trace():
# NB: this trick is required to make class-separate NMS using ONNX NMS operation;
# the ONNX NMS operation supports another way of class separation (class-separate scores), but this is not used here.
# Note that `not_nncf_trace` is required here, since this trick causes accuracy degradation in case of int8 quantization:
# if the output of the addition below is quantized, the maximal output value is about
# ~ max_value_in_inds * max_coordinate,
# usually this value is big, so after int8-quantization different small bounding
# boxes may be squashed into the same bounding box, this may cause accuracy degradation.
# TODO: check if it is possible in this architecture use class-separate scores that are supported in ONNX NMS.
bboxes_for_nms = bboxes + offsets[:, None]
nms_type = nms_cfg_.pop('type', 'nms')
nms_op = eval(nms_type)
dets, keep = nms_op(torch.cat([bboxes_for_nms, scores[:, None]], -1),
**nms_cfg_)
bboxes = bboxes[keep]
scores = dets[:, -1]
return torch.cat([bboxes, scores[:, None]], -1), keep
def nms_core(dets, iou_thr, score_thr, max_num):
if is_in_onnx_export():
with no_nncf_trace():
valid_dets_mask = dets[:, 4] > score_thr
dets = dets[valid_dets_mask]
if dets.shape[0] == 0:
inds = dets.new_zeros(0, dtype=torch.long)
else:
inds = nms_ext.nms(dets, iou_thr)
if is_in_onnx_export():
inds = inds[:max_num]
return inds
def multiclass_nms_core(multi_bboxes,
multi_scores,
score_thr,
nms_cfg,
max_num=-1):
num_classes = multi_scores.size(1)
if multi_bboxes.shape[1] > 4:
bboxes = multi_bboxes.view(multi_scores.size(0), -1, 4)
else:
bboxes = multi_bboxes[:, None].expand(multi_scores.size(0),
num_classes, 4)
scores = multi_scores
if is_in_onnx_export():
labels = torch.arange(num_classes, dtype=torch.long, device=scores.device) \
.unsqueeze(0) \
.expand_as(scores) \
.reshape(-1)
bboxes = bboxes.reshape(-1, 4)
scores = scores.reshape(-1)
assert nms_cfg[
'type'] == 'nms', 'Only vanilla NMS is compatible with ONNX export'
nms_cfg['score_thr'] = score_thr
nms_cfg['max_num'] = max_num if max_num > 0 else sys.maxsize
else:
with no_nncf_trace():
valid_mask = scores > score_thr
bboxes = bboxes[valid_mask]
scores = scores[valid_mask]
labels = valid_mask.nonzero()[:, 1]
if bboxes.numel() == 0:
dets = multi_bboxes.new_zeros((0, 6))
return dets
dets, keep = batched_nms(bboxes, scores, labels, nms_cfg)
labels = labels[keep]
dets = torch.cat([dets, labels.to(dets.dtype).unsqueeze(-1)], dim=1)
if not is_in_onnx_export() and max_num > 0:
dets = dets[:max_num]
return dets
def forward_train(self,
x,
img_metas,
gt_bboxes,
gt_labels=None,
gt_bboxes_ignore=None,
proposal_cfg=None,
**kwargs):
outs = self(x)
if gt_labels is None:
loss_inputs = outs + (gt_bboxes, img_metas)
else:
loss_inputs = outs + (gt_bboxes, gt_labels, img_metas)
with no_nncf_trace():
losses = self.loss(*loss_inputs, gt_bboxes_ignore=gt_bboxes_ignore)
if proposal_cfg is None:
return losses
else:
proposal_list = self.get_bboxes(*outs, img_metas, cfg=proposal_cfg)
return losses, proposal_list