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

问题描述

我正在尝试对我的自定义 YOlov5 模型进行推理。 official documentation 使用认的 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()

解决方法

检查 detection.y 文件以获得使用该模型的明确说明。您必须在代码中进行模型推断后手动进行分数阈值处理。

这意味着检测分数的范围从 0.0 到 1.0,您必须使用置信度阈值对其进行过滤。