示例#1
0
    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)
示例#2
0
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
示例#3
0
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
示例#4
0
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
示例#5
0
 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