Guide - How to interpet onnx predictions from detection model

[romanvelichkin]

I converted fine-tuned yolo11 to onnx with input_size=640. I have 4 classes to detect.
Onnx model returns prediction of (1, 1, 8, 8400) size.

This 8 dimension contains x, y, w, h, class probabilities
Your main target is to extract those values and find most confident boxes.

# Run the model
ort_inputs = {ort_session.get_inputs()[0].name: img_array}
# Get result, remove redundant dimensions
ort_outputs = ort_session.run(None, ort_inputs)[0][0]

# Convert result to numpy and rotate it (8, 8400) -> (8400, 8)
output = np.array(ort_outputs)
# Now each row contains coords: x_center, y_center, w, h and class confidences
output = output.T

# Extract bounding box coordinates and objectness score
boxes = output[:, :4]  # First 4 cols for x, y, w, h; shape (8400, 4)
class_probs = output[:, 4:]  # Next cols for class probabilities; shape (8400, 4)

# Get the class indices with the highest confidence scores
class_indices = np.argmax(class_probs, axis=1)  # Shape (8400,)

# Get the maximum confidence scores
max_confidence_scores = np.max(class_probs, axis=1)  # Shape (8400,)

# Combine the results into a single array
results = np.column_stack((boxes, class_indices, max_confidence_scores))  # Shape (8400, 6)

max_confidence_index = np.argmax(results[:, 5])  # Get the index of the highest confidence score

# Extract the box with the highest confidence
best_box = results[max_confidence_index]  # This will be shape (6,)

# Split the details for clarity
x, y, w, h = best_box[:4]
class_index = int(best_box[4]) 
confidence_score = best_box[5]

# Display the results
print("Box with highest confidence:")
print(f"Coordinates: (x: {x}, y: {y}), Width: {w}, Height: {h}")
print(f"Class Index: {class_index}, Confidence Score: {confidence_score}")

[UltralyticsAssistant]

👋 Hello @romanvelichkin, thank you for sharing your conversion process and interpretation guide for YOLO ONNX predictions 🚀! This level of detail is super insightful for the community.

For clarity and to ensure we address any potential issues with using the ultralytics tools, here are a few pointers:

• If you encountered any 🐛 bugs during the ONNX conversion or model inference process, please provide a minimum reproducible example (MRE), including your model conversion code, details of the fine-tuned model, and the ONNX export/export configurations you used. This will help us debug more effectively.
• If you have specific ❓ questions about interpreting the ONNX outputs or need further guidance on ONNX workflows, please let us know and include more details, such as sample outputs or the shape/content of input arrays.

Additionally, ensure your ultralytics package is up-to-date to avoid any version-related issues. Upgrade using:

pip install -U ultralytics

For more guidance:

• Check out the Ultralytics Docs which provide comprehensive examples and methodologies for using YOLO models.
• Join the Ultralytics community where discussions flourish! 👇
  • Discord for real-time chat 🎧
  • Discourse for in-depth discussions
  • Subreddit for knowledge-sharing

For simplicity, for ONNX workflows and deployment environments, also review our tested setups:

• Notebooks for easy environment setup:
  
• Docker Support if needed for deployment.

Lastly, an Ultralytics engineer will be here soon to dive deeper into your workflow 🚀✨. Thank you for contributing to the ONNX inference space for YOLO models!

[glenn-jocher]

@romanvelichkin your approach is correct for interpreting ONNX predictions from a YOLO11 detection model. To ensure accurate results, apply non-maximum suppression (NMS) to filter overlapping boxes and retain the most confident detections. For further guidance, refer to the ONNX Runtime documentation or YOLO predict examples.