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]][source]¶
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.
- Parameters
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).
- Returns
kept dets (boxes and scores) and indice, which always have the same data type as the input.
- Return type
Example
>>> 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