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.