python - 如何使用自定义模型从 YoloV5 推理中获取边界框坐标？

Question

我目前正在尝试使用我自己的脚本而不是 detect.py 从我的图像和我的自定义模型中获取边界框坐标。我想获得在图像上绘制边界框所需的坐标。有人可以帮我吗？

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

results = model(img, size=416)

Answer 1

尝试这个 ：

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)