← back

Non-Maximum Suppression for Object Detection

#242 · Computer Vision · Hard

⊣ Solve on deep-ml.com

Problem

Implement Non-Maximum Suppression (NMS) for object detection. Given a set of bounding boxes with confidence scores, remove overlapping boxes that likely refer to the same object.

Solution

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
def nms(
    boxes: list[list[float]],
    scores: list[float],
    iou_threshold: float = 0.5,
) -> list[int]:
    """
    boxes: [[x1, y1, x2, y2], ...] bounding box coordinates
    scores: confidence score for each box
    Returns: indices of boxes to keep
    """
    if not boxes:
        return []

    n = len(boxes)
    indices = sorted(range(n), key=lambda i: scores[i], reverse=True)

    def compute_iou(i, j):
        x1 = max(boxes[i][0], boxes[j][0])
        y1 = max(boxes[i][1], boxes[j][1])
        x2 = min(boxes[i][2], boxes[j][2])
        y2 = min(boxes[i][3], boxes[j][3])

        inter_w = max(0, x2 - x1)
        inter_h = max(0, y2 - y1)
        intersection = inter_w * inter_h

        area_i = (boxes[i][2] - boxes[i][0]) * (boxes[i][3] - boxes[i][1])
        area_j = (boxes[j][2] - boxes[j][0]) * (boxes[j][3] - boxes[j][1])
        union = area_i + area_j - intersection

        if union <= 0:
            return 0.0
        return intersection / union

    keep = []
    suppressed = set()

    for idx in indices:
        if idx in suppressed:
            continue
        keep.append(idx)
        for other in indices:
            if other in suppressed or other == idx:
                continue
            if compute_iou(idx, other) > iou_threshold:
                suppressed.add(other)

    return keep

Explanation

  1. Sort boxes by confidence score in descending order.
  2. Pick the highest-scoring box and add it to the keep list.
  3. Compute the IoU (Intersection over Union) between the selected box and all remaining boxes.
  4. Suppress (remove) any box with IoU above the threshold, as it likely detects the same object.
  5. Repeat until all boxes are either kept or suppressed.

Complexity

  • Time: O(n^2) where n is the number of boxes (pairwise IoU comparisons)
  • Space: O(n) for the suppressed set and keep list
← #241#243 →