如何从YOLOv5预测中得到类和包围盒坐标？

Question

我正在尝试对我的自定义YOLOv5模型执行推断。官方文档使用默认的detect.py脚本进行推断。我已经编写了自己的python脚本，但是我无法从模型的输出中访问预测类和边界框坐标。下面是我的代码：

import torch
    
model = torch.hub.load('ultralytics/yolov5', 'custom', path_or_model='best.pt') 
predictions = model("my_image.png")

print(predictions)

Answer 1

results = model(input_images)
labels, cord_thres = results.xyxyn[0][:, -1].numpy(), results.xyxyn[0][:, :-1].numpy()

这将为检测到的每个对象提供标签、坐标和阈值，您可以使用它来绘制边界框。您可以查看这个回购更详细的代码。

https://github.com/akash-agni/real-time-object-detection