mmcv.ops.soft_nms

mmcv.ops.soft_nms(boxes: Union[torch.Tensor, numpy.ndarray], scores: Union[torch.Tensor, numpy.ndarray], iou_threshold: float = 0.3, sigma: float = 0.5, min_score: float = 0.001, method: str = 'linear', offset: int = 0) → Tuple[Union[torch.Tensor, numpy.ndarray], Union[torch.Tensor, numpy.ndarray]]

Dispatch to only CPU Soft NMS implementations.

The input can be either a torch tensor or numpy array. The returned type will always be the same as inputs.

参数

• boxes (torch.Tensor or np.ndarray) – boxes in shape (N, 4).

• scores (torch.Tensor or np.ndarray) – scores in shape (N, ).

• iou_threshold (float) – IoU threshold for NMS.

• sigma (float) – hyperparameter for gaussian method

• min_score (float) – score filter threshold

• method (str) – either ‘linear’ or ‘gaussian’

• offset (int, 0 or 1) – boxes’ width or height is (x2 - x1 + offset).

返回

kept dets (boxes and scores) and indice, which always have the same data type as the input.

返回类型

tuple

示例

>>> boxes = np.array([[4., 3., 5., 3.],
>>>                   [4., 3., 5., 4.],
>>>                   [3., 1., 3., 1.],
>>>                   [3., 1., 3., 1.],
>>>                   [3., 1., 3., 1.],
>>>                   [3., 1., 3., 1.]], dtype=np.float32)
>>> scores = np.array([0.9, 0.9, 0.5, 0.5, 0.4, 0.0], dtype=np.float32)
>>> iou_threshold = 0.6
>>> dets, inds = soft_nms(boxes, scores, iou_threshold, sigma=0.5)
>>> assert len(inds) == len(dets) == 5