Non-maximum suppression in PyTorch

Introduction

Non-maximum suppression helps pick the best boxes when detecting objects, so we don't get many overlapping boxes for the same object.

When you have many overlapping boxes from an object detector and want to keep only the best ones.

To reduce duplicate detections in images with multiple objects.

When you want to clean up predictions from models like YOLO or Faster R-CNN.

To improve the clarity of object detection results before showing them to users.

Syntax

PyTorch

torchvision.ops.nms(boxes, scores, iou_threshold) -> indices

boxes: Tensor of shape [N, 4] with box coordinates in (x1, y1, x2, y2) format.

scores: Tensor of shape [N] with confidence scores for each box.

iou_threshold: float between 0 and 1 to decide when boxes overlap too much.

Examples

Keeps boxes with less than 0.5 overlap and removes others.

PyTorch

indices = torchvision.ops.nms(boxes, scores, 0.5)

Uses a stricter overlap threshold, removing more boxes.

PyTorch

indices = torchvision.ops.nms(boxes, scores, 0.3)

Allows more overlap, keeping more boxes.

PyTorch

indices = torchvision.ops.nms(boxes, scores, 0.7)

Sample Model

This code creates 4 boxes with scores. It uses NMS to keep only boxes that don't overlap too much. Boxes 0 and 2 are kept because boxes 1 and 3 overlap too much with them and have lower scores.

PyTorch

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()}")

Important Notes

NMS works by sorting boxes by score and removing boxes that overlap too much with higher scored boxes.

The IoU threshold controls how much overlap is allowed; lower means stricter filtering.

Make sure boxes are in (x1, y1, x2, y2) format before using NMS.

Summary

Non-maximum suppression removes overlapping boxes to keep only the best detections.

It uses box coordinates and scores to decide which boxes to keep.

The IoU threshold controls how strict the overlap filtering is.