mmcv.ops.nms¶

mmcv.ops.nms(boxes: Union[torch.Tensor, numpy.ndarray], scores: Union[torch.Tensor, numpy.ndarray], iou_threshold: float, offset: int = 0, score_threshold: float = 0, max_num: int = - 1) → Tuple[Union[torch.Tensor, numpy.ndarray], Union[torch.Tensor, numpy.ndarray]][source]¶

Dispatch to either CPU or GPU NMS implementations.

The input can be either torch tensor or numpy array. GPU NMS will be used if the input is gpu tensor, otherwise CPU NMS will be used. 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.
offset (int, 0 or 1) – boxes’ width or height is (x2 - x1 + offset).
score_threshold (float) – score threshold for NMS.
max_num (int) – maximum number of boxes after NMS.

Returns

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

Return type

tuple

Example

>>> boxes = np.array([[49.1, 32.4, 51.0, 35.9],
>>>                   [49.3, 32.9, 51.0, 35.3],
>>>                   [49.2, 31.8, 51.0, 35.4],
>>>                   [35.1, 11.5, 39.1, 15.7],
>>>                   [35.6, 11.8, 39.3, 14.2],
>>>                   [35.3, 11.5, 39.9, 14.5],
>>>                   [35.2, 11.7, 39.7, 15.7]], dtype=np.float32)
>>> scores = np.array([0.9, 0.9, 0.5, 0.5, 0.5, 0.4, 0.3],               dtype=np.float32)
>>> iou_threshold = 0.6
>>> dets, inds = nms(boxes, scores, iou_threshold)
>>> assert len(inds) == len(dets) == 3