Non-Maximum Suppression (NMS) - CNNs & Computer Vision

Non-Maximum Suppression (NMS)

What you'll build

Greedy Non-Maximum Suppression: from a set of scored boxes, keep the best box for each object and drop the overlapping duplicates.

Background

Detectors predict a box at every anchor or grid cell, so one object usually gets many near-identical boxes. NMS is the post-processing step (in Faster R-CNN, SSD, RetinaNet and YOLO up to v8) that keeps the highest-scoring box and removes boxes that overlap it too much.

Algorithm

1. Sort the box indices by score, highest first.
2. Take the first remaining index i and append it to keep.
3. Remove every remaining box whose IoU with box i is greater than iou_threshold.
4. Repeat steps 2 and 3 until no boxes remain, then return keep.

IoU uses the xyxy convention of the IoU problem: intersection / (area_a + area_b - intersection), with the intersection width and height clamped at 0.

Inputs

• boxes: array or list of shape (N, 4), each row [x1, y1, x2, y2] (top-left, bottom-right).
• scores: array or list of shape (N,), one confidence per box.
• iou_threshold: float in [0, 1], default 0.5.

Expected return format

A Python list of int indices into the original boxes, in the order they were kept (decreasing score).

Edge cases and conventions

• A box with IoU exactly equal to the threshold is kept; only IoU > iou_threshold suppresses.
• Boxes are compared with the box just kept, not only with the overall top box.
• Empty input returns [].
• Scores are distinct in the tests; with ties, a stable sort keeps the lower index first.
• This is class-agnostic NMS: all boxes compete with each other.

Constraints

NumPy or plain Python; no torchvision.ops.nms. Vectorising the IoU of the kept box against all remaining boxes is a nice touch.

Follow-up

• What is the worst-case time complexity, and how do you make NMS class-aware with one call (the coordinate-offset trick)?
• How does Soft-NMS differ (decay the scores instead of deleting boxes), and when does it help?

Input: boxes = [[0, 0, 4, 4], [1, 1, 4, 4], [6, 6, 8, 8]], scores = [0.9, 0.75, 0.8], iou_threshold = 0.5
Output: [0, 2]

Input: boxes = [[0, 0, 1, 1], [3, 3, 4, 4], [6, 6, 7, 7]], scores = [0.5, 0.9, 0.7]
Output: [1, 2, 0]

import numpy as np

def nms(boxes, scores, iou_threshold=0.5):
    """
    Greedy Non-Maximum Suppression (class-agnostic).

    Args:
        boxes: (N, 4) boxes in xyxy format [x1, y1, x2, y2]
        scores: (N,) confidence score per box
        iou_threshold: suppress boxes with IoU > iou_threshold (default 0.5)

    Returns:
        List of kept indices (ints), in decreasing score order
    """
    # Your code here
    pass