How to get bounding box coordinates from YoloV5 inference with a custom model?

Question

I am currently trying to get the bounding box coordinates from my image with my custom model by using my own script and not the detect.py . I would like to get the coordinates needed to draw bounding boxes on the image. Could someone help me please?

model = torch.hub.load('ultralytics/yolov5', 'custom', 'best.pt')
model = model.autoshape()

results = model(img, size=416)

Answer 1

Try this :

import torch
from matplotlib import pyplot as plt
import cv2
from PIL import Image

#Model
model = torch.hub.load('path to yolov5 folder', 'custom', path='path to model/best.pt', source='local')  # local repo
# Images
img = cv2.imread('yolov5/esa.jpeg')
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY )
# Inference
results = model(img, size=328)  # includes NMS

# Results
results.print()  
#results.show()  # or .show()

#results = results.xyxy[0]  # img1 predictions (tensor)
boxes = results.pandas().xyxy[0]  # img1 predictions (pandas)