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]