| import torch | |
| from torchvision.ops.boxes import box_iou | |
| def nms(bboxes: torch.Tensor, scores: torch.Tensor, iou_threshold: float) -> torch.Tensor: | |
| order = torch.argsort(-scores).to(bboxes.device) | |
| indices = torch.arange(bboxes.shape[0]).to(bboxes.device) | |
| keep = torch.ones_like(indices, dtype=torch.bool).to(bboxes.device) | |
| for i in indices: | |
| if keep[i]: | |
| bbox = bboxes[order[i]] | |
| iou = box_iou(bbox[None, ...], (bboxes[order[i + 1:]]) * keep[i + 1:][..., None]) | |
| overlapped = torch.nonzero(iou > iou_threshold) | |
| keep[overlapped + i + 1] = 0 | |
| return order[keep] | |