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

wpcxdonn  于 2023-01-29  发布在  Python
关注(0)|答案(2)|浏览(428)

我尝试在我的自定义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)
qhhrdooz

qhhrdooz1#

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

zvokhttg

zvokhttg2#

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

简单推理示例

本例从PyTorch Hub加载一个预训练的YOLOv5s模型model,并传递一个图像进行推理。'yolov5s'是YOLOv5的“小”模型。有关所有可用模型的详细信息,请参阅自述文件。也可以加载自定义模型,包括自定义训练的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

相关问题