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

Question

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

Answer 1

这个对我有用：

model.conf = 0.25  # confidence threshold (0-1)
model.iou = 0.45  # NMS IoU threshold (0-1)  

更多信息：

https://github.com/ultralytics/yolov5/issues/36

Answer 2

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

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