[Python] yolov8n onnx for object detection (ultralytics library) -> how to postprocess the...

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

Continue reading...