r"""Performs non-maximum suppression (NMS) on the boxes according to their intersection-over-union(IoU). Args: boxes: tensor of shape ``(N, 4)``; the boxes to perform nms on; each box is expected to be in ``(x1, y1, x2, y2)`` format. iou_thresh: IoU threshold for overlapping.
(
boxes: Tensor, scores: Tensor, iou_thresh: float, max_output: Optional[int] = None
)
| 245 | |
| 246 | |
| 247 | def nms( |
| 248 | boxes: Tensor, scores: Tensor, iou_thresh: float, max_output: Optional[int] = None |
| 249 | ) -> Tensor: |
| 250 | r"""Performs non-maximum suppression (NMS) on the boxes according to their intersection-over-union(IoU). |
| 251 | |
| 252 | Args: |
| 253 | boxes: tensor of shape ``(N, 4)``; the boxes to perform nms on; each box is expected to be in ``(x1, y1, x2, y2)`` format. |
| 254 | iou_thresh: IoU threshold for overlapping. |
| 255 | scores: tensor of shape ``(N,)``, the score of boxes. |
| 256 | max_output: the maximum number of boxes to keep; it is optional if this operator is not traced |
| 257 | otherwise it required to be specified; if it is not specified, all boxes are kept. |
| 258 | |
| 259 | Returns: |
| 260 | indices of the elements that have been kept by NMS, sorted by scores. |
| 261 | |
| 262 | Note: |
| 263 | max_output should be specified and should have valid positive value under tracing. |
| 264 | |
| 265 | Examples: |
| 266 | >>> import numpy as np |
| 267 | >>> x = np.zeros((100,4)) |
| 268 | >>> np.random.seed(42) |
| 269 | >>> x[:,:2] = np.random.rand(100,2)*20 |
| 270 | >>> x[:,2:] = np.random.rand(100,2)*20 + 100 |
| 271 | >>> scores = Tensor(np.random.rand(100)) |
| 272 | >>> inp = Tensor(x) |
| 273 | >>> F.vision.nms(inp, scores, iou_thresh=0.7) |
| 274 | Tensor([75 69], dtype=int32, device=xpux:0) |
| 275 | """ |
| 276 | assert ( |
| 277 | boxes.ndim == 2 and boxes.shape[1] == 4 |
| 278 | ), "the expected shape of boxes is (N, 4)" |
| 279 | assert scores.ndim == 1, "the expected shape of scores is (N,)" |
| 280 | assert ( |
| 281 | boxes.shape[0] == scores.shape[0] |
| 282 | ), "number of boxes and scores are not matched" |
| 283 | |
| 284 | boxes = boxes.detach() |
| 285 | scores = scores.detach() |
| 286 | sorted_idx = argsort(scores, descending=True) |
| 287 | boxes = boxes[sorted_idx] |
| 288 | |
| 289 | if max_output is None: |
| 290 | max_output = boxes.shape[0] |
| 291 | |
| 292 | op = builtin.NMSKeep(iou_thresh, max_output) |
| 293 | inp = (boxes.reshape(1, -1, 4),) |
| 294 | indices, count = apply(op, *inp) |
| 295 | indices = indices[0][: count[0]] |
| 296 | keep_inds = sorted_idx[indices] |
| 297 | return keep_inds |
| 298 | |
| 299 | |
| 300 | def remap( |