时间戳

首页 > 解决方案 > 如何从自定义 YOLOv5 模型中设置和获取置信度阈值？

python - 如何从自定义 YOLOv5 模型中设置和获取置信度阈值？

问题描述

我正在尝试对我的自定义 YOLOv5 模型进行推理。官方文档使用默认detect.py脚本进行推理。

例子：python detect.py --source data/images --weights yolov5s.pt --conf 0.25

我已经编写了自己的 python 脚本，但我既不能在初始化期间设置置信度阈值，也不能从模型的预测中检索它。我只能获取标签和边界框坐标。这是我的代码：

import torch

model = torch.hub.load('ultralytics/yolov5', 'custom', path_or_model='best.pt') 
results = model("my_image.png")
labels, cord_thres = results.xyxyn[0][:, -1].numpy(), results.xyxyn[0][:, :-1].numpy()

标签： pythonpytorchinferenceyolov5

解决方案

要设置自定义训练的 YOLOv5 模型的置信度阈值，请使用以下命令：

import torch
model = torch.hub.load('ultralytics/yolov5', 'custom',
                       path='absolute/path/to/.pt/file', source='local')
model.conf = 0.25

要检索推理图像和其他值，请使用

# Inference Result
img = "my_image.png"
results = model(img)

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

推荐阅读