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

【问题标题】：How to get bounding box coordinates from YoloV5 inference with a custom model?如何使用自定义模型从 YoloV5 推理中获取边界框坐标？
【发布时间】：2021-09-01 15:22:06
【问题描述】：

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


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

results = model(img, size=416)

【问题讨论】：

标签： python object-detection yolo yolov5

【解决方案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)

【讨论】：

是否有一些选项可以知道边界框的宽度
是的，您可以使用每个边界框的 xmax 和 xmin 来计算它（尝试将“boxes”列表转换为 pandas Dataframe