Skip to content
Advertisement

How to get class and bounding box coordinates from YOLOv5 predictions?

I am trying to perform inference on my custom YOLOv5 model. The official documentation uses the default detect.py script for inference. I have written my own python script but I cannot access the predicted class and the bounding box coordinates from the output of the model. Here is my code:

import torch
    
model = torch.hub.load('ultralytics/yolov5', 'custom', path_or_model='best.pt') 
predictions = model("my_image.png")

print(predictions)

Advertisement

Answer

results = model(input_images)
labels, cord_thres = results.xyxyn[0][:, -1].numpy(), results.xyxyn[0][:, :-1].numpy()

This will give you labels, coordinates, and thresholds for each object detected, you can use it to plot bounding boxes. You can check out this repo for more detailed code.

https://github.com/akash-agni/Real-Time-Object-Detection

User contributions licensed under: CC BY-SA
5 People found this is helpful
Advertisement