Inference with ONNX exported model

[karlhm76]

Hello,

I have trained a model and I have tested it with detect.py and it works very well. I have exported the model to ONNX format using the command line:
python export.py --weights "best.pt" --img-size 640 --batch-size 1 --simplify --optimize --opset 14 --include onnx

and it produced an ONNX file.

I attempt to use the ONNX file for inference:

import torch
import torchvision
import onnx
import onnxruntime as ort
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

# Load the YOLOv5 model in ONNX format
ort_session = ort.InferenceSession("best.onnx")

# Preprocess the input image for YOLOv5
image_path = "20221128_00001L_00423.jpg"
image = Image.open(image_path)  # Load the image using PIL
original_width = image.width
original_height = image.height

resized_image = image.resize((640, 640))  # Resize to 640x640

# Convert PIL image to tensor and normalize
transform = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor(),  # Convert PIL image to tensor
    torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

input_tensor = transform(resized_image)
input_tensor = input_tensor.unsqueeze(0)  # Add a batch dimension

# Run the YOLOv5 model
outputs = ort_session.run(None, {'images': input_tensor.numpy()})

# Post-process the output tensor to get bounding box coordinates and confidence scores
batch = outputs[0]
boxes = outputs[0][0]

# Visualize the predicted bounding boxes on the image
plt.imshow(resized_image)
current_axis = plt.gca()

for box in boxes:
    # Convert box format from [x_center, y_center, width, height] to [x_min, y_min, x_max, y_max]
    x, y, w, h, c, l = box

    x = sigmoid(x)
    y = sigmoid(y)
    w = sigmoid(w)
    h = sigmoid(h)
    c = sigmoid(c)

    x_min, y_min = (x - w / 2) * original_width, (y - h / 2) * original_height
    x_max, y_max = (x + w / 2) * original_width, (y + h / 2) * original_height

    # Draw the bounding box on the image
    current_axis.add_patch(plt.Rectangle((x_min, y_min), x_max - x_min, y_max - y_min, color='red', fill=False, linewidth=2))
    current_axis.text(x_min, y_min, f'Confidence: {c:.2f}', bbox={'facecolor': 'white', 'alpha': 0.7})


plt.show()

and the bounding boxes are all moved to one side of the image, all confidences are 0.5 - this is a result of the sigmoid, obviously.
When using this same image with detect.py and the best.pt weights, it works perfectly.

I have tried removing the sigmoid, but that doesn't really help at all, but the bounding boxes are less skewed.

the first result of my output looks like this:

[4.1635551e+00 9.4634151e-01 1.2783272e+01 6.4645405e+00 1.4603138e-06 9.9997914e-01]

I have also attempted the solution here: [https://github.com//discussions/2170] with no success

How can I interpret the results of the inference?

[karlhm76]

I have worked it out

The model input doesn't need to be normalised.

The outputs don't need to be adjusted with sigmoid function.

I have another issue now, but I don't know whether it is a problem with YOLO or ORT.

The final code was:

import torch
import torchvision
import onnx
import onnxruntime as ort
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

# Load the YOLOv5 model in ONNX format
ort_session = ort.InferenceSession("best.onnx")

# Preprocess the input image for YOLOv5
image_path = "20221128_00001L_00423.jpg"
image = Image.open(image_path)  # Load the image using PIL
original_width = image.width
original_height = image.height

scale_x = image.width / 640
scale_y = image.height / 640

resized_image = image.resize((640, 640))  # Resize to 640x640

# Convert PIL image to tensor and normalize
transform = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor(),  # Convert PIL image to tensor
    #torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) Normalisation not required
])

input_tensor = transform(resized_image)
input_tensor = input_tensor.unsqueeze(0)  # Add a batch dimension

# Run the YOLOv5 model
outputs = ort_session.run(None, {'images': input_tensor.numpy()})

# Post-process the output tensor to get bounding box coordinates and confidence scores
batch = outputs[0]
boxes = outputs[0][0]
confidences = outputs[0][0][4]  # Confidence scores are stored in the 5th element

# Visualize the predicted bounding boxes on the image
plt.imshow(image)
current_axis = plt.gca()

max_confidence = 0
max_confidence_index = 0

for ix, box in enumerate(boxes):
    # Convert box format from [x_center, y_center, width, height] to [x_min, y_min, x_max, y_max]
    x, y, w, h, c, l = box

    #x = sigmoid(x) Sigmoid not required
    #y = sigmoid(y)
    #w = sigmoid(w)
    #h = sigmoid(h)
    #c = sigmoid(c) #Sigmoid not required for confidence either

    if c > max_confidence:
        max_confidence = c
        max_confidence_index = ix

x, y, w, h, c, l = boxes[max_confidence_index]

x_min, y_min = (x - w / 2) * scale_x, (y - h / 2) * scale_y
x_max, y_max = (x + w / 2) * scale_x, (y + h / 2) * scale_y

# Draw the bounding box on the image
current_axis.add_patch(plt.Rectangle((x_min, y_min), x_max - x_min, y_max - y_min, color='red', fill=False, linewidth=2))
current_axis.text(x_min, y_min, f'Confidence: {c:.2f}', bbox={'facecolor': 'white', 'alpha': 0.7})


plt.show()

[karlhm76]

I also discovered today that the confidences don't need to be run through sigmoid either!