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]][源代码]¶
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.
- 参数
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.
- 返回
kept dets (boxes and scores) and indice, which always have the same data type as the input.
- 返回类型
示例
>>> 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