YOLOv8 Tutorial¶

This tutorial will guide you through using trtutils with YOLOv8 models. We will cover:

Exporting ONNX weights from YOLOv8
Building a TensorRT engine
Running inference with the engine

Exporting ONNX Weights¶

YOLOv8 is implemented by Ultralytics and requires a two-step process for end-to-end ONNX export. First, export the basic ONNX weights:

# Install ultralytics if you haven't already
$ pip install ultralytics

# Export ONNX weights
# This will save the ONNX file in the same directory as your PyTorch weights
$ yolo export model=TORCH_WEIGHTS format=onnx

Building TensorRT Engine¶

Once you have the ONNX weights, build a TensorRT engine:

# build_yolo is an alias for the 'build' command with '--yolo' passed to it
python3 -m trtutils build_yolo \
    --onnx PATH_TO_WEIGHTS \
    --output PATH_TO_OUTPUT \
    --fp16 \
    --num_classes 80 \
    --iou_threshold 0.5 \
    --conf_threshold 0.25 \
    --top_k 100

Alternatively, if you want to export the engine using the Python API:

from trtutils.builder import build_engine, hooks

build_engine(
    onnx="yolov8.onnx",
    output="yolov8.engine",
    fp16=True,
    hooks=[hooks.yolo_efficient_nms_hook(
        num_classes=80,
        iou_threshold=0.5,
        conf_threshold=0.25,
        top_k=100,
    )]
)

Running Inference¶

The YOLO class provides a high-level interface for running YOLOv8 inference:

import cv2
from trtutils.models import YOLO, YOLOv8

# Load the YOLOv8 model
yolo = YOLO("yolov8.engine")

# OR, use the YOLOv8 class
yolo = YOLOv8("yolov8.engine")

# Read and process an image
img = cv2.imread("example.jpg")
detections = yolo.end2end(img)

# Print results
for bbox, confidence, class_id in detections:
    print(f"Class: {class_id}, Confidence: {confidence}")
    print(f"Bounding Box: {bbox}")