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:
JavaScript
x
7
1
import torch
2
3
model = torch.hub.load('ultralytics/yolov5', 'custom', path_or_model='best.pt')
4
predictions = model("my_image.png")
5
6
print(predictions)
7
Advertisement
Answer
JavaScript
1
3
1
results = model(input_images)
2
labels, cord_thres = results.xyxyn[0][:, -1].numpy(), results.xyxyn[0][:, :-1].numpy()
3
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.