![](/img/trans.png)
[英]How to get bounding box coordinates from YoloV5 inference with a custom model?
[英]How to get class and bounding box coordinates from YOLOv5 predictions?
我正在尝试对我的自定义 YOLOv5 model 进行推理。官方文档使用默认的detect.py
脚本进行推理。 我已经编写了自己的 python 脚本,但我无法访问预测的 class 和来自 model 的 output 的边界框坐标。这是我的代码:
import torch
model = torch.hub.load('ultralytics/yolov5', 'custom', path_or_model='best.pt')
predictions = model("my_image.png")
print(predictions)
YOLOv5 PyTorch 集线器模型允许在不使用detect.py
的情况下在纯 python 环境中进行简单的 model 加载和推理。
此示例从 PyTorch Hub 加载预训练的 YOLOv5s model 作为model
并传递图像以进行推理。 'yolov5s'
是 YOLOv5 'small' model。有关所有可用模型的详细信息,请参阅自述文件。 还可以加载自定义模型,包括自定义训练的 PyTorch 模型及其导出的变体,即 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 教程。
results = model(input_images)
labels, cord_thres = results.xyxyn[0][:, -1].numpy(), results.xyxyn[0][:, :-1].numpy()
这将为您检测到的每个 object 提供标签、坐标和阈值,您可以将其用于 plot 边界框。 您可以查看此 repo 以获取更详细的代码。
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.