如何从 YOLOv5 预测中获取类和边界框坐标？

Posted 2023-03-12

【中文标题】如何从 YOLOv5 预测中获取类和边界框坐标？

【英文标题】：How to get class and bounding box coordinates from YOLOv5 predictions?

【发布时间】：2021-07-18 12:08:44

【问题描述】：

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

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

print(predictions)

【参考方案1】：

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

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

https://github.com/akash-agni/Real-Time-Object-Detection

【参考方案2】：

YOLOv5 ? PyTorch Hub 模型允许在纯 Python 环境中进行简单的模型加载和推理，而无需使用 detect.py。

简单推理示例

此示例从 PyTorch Hub 加载预训练的 YOLOv5s 模型作为 model 并传递图像以进行推理。 'yolov5s' 是 YOLOv5 的“小”模型。有关所有可用型号的详细信息，请参阅README。还可以加载自定义模型，包括自定义训练的 PyTorch 模型及其exported 变体，即 ONNX、TensorRT、TensorFlow、OpenVINO YOLOv5 模型。

import torch

# Model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # or yolov5m, yolov5l, yolov5x, etc.
# model = torch.hub.load('ultralytics/yolov5', 'custom', 'path/to/best.pt')  # custom trained model

# Images
im = 'https://ultralytics.com/images/zidane.jpg'  # or file, Path, URL, PIL, OpenCV, numpy, list

# Inference
results = model(im)

# Results
results.print()  # or .show(), .save(), .crop(), .pandas(), etc.

results.xyxy[0]  # im predictions (tensor)
results.pandas().xyxy[0]  # im predictions (pandas)
#      xmin    ymin    xmax   ymax  confidence  class    name
# 0  749.50   43.50  1148.0  704.5    0.874023      0  person
# 2  114.75  195.75  1095.0  708.0    0.624512      0  person
# 3  986.00  304.00  1028.0  420.0    0.286865     27     tie

详情请见YOLOv5 PyTorch Hub Tutorial。