How to filter YOLOv5 detections

Before a model goes to production, you need to build logic around the model. This often involves filtering detections from the model. For instance, you may want to retrieve detections from a model that are above a specified confidence level, or detections in a particular region of an image or video.

Using supervision, you can filter detections from computer vision models in a few lines of code.

supervision is an open-source Python package with a range of utilities for working with computer vision models. For this tutorial, we will focus specifically on filtering detections with supervision, but you can also use supervision for model evaluation, visualizing detections, and more.

Below, we show how to filter

YOLOv5

detections with supervision. We will:

1. Install supervision
2. Load

YOLOv5

model detections
3. Show how to filter by confidence, area, and more

Without further ado, let's get started!

Step #1: Install supervision

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.

Step #2: Load Data

supervision has a standard class called supervision.Detections() that provides a range of functions for working with bounding box and segmentation detections.

You can load predictions into a Detections() object in a single line of code using the supervision "from" loaders.

For

YOLOv5

predictions, we will use the

from_yolov5()

data loader:
‍


import cv2
import torch
import supervision as sv

# load a pre-trained or custom-trained model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

# load an image and run inference
image = cv2.imread(SOURCE_IMAGE_PATH)
result = model(image)

# process detections
detections = sv.Detections.from_yolov5(result)

Step #3: Filter Detections

With supervision, you can filter detections by:

A specific class
A set of classes
Confidence
Bounding box area
Box dimensions
A specific zone

You can combine these filters to build the logic you need.

For this guide, we will focus on filtering by classes and confidence. To learn about the other filters available in supervision, check out the Detections() quickstart guide.

Filter by Class

To filter

YOLOv5

detections by class, use the following code:


detections = detections[detections.class_id == 0]

Replace the number 0 with the ID of the class whose predictions you want to retriev.e

Filter by Set of Classes

To retrieve

YOLOv5

detections by class, use the following code:


selected_classes = [0, 2, 3]
detections = detections[np.isin(detections.class_id, selected_classes)]

Where selected_classes contains the list of class IDs you want to filter.

Filter by Confidence

To retrieve

YOLOv5

detections over a specified confidence level, use the following code:


detections = detections[detections.confidence > 0.5]

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
4. Plot bounding boxes and segmentation masks

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