yolov8n onnx for object detection -> how to postprocess the output?

[dominicnoy]

When using the Ultralytics library to perform model inference directly on an image, the model correctly predicts the bounding box positions and class identities:

    !pip install ultralytics opencv-python
    from ultralytics import YOLO
    import cv2
    import matplotlib.pyplot as plt
    
    # Load the YOLOv8n model
    model = YOLO('yolov8n.pt')
    
    # Path to your image
    image_path = '/path-to-image/example_image.png'
    
    # Run inference
    results = model(image_path)[0]  # get the first result
    
    # Load image with OpenCV
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    
    # Draw bounding boxes
    for box in results.boxes:
        x1, y1, x2, y2 = map(int, box.xyxy[0])
        conf = float(box.conf)
        class_id = int(box.cls)
        label = f"{model.names[class_id]} {conf:.2f}"
    
        # Draw rectangle and label
        cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(image, label, (x1, y1 - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2)
    
    # Show the image with bounding boxes
    plt.figure(figsize=(12, 8))
    plt.imshow(image)
    plt.axis('off')
    plt.title("YOLOv8n Object Detection")
    plt.show()

Next, I exported the model to onnx:

    # Load YOLOv8n
    model = YOLO('yolov8n.pt')
    
    # Export to ONNX
    model.export(format='onnx', opset=12, simplify=True)

When exporting the model to onnx, I have to do the post-processing:

    import onnxruntime as ort
    import numpy as np
    import cv2
    import matplotlib.pyplot as plt
    from scipy.special import softmax
    
    
    # Load ONNX model
    session = ort.InferenceSession("yolov8n.onnx", providers=["CPUExecutionProvider"])
    
    # Path to your image
    image_path = '/path-to-image/example_image.png'
    
    # Load and preprocess image
    image = cv2.imread(image_path)
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    resized = cv2.resize(image_rgb, (640, 640))
    input_tensor = resized.astype(np.float32) / 255.0  # normalize to [0, 1]
    input_tensor = np.transpose(input_tensor, (2, 0, 1))  # HWC to CHW
    input_tensor = np.expand_dims(input_tensor, axis=0)  # Add batch dim
    
    # Run inference
    inputs = {session.get_inputs()[0].name: input_tensor}
    outputs = session.run(None, inputs)[0]
    
    # Post-process outputs
    # Output shape: (1, num_outputs, 84), where 84 = 4 (box) + 1 (objectness) + 80 (class scores)
    preds = outputs[0].transpose(1,0)  # shape: (N, 84)
    
    
    # COCO 80 class names
    class_names = [  # only a few listed for brevity, full list has 80
        "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train",
        "truck", "boat", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench",
        "bird", "cat", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe",
        "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard",
        "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
        "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl",
        "banana", "apple", "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza",
        "donut", "cake", "chair", "couch", "potted plant", "bed", "dining table", "toilet",
        "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven",
        "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
        "hair drier", "toothbrush"
    ]
    
    # Confidence threshold and NMS
    **conf_thresh = 0.007**
    
    boxes = []
    scores = []
    class_ids = []
    
    for pred in preds:
        x1, y1, w, h = pred[0:4]
    
        objectness = 1 / (1 + np.exp(-pred[4]))
        class_scores = softmax(pred[5:])
    
        class_id = np.argmax(class_scores)
        class_score = class_scores[class_id]
    
        conf = objectness * class_score
    
        if conf > conf_thresh:
            boxes.append([x1, y1, w, h])
            scores.append(conf)
            class_ids.append(class_id)
    
    print(image.shape[1])
    print(image.shape[0])
    # Scale boxes to original image size
    scale_x = image.shape[1] / 640
    scale_y = image.shape[0] / 640
    
    print(scale_x)
    print(scale_y)
    
    
    for box, conf, cls_id in zip(boxes, scores, class_ids):
        cx, cy, w, h = box
        x1 = int((cx - w / 2) * scale_x)
        y1 = int((cy - h / 2) * scale_y)
        x2 = int((cx + w / 2) * scale_x)
        y2 = int((cy + h / 2) * scale_y)
    
        label = f"{class_names[cls_id]} {conf:.2f}"
        cv2.rectangle(image_rgb, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(image_rgb, label, (x1, y1 - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2)
    
    # Show result
    plt.figure(figsize=(12, 8))
    plt.imshow(image_rgb)
    plt.axis('off')
    plt.title("ONNX Inference with YOLOv8n")
    plt.show()

This code is able to display bounding boxes around the correct objects. However, this only works when I reduce the conf_thres to very low values, such as 0.007.
Moreover, the predicted class_ids are incorrect.

I haven't added non-maximum suppression yet, because I believe the issue occurs before that step. (but code with non-maxiumum suppression can be found in the colab code below)

Can anyone help me understand what I'm doing wrong?

here is the colab link:
https://colab.research.google.com/drive/1T-LFXvyZS3hKSfICp_W9Re1A7pW-mNdv

[UltralyticsAssistant]

👋 Hello @dominicnoy, thank you for sharing your detailed workflow and for using Ultralytics 🚀! For new users or those looking to deepen their understanding, we recommend checking out the Docs, which include Python and CLI examples that may help clarify common questions.

Since this appears to be a 🐛 Bug Report regarding ONNX model output and post-processing, please provide a minimum reproducible example (MRE) if you haven't already. This helps us debug more efficiently.

If your question is about custom training or inference, please share additional details such as dataset examples, full inference logs, or any modifications made to the pipeline. Also, review our Tips for Best Training Results to ensure optimal outcomes.

You're welcome to join the Ultralytics community in the way that suits you best. For real-time support, visit our Discord 🎧. For focused discussions, use Discourse. Or connect with peers on our Subreddit.

Upgrade

First, please upgrade to the latest ultralytics package with all requirements in a Python>=3.8 environment and PyTorch>=1.8 to ensure your issue isn't already resolved in a recent release:

pip install -U ultralytics

Environments

YOLO can be run in any of the following up-to-date verified environments (with all dependencies including CUDA/CUDNN, Python and PyTorch preinstalled):

• Notebooks with free GPU:
• Google Cloud Deep Learning VM. See GCP Quickstart Guide
• Amazon Deep Learning AMI. See AWS Quickstart Guide
• Docker Image. See Docker Quickstart Guide

Status

If this badge is green, all Ultralytics CI tests are passing. CI tests verify correct operation of all YOLO Modes and Tasks on macOS, Windows, and Ubuntu every 24 hours and on every commit.

This is an automated response 🤖. An Ultralytics engineer will review your discussion and offer further assistance soon!

[glenn-jocher]

Hi @dominicnoy,

Looking at your ONNX inference code, I see a few issues with how you're processing the model output:

1. The output format for YOLOv8 ONNX models is different from what you're expecting. The raw output needs to be transposed from (1, 84, N) to (N, 84), which you're doing, but the data interpretation is incorrect.

2. In YOLOv8's ONNX output, the boxes are already in xywh format and class confidences don't need the softmax function. The model has already applied the sigmoid activation.

3. You're missing the non-maximum suppression step which is crucial.

Here's how to correctly process the outputs:

import numpy as np
import cv2
from ultralytics.utils.ops import non_max_suppression
import torch

# After running inference with ONNX
outputs = session.run(None, inputs)[0]  # shape: (1, 84, N)

# Convert to PyTorch tensor and proper format for NMS
outputs = torch.from_numpy(outputs)
outputs = outputs.permute(0, 2, 1)  # Reshape to (1, N, 84)

# Apply NMS (handles confidence thresholding and box decoding)
conf_thres = 0.25  # More reasonable threshold
iou_thres = 0.45
results = non_max_suppression(outputs, conf_thres, iou_thres)

# Process detections
for i, det in enumerate(results):
    if len(det):
        # Rescale boxes to original image
        det[:, :4] = scale_boxes(input_shape, det[:, :4], original_shape)
        
        for *xyxy, conf, cls in det:
            x1, y1, x2, y2 = [int(x) for x in xyxy]
            label = f"{class_names[int(cls)]} {conf:.2f}"
            # Draw box and label

The non_max_suppression function handles the box decoding and confidence thresholding properly. This should fix both your confidence threshold and class prediction issues.

[dominicnoy]

Hi @glenn-jocher,

Thanks for your help, it's working now!

Just one note: I had to comment out outputs = outputs.permute(0, 2, 1) to get it working.

[glenn-jocher]

Hi @dominicnoy,

I'm glad to hear it's working now! Thanks for pointing out that you needed to comment out the permute operation. This makes sense - ONNX model outputs can sometimes vary in format depending on the export settings and version.

The important part is that the non_max_suppression function expects inputs in the format (batch_size, num_classes + 4 + num_masks, num_boxes), and it seems your model was already providing outputs in the correct format for direct processing.

If you're working with different YOLO models in the future, you might need to adjust the pre-processing of outputs depending on how they're structured. The key indicators are typically the shape and ordering of the tensors.