How to apply Non Maximum Suppression to YOLOv5 detections

Learn how to apply NMS to computer vision model detections.

Overview

Non-Maximum Suppression (NMS) allows you to remove duplicate, overlapping bounding boxes from predictions returned by a computer vision model.

In this guide, we will show you how to apply NMS to

YOLOv5

predictions in a few lines of code.

We will:

1. Install supervision
2. Load data
3. Apply NMS with the .with_nms() method

Without further ado, let's get started!

YOLOv5 and Image Annotation Resources

Explore these resources to enhance your understanding of YOLOv5 and image annotation techniques.

Import data into Roboflow Annotate

First, install the supervision pip package:

pip install supervision


Once you have installed supervision, you are ready to load your data and start writing logic to filter detections.

Load Data and Apply NMS

First, we are going to load our predictions into a supervision Detections() object. You can apply NMS to the detections when loaded into the object. To load our detections, we will use the

YOLOv5

data loader.

You can load data using the following code:


import cv2
import supervision as sv
from ultralytics import YOLO

image = cv2.imread("SOURCE_IMAGE_PATH")
model = YOLO('yolov5s.pt')

result = model(image)[0]
detections = sv.Detections.from_ultralytics(result).with_nms(threshold=0.5, class_agnostic=False)

This code will return a Detections() object with detections to which NMS was applied.

You can update the code above to adjust the threshold by which two or more detections need to overlap in order for NMS to be applied to those detections. You can also choose whether to apply NMS while considering the classes of overlapping bounding boxes.

Next Steps

supervision provides an extensive range of functionalities for working with computer vision models. With supervision, you can:

1. Process and filter detections and segmentation masks from a range of popular models (YOLOv5, Ultralytics YOLOv8, MMDetection, and more).
2. Process and filter classifications.
3. Compute confusion matrices.

And more! To learn about the full range of functionality in supervision, check out the supervision documentation.