def nms_core(dets, iou_thr, score_thr, max_num):
if is_in_onnx_export():
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 clamp(x, min, max):
if is_in_onnx_export():
is_min_tensor = isinstance(min, torch.Tensor)
is_max_tensor = isinstance(max, torch.Tensor)
if is_min_tensor and is_max_tensor:
y = x.clamp(min=min, max=max)
else:
device = x.device
dtype = x.dtype
y = x
d = len(y.shape)
min_val = torch.as_tensor(min, dtype=dtype, device=device)
y = torch.stack([y, min_val.view([
1,
] * y.dim()).expand_as(y)],
dim=d)
y = torch.max(y, dim=d, keepdim=False)[0]
max_val = torch.as_tensor(max, dtype=dtype, device=device)
y = torch.stack([y, max_val.view([
1,
] * y.dim()).expand_as(y)],
dim=d)
y = torch.min(y, dim=d, keepdim=False)[0]
else:
y = x.clamp(min=min, max=max)
return y
def topk(x, k, dim=None, **kwargs):
from torch.onnx import operators, is_in_onnx_export
if dim is None:
dim = x.dim() - 1
if is_in_onnx_export():
n = operators.shape_as_tensor(x)[dim].unsqueeze(0)
if not isinstance(k, torch.Tensor):
k = torch.tensor([k], dtype=torch.long)
# Workaround for ONNXRuntime: convert values to int to get minimum.
n = torch.min(torch.cat((k, n), dim=0).int()).long()
# ONNX OpSet 10 does not support non-floating point input for TopK.
original_dtype = x.dtype
require_cast = original_dtype not in {
torch.float16, torch.float32, torch.float64
}
if require_cast:
x = x.to(torch.float32)
values, keep = torch.topk(x, n, dim=dim, **kwargs)
if require_cast:
values = values.to(original_dtype)
else:
values, keep = torch.topk(
x, min(int(k), x.shape[dim]), dim=dim, **kwargs)
return values, keep
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 get_bboxes(self,
cls_scores,
bbox_preds,
centernesses,
img_metas,
cfg,
rescale=False):
from torch.onnx import is_in_onnx_export
if is_in_onnx_export():
from ...utils.deployment import TracerStub
assert len(cls_scores) == len(bbox_preds)
num_levels = len(cls_scores)
device = cls_scores[0].device
# FIXME. Workaround for OpenVINO-friendly export:
# passing scores tensor itself instead of its spatial size.
pass_tensor = is_in_onnx_export() and isinstance(self.anchor_generators[0].grid_anchors, TracerStub)
mlvl_anchors = [
self.anchor_generators[i].grid_anchors(
cls_scores[i] if pass_tensor else cls_scores[i].size()[-2:],
self.anchor_strides[i],
device=device) for i in range(num_levels)
]
result_list = []
for img_id in range(len(img_metas)):
cls_score_list = [
cls_scores[i][img_id].detach() for i in range(num_levels)
]
bbox_pred_list = [
bbox_preds[i][img_id].detach() for i in range(num_levels)
]
centerness_pred_list = [
centernesses[i][img_id].detach() for i in range(num_levels)
]
img_shape = img_metas[img_id]['img_shape']
scale_factor = img_metas[img_id]['scale_factor']
proposals = self.get_bboxes_single(cls_score_list, bbox_pred_list,
centerness_pred_list,
mlvl_anchors, img_shape,
scale_factor, cfg, rescale)
result_list.append(proposals)
return result_list
def get_bboxes(self,
cls_scores,
bbox_preds,
img_metas,
cfg,
rescale=False):
"""
Transform network output for a batch into labeled boxes.
Args:
cls_scores (list[Tensor]): Box scores for each scale level
Has shape (N, num_anchors * num_classes, H, W)
bbox_preds (list[Tensor]): Box energies / deltas for each scale
level with shape (N, num_anchors * 4, H, W)
img_metas (list[dict]): size / scale info for each image
cfg (mmcv.Config): test / postprocessing configuration
rescale (bool): if True, return boxes in original image space
Returns:
list[tuple[Tensor, Tensor]]: each item in result_list is 2-tuple.
The first item is an (n, 5) tensor, where the first 4 columns
are bounding box positions (tl_x, tl_y, br_x, br_y) and the
5-th column is a score between 0 and 1. The second item is a
(n,) tensor where each item is the class index of the
corresponding box.
Example:
>>> import mmcv
>>> self = AnchorHead(num_classes=9, in_channels=1)
>>> img_metas = [{'img_shape': (32, 32, 3), 'scale_factor': 1}]
>>> cfg = mmcv.Config(dict(
>>> score_thr=0.00,
>>> nms=dict(type='nms', iou_thr=1.0),
>>> max_per_img=10))
>>> feat = torch.rand(1, 1, 3, 3)
>>> cls_score, bbox_pred = self.forward_single(feat)
>>> # note the input lists are over different levels, not images
>>> cls_scores, bbox_preds = [cls_score], [bbox_pred]
>>> result_list = self.get_bboxes(cls_scores, bbox_preds,
>>> img_metas, cfg)
>>> det_bboxes, det_labels = result_list[0]
>>> assert len(result_list) == 1
>>> assert det_bboxes.shape[1] == 5
>>> assert len(det_bboxes) == len(det_labels) == cfg.max_per_img
"""
from torch.onnx import is_in_onnx_export
if is_in_onnx_export():
from ...utils.deployment import TracerStub
assert len(cls_scores) == len(bbox_preds)
num_levels = len(cls_scores)
device = cls_scores[0].device
# FIXME. Workaround for OpenVINO-friendly export:
# passing scores tensor itself instead of its spatial size.
pass_tensor = is_in_onnx_export() and isinstance(
self.anchor_generators[0].grid_anchors, TracerStub)
mlvl_anchors = [
self.anchor_generators[i].grid_anchors(
cls_scores[i] if pass_tensor else cls_scores[i].size()[-2:],
self.anchor_strides[i],
device=device) for i in range(num_levels)
]
result_list = []
for img_id in range(len(img_metas)):
cls_score_list = [
cls_scores[i][img_id].detach() for i in range(num_levels)
]
bbox_pred_list = [
bbox_preds[i][img_id].detach() for i in range(num_levels)
]
img_shape = img_metas[img_id]['img_shape']
scale_factor = img_metas[img_id]['scale_factor']
proposals = self.get_bboxes_single(cls_score_list, bbox_pred_list,
mlvl_anchors, img_shape,
scale_factor, cfg, rescale)
result_list.append(proposals)
return result_list
def _get_bboxes_single(self,
cls_scores,
bbox_preds,
mlvl_anchors,
img_shape,
scale_factor,
cfg,
rescale=False):
cfg = self.test_cfg if cfg is None else cfg
# bboxes from different level should be independent during NMS,
# level_ids are used as labels for batched NMS to separate them
level_ids = []
mlvl_scores = []
mlvl_bbox_preds = []
mlvl_valid_anchors = []
for idx in range(len(cls_scores)):
rpn_cls_score = cls_scores[idx]
rpn_bbox_pred = bbox_preds[idx]
assert rpn_cls_score.size()[-2:] == rpn_bbox_pred.size()[-2:]
rpn_cls_score = rpn_cls_score.permute(1, 2, 0)
if self.use_sigmoid_cls:
rpn_cls_score = rpn_cls_score.reshape(-1)
scores = rpn_cls_score.sigmoid()
else:
rpn_cls_score = rpn_cls_score.reshape(-1, 2)
# we set FG labels to [0, num_class-1] and BG label to
# num_class in other heads since mmdet v2.0, However we
# keep BG label as 0 and FG label as 1 in rpn head
scores = rpn_cls_score.softmax(dim=1)[:, 1]
rpn_bbox_pred = rpn_bbox_pred.permute(1, 2, 0).reshape(-1, 4)
anchors = mlvl_anchors[idx]
if cfg.nms_pre > 0:
scores, topk_inds = topk(scores, cfg.nms_pre, dim=0)
rpn_bbox_pred = rpn_bbox_pred[topk_inds]
anchors = anchors[topk_inds]
mlvl_scores.append(scores)
mlvl_bbox_preds.append(rpn_bbox_pred)
mlvl_valid_anchors.append(anchors)
level_ids.append(
torch.full((scores.size(0), ),
idx,
dtype=torch.long,
device=scores.device))
scores = torch.cat(mlvl_scores)
anchors = torch.cat(mlvl_valid_anchors)
rpn_bbox_pred = torch.cat(mlvl_bbox_preds)
proposals = self.bbox_coder.decode(anchors,
rpn_bbox_pred,
max_shape=img_shape)
ids = torch.cat(level_ids)
if cfg.min_bbox_size > 0:
w = proposals[:, 2] - proposals[:, 0]
h = proposals[:, 3] - proposals[:, 1]
valid_inds = torch.nonzero((w >= cfg.min_bbox_size)
& (h >= cfg.min_bbox_size),
as_tuple=False).squeeze()
if valid_inds.sum().item() != len(proposals) or is_in_onnx_export(
):
proposals = proposals[valid_inds, :]
scores = scores[valid_inds]
ids = ids[valid_inds]
# TODO: remove the hard coded nms type
nms_cfg = dict(type='nms', iou_thr=cfg.nms_thr)
dets, keep = batched_nms(proposals, scores, ids, nms_cfg)
return dets[:cfg.nms_post]
