HOW TO GUIDE

Save YOLO-NAS Predictions to CSV

Learn how to save model predictions to a file.

Overview

When you are working with computer vision models, you may want to save your detections to CSV or JSON for further processing. This is especially useful in testing and debugging scripts, or applications where you want to log all results from your model to a plain text file.

Using the supervision Python package, you can

predictions in a few lines of code. In this guide, we will show how to plot and visualize model predictions.

We will:

1. Install supervision
2. Load a model and save predictions with the supervision Sink API
‍
Without further ado, let's get started!

YOLO-NAS and Image Annotation Resources

Explore these resources to enhance your understanding of YOLO-NAS and image annotation techniques.

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 save detections.

Step 2

Save Data

First, we need to load our computer vision model. Then, we can load detections into the supervision Detections() API. With our detections in the Detections() API, we can use the

CSV

sink to save our data.

You can save data using the following code:
‍


from super_gradients.training import models
import supervision as sv

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

csv_sink = sv.CSVSink("RESULT_CSV_FILE_PATH")
frames_generator = sv.get_video_frames_generator("SOURCE_VIDEO_PATH>")

with csv_sink as sink:
    for frame in frames_generator:
    		# add your model inference logic
        result = list(model.predict(image, conf=0.35))[0]
        detections = sv.Detections.from_yolo_nas(result)
        sink.append(detections, custom_data={'CUSTOM_LABEL':'CUSTOM_DATA'})

Step 3

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.