Introduction
Non-maximum suppression helps pick the best boxes when detecting objects, so we don't get many overlapping boxes for the same object.
torchvision.ops.nms(boxes, scores, iou_threshold) -> indices
indices = torchvision.ops.nms(boxes, scores, 0.5)indices = torchvision.ops.nms(boxes, scores, 0.3)indices = torchvision.ops.nms(boxes, scores, 0.7)import torch import torchvision.ops # Define 4 boxes (x1, y1, x2, y2) boxes = torch.tensor([ [10, 10, 50, 50], [12, 12, 48, 48], [100, 100, 150, 150], [110, 110, 140, 140] ], dtype=torch.float32) # Confidence scores for each box scores = torch.tensor([0.9, 0.8, 0.75, 0.6]) # Apply Non-maximum suppression with IoU threshold 0.5 keep_indices = torchvision.ops.nms(boxes, scores, 0.5) print("Boxes kept after NMS:") for i in keep_indices: print(f"Box {i.item()}: {boxes[i].tolist()}, score: {scores[i].item()}")
import torch boxes = torch.tensor([[10, 10, 20, 20], [11, 11, 21, 21], [30, 30, 40, 40]], dtype=torch.float) scores = torch.tensor([0.9, 0.8, 0.7]) iou_threshold = 0.5 indices = torch.ops.torchvision.nms(boxes, scores, iou_threshold) print(indices.tolist())
import torch boxes = torch.tensor([[0, 0, 10, 10], [1, 1, 11, 11]], dtype=torch.float) scores = torch.tensor([0.5, 0.6]) iou_threshold = 0.3 indices = torch.nms(boxes, scores, iou_threshold) print(indices)
Pass the test and this pattern moves unknown