Selecting the Optimal Bounding Box with Non-Max Suppression: A Comprehensive Guide

ByJordan Brown

Introduction

Object detection is a pivotal task in computer vision with wide-ranging applications from autonomous driving to video surveillance. The goal is to localize and classify objects of interest within an image, typically by drawing bounding boxes around them. Modern object detectors based on deep convolutional neural networks (CNNs) have achieved remarkable accuracy, with state-of-the-art models like YOLOv4 and EfficientDet pushing mean Average Precision (mAP) above 50% on challenging benchmarks like COCO.

However, the raw output of these detectors often contains multiple overlapping bounding boxes for each object, as shown in Figure 1. These redundant detections can significantly hinder the detector‘s precision and computational efficiency. This is where non-maximum suppression (NMS) comes in – a simple yet crucial post-processing step that filters out duplicate detections and leaves only the most confident box for each object.

Detections before and after NMS
Figure 1: Object detector output before and after applying NMS. Source: Author.

In this article, we‘ll dive deep into the technical details of NMS and its central role in the object detection pipeline. We‘ll discuss its strengths, weaknesses, and variants, and walk through a concrete implementation in PyTorch. By the end, you‘ll have a solid grasp of this core technique and how to leverage it to build high-quality detection systems.

Object Detection Pipeline Overview

Before we examine NMS in depth, let‘s briefly review the typical CNN-based object detection pipeline, illustrated in Figure 2. The main stages are:

  1. Backbone network: A pre-trained CNN like ResNet extracts convolutional features from the input image. These features encode high-level semantic information about the image content.

  2. Detection head: The backbone features are fed to a detection head that makes two types of predictions: (1) bounding box coordinates for potential objects, and (2) class probabilities indicating the object category. The head is usually a set of convolutional and fully-connected layers.

  3. Anchor boxes: Most detectors use pre-defined anchor boxes of various scales and aspect ratios tiled regularly across the image. The model predicts offsets to these anchor boxes. This allows detecting objects of different sizes and shapes.

  4. Non-max suppression: The raw detector output contains many overlapping boxes for each object. NMS filters these boxes based on their confidence scores and overlap, returning a single box per object.

Object detection pipeline
Figure 2: High-level object detection pipeline. Source: Author, inspired by [Zou et al. 2019](https://arxiv.org/abs/1809.02165).

The final output is a set of bounding boxes with class labels, ready for downstream tasks like tracking or counting. Stages 1-3 are the core components of the detector itself, while NMS is a universal post-processing step applied to the model predictions. Now let‘s zoom in on the details of the NMS algorithm.

Non-Max Suppression Algorithm

NMS is a greedy algorithm that selects high-scoring boxes and suppresses overlapping lower-scoring ones. The key steps are:

  1. Sort the predicted boxes in descending order of their confidence scores.
  2. Take the box with the highest score and add it to the final output.
  3. Compute the overlap (Intersection over Union) between this box and all remaining boxes.
  4. Remove any boxes with IoU above a set threshold (e.g., 0.5).
  5. Repeat steps 2-4 with the next highest-scoring box until all boxes are either selected or suppressed.

Mathematically, let $B = {b_1, b_2, …, b_N}$ be the set of predicted boxes and $S = {s_1, s_2, …, s_N}$ be their corresponding confidence scores. The pseudocode for NMS is:

function NMS(B, S, iou_threshold):
    final_boxes = []
    while B is not empty:
        max_idx = argmax(S)
        max_box = B[max_idx]
        final_boxes.append(max_box)

        # Remove max_box from B and S
        B.remove(max_box)
        S.remove(S[max_idx])

        # Compute IoU between max_box and remaining boxes
        for box in B:
            iou = compute_iou(max_box, box)
            if iou > iou_threshold:
                B.remove(box)
                S.remove(corresponding score)

    return final_boxes

The iou_threshold is a key hyperparameter that controls the suppression strength. Lower values lead to more aggressive suppression, while higher values allow more overlapping boxes to survive. In practice, a threshold of 0.5-0.7 works well for most detectors.

The main computational bottleneck in NMS is computing the IoUs between each selected box and all remaining boxes. This has a worst-case complexity of $O(N^2)$ for $N$ input boxes. However, in practice, most boxes are suppressed quickly and the quadratic behavior is not a major issue. Nonetheless, several more efficient variants of NMS have been proposed, which we‘ll discuss later.

Table 1 shows the impact of NMS on detector performance for various IoU thresholds. Without NMS, the detector suffers from very low precision due to many false positive boxes. NMS dramatically improves precision by removing duplicate detections, with a peak mAP around 0.5-0.6 IoU threshold. This demonstrates the critical importance of NMS for obtaining usable detector outputs.

Table of mAP scores with and without NMS
Table 1: Impact of NMS on YOLOv3 performance on COCO val set. Source: Author.

Next, we‘ll see how to implement NMS in PyTorch and visualize its effect on real detector outputs.

Implementing NMS in PyTorch

PyTorch provides a convenient torchvision.ops.nms function that performs NMS on a CPU or GPU tensor of bounding boxes. Here‘s a minimal example:

import torch
from torchvision import ops

def nms(boxes, scores, iou_threshold):
    """
    Performs NMS on a set of bounding boxes and returns the selected indices.

    Args:
        boxes (Tensor[N, 4]): boxes in (x1, y1, x2, y2) format
        scores (Tensor[N]): box confidence scores
        iou_threshold (float): IoU threshold for suppression
    Returns:
        keep (Tensor): indices of boxes to keep after NMS
    """
    return ops.nms(boxes, scores, iou_threshold)

To visualize the effect of NMS, we can plot the raw detector outputs and the post-NMS outputs on an example image. Figure 3 shows the results for YOLOv3 on a sample COCO image.

YOLOv3 outputs before and after NMS
Figure 3: YOLOv3 detections on a COCO image before and after NMS. Source: Author.

Before NMS, there are many highly overlapping boxes for each person and surfboard. NMS removes these near-duplicate boxes and returns a clean final output with one box per distinct object. The choice of IoU threshold allows trading off recall and precision – a lower threshold gives fewer but more confident detections.

Advanced NMS Variants

While standard NMS is effective and widely used, it has some limitations. In particular, it can sometimes suppress true positive boxes when two objects have very high overlap (e.g., pedestrians in a crowd). It also tends to favor larger, more confident boxes over smaller ones.

Several advanced NMS variants have been proposed to address these issues:

  • Soft-NMS (Bodla et al., 2017): Instead of completely removing boxes above the IoU threshold, Soft-NMS decays their scores according to a continuous function of their overlap. This allows retaining multiple boxes for highly overlapping objects.

  • DIoU-NMS (Zheng et al., 2020): This variant uses a distance-IoU metric that considers both the overlap and the distance between box centers. DIoU-NMS can more effectively suppress distant false positives.

  • Weighted NMS (Zhou et al., 2017): Weighted NMS assigns different suppression thresholds to different object categories based on their typical sizes. This helps retain smaller objects that would otherwise be suppressed.

Here‘s a PyTorch implementation of Soft-NMS:

def soft_nms(boxes, scores, sigma=0.5, score_thresh=0.001):
    """
    Performs Soft-NMS on a set of bounding boxes.

    Args:
        boxes (Tensor[N, 4]): boxes in (x1, y1, x2, y2) format
        scores (Tensor[N]): box scores
        sigma (float): standard deviation for Gaussian decay function
        score_thresh (float): score threshold for stopping decay
    Returns:
        keep (Tensor): indices of boxes to keep after Soft-NMS
    """
    keep = []
    idxs = scores.argsort()
    while idxs.numel() > 0:
        max_idx = idxs[-1]
        keep.append(max_idx)
        if idxs.size(0) == 1:
            break
        idxs = idxs[:-1]
        other_boxes = boxes[idxs]
        ious = ops.box_iou(boxes[max_idx].unsqueeze(0), other_boxes)[0]
        # Decay scores of overlapping boxes
        decay = torch.exp(-(ious ** 2) / sigma)
        scores[idxs] *= decay
        # Remove boxes with decayed score below threshold
        keep_idxs = scores[idxs] > score_thresh
        idxs = idxs[keep_idxs]
    return torch.tensor(keep)

Soft-NMS can improve mAP by 1-2% over standard NMS on datasets like COCO. However, it does increase the number of output boxes and hence the downstream processing cost. The choice of NMS variant depends on the specific application requirements.

Conclusion

In this article, we took a deep dive into non-maximum suppression and its central role in modern object detection pipelines. We saw how NMS filters raw detector outputs to produce a clean set of final bounding boxes, greatly improving precision. We discussed the core NMS algorithm, its efficient implementation in PyTorch, and several advanced variants that address its limitations.

To summarize the key takeaways:

  • NMS is a simple yet crucial post-processing step in nearly all CNN-based object detectors. It removes redundant detections and dramatically improves mAP.
  • The core idea is to select high-scoring boxes and suppress lower-scoring ones with high overlap. The IoU threshold controls the suppression strength.
  • Standard NMS can struggle with highly overlapping objects and heavily favor larger boxes. Advanced variants like Soft-NMS and DIoU-NMS aim to mitigate these issues.
  • Implementing NMS is straightforward in libraries like PyTorch. However, more efficient algorithms are needed for real-time applications.

As object detection technology continues to advance, NMS remains an indispensable tool for refining detector outputs. Novel NMS variants and learning-based alternatives are active research areas. Nonetheless, standard NMS is still a reliable baseline and a good starting point for any detection pipeline.

I hope this guide provided a comprehensive overview of NMS and its importance in object detection. Feel free to experiment with the provided code and try out different NMS variants on your own detection tasks. For further reading, I recommend the following influential papers:

Object detection is a vast and actively evolving field. NMS is just one piece of the puzzle, but a critical one for obtaining high-quality results. By mastering this core technique, you‘ll be well-equipped to tackle a wide range of exciting detection problems.

How useful was this post?

Click on a star to rate it!

Average rating 0 / 5. Vote count: 0

No votes so far! Be the first to rate this post.

You May Like to Read,

Similar Posts