How to create a YOLO-NAS confusion matrix

Evaluation is an essential part of the computer vision model development process. While you are striving to build the first version of your model, model evaluation will help you understand baseline performance and judge how close your model is to being ready for production. When working on future versions of a model, evaluation helps you understand the impact of each change you make.

One key part of evaluating models is computing confusion matrices. A confusion matrix is a visualization that shows how your model is performing on the classes on which it was trained.

In this guide, we are going to show you how to use the open source supervision Python package to create a

YOLO-NAS

confusion matrix.

We will:
‍
1. Install supervision
2. Run inference on a dataset using a

YOLO-NAS

model
3. Create and plot a confusion matrix for the model

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 and Compute Matrix

First, we are going to load our dataset into a supervision.DetectionDataset() object. This object will contain information about all the images in a dataset. You can load datasets from many different model types, from YOLO to MMDetection. For this guide, we will use the

YOLO-NAS

data loader.

Once we load data, we will define a callback function that runs inference on a

YOLO-NAS

model. We will use that callback to run inference on every image in our dataset, and compute a confusion matrix that shows how the model performs on the dataset.

Create a new Python file and add the following code:
‍


import cv2
from super_gradients.training import models
import supervision as sv

image = cv2.imread(SOURCE_IMAGE_PATH)
model = models.get('yolo_nas_l', pretrained_weights="coco")

def callback(image: np.ndarray) -> sv.Detections:
  result = list(model.predict(image, conf=0.35))[0]
  return sv.Detections.from_yolo_nas(result)

confusion_matrix = sv.ConfusionMatrix.benchmark(
   dataset = dataset,
   callback = callback
)

confusion_matrix.plot()

Set the DATASET value as the path to the folder where your dataset is stored.

Then, run the code to create the confusion matrix.

Step #3: Plot Confusion Matrix

We can plot the confusion matrix showing the results of the

YOLO-NAS

model evaluation using the following line of code


confusion_matrix.plot()

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

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

Learn how to create confusion matrices for other models

Below, you can find our guides on how to create confusion matrices for other models.