从零搭建安全帽检测（5）— yolov8模型推理

Billy_HF

于 2025-08-25 11:50:54 发布

阅读量148

点赞数 6

CC 4.0 BY-SA版权

分类专栏：安全帽AI检测文章标签：人工智能计算机视觉深度学习

本文链接：https://blog.csdn.net/weixin_41957687/article/details/150763345

安全帽AI检测专栏收录该内容

5 篇文章

订阅专栏

在成功完成模型训练后，我们迎来了项目初始阶段的最终环节—模型推理部署。本章将重点介绍如何将训练好的YOLOv8模型应用于实际的安全帽检测场景，实现从训练到实际应用的完整闭环。

从训练到应用的转化之旅

模型训练让我们获得了具备识别能力的权重文件，而模型推理则是将这些"沉睡的智能"唤醒，让其在真实场景中发挥价值的关键步骤。如果说训练是"教模型学习"，那么推理就是"让模型工作"，是将算法能力转化为实际生产力的重要过程。

YOLOv8不仅在训练效率上表现优异，在推理部署方面同样提供了强大的支持。其优化的网络结构和高效的推理引擎，使其能够在各种硬件环境下实现实时检测，满足工业场景对速度和精度的双重需求。

推理：智能检测的实际绽放

模型推理是将训练好的模型应用于新数据的过程，在这个过程中我们将：

性能验证：在测试集上全面评估模型的准确率、召回率、mAP等关键指标
速度优化：通过模型压缩、量化等技术提升推理速度
可视化分析：直观展示检测结果，分析模型在不同场景下的表现
部署准备：为后续的实际应用部署做好技术准备

一、推理预配置

接上文，yolov8二次训练的结果保存在了runs/detect/hardhat_detection_v1/weights/文件夹下，二次训练出来的模型名称为best.pt，本文在推理时使用此pt文件进行安全帽检测。

二、模型推理

import os
from ultralytics import YOLO

# --- 网络模型配置 ---
TRAINED_MODEL_PATH = 'runs/detect/hardhat_detection_v1/weights/best.pt'

# --- 测试图片路径 ---
IMAGE_TO_PREDICT = 'data/1.jpg' 

# --- 置信度阈值 (confidence threshold) ---
CONF_THRESHOLD = 0.2 

# --- IoU 阈值 (Intersection over Union threshold) 用于非极大值抑制 (NMS) ---
#    用于过滤掉重叠的边界框。
#    您可以根据需要调整，例如 0.7 (默认)。
IOU_THRESHOLD = 0

# --- 脚本开始 ---

print(f"正在加载模型: {TRAINED_MODEL_PATH}")
try:
    # 加载训练好的 YOLOv8 模型
    model = YOLO(TRAINED_MODEL_PATH)
    print("模型加载成功！")
except Exception as e:
    print(f"错误：无法加载模型 '{TRAINED_MODEL_PATH}'。请检查路径是否正确，以及模型文件是否存在。")
    print(f"详细错误信息: {e}")
    exit() # 退出脚本

print(f"正在对图片 '{IMAGE_TO_PREDICT}' 进行推理...")
if not os.path.exists(IMAGE_TO_PREDICT):
    print(f"错误：图片文件 '{IMAGE_TO_PREDICT}' 不存在。请检查路径是否正确。")
    exit() # 退出脚本

# 进行推理
# save=True: 会将带有检测框的图片保存到 runs/detect/predict 目录下。
# show=False: 不会弹出图片显示窗口（在服务器或没有图形界面的环境下很有用）。
#             如果您在带有图形界面的桌面环境运行，可以设置为 True 来实时查看。
results = model.predict(
    source=IMAGE_TO_PREDICT, 
    save=True, 
    show=False, 
    conf=CONF_THRESHOLD, 
    iou=IOU_THRESHOLD,
    name='single_image_prediction' # 为本次推理结果指定一个名称，方便管理
)

# 打印推理结果的详细信息
print("\n--- 推理结果详情 ---")
for r in results:
    # r 是一个 Results 对象，包含了当前图片的推理结果
    print(f"图片路径: {r.path}")
    print(f"检测到 {len(r.boxes)} 个目标。")

    if len(r.boxes) == 0:
        print("  未检测到任何目标。")
    else:
        for box in r.boxes:
            cls_id = int(box.cls[0])           # 类别 ID (例如 0 或 1)
            conf = float(box.conf[0])          # 置信度
            xyxy = box.xyxy[0].tolist()        # 边界框坐标 [x1, y1, x2, y2]
            
            class_name = model.names[cls_id]   # 根据 ID 获取类别名称 (例如 'safe', 'dangerous')

            print(f"  - 类别: {class_name} (ID: {cls_id}), 置信度: {conf:.2f}, 坐标: "
                  f"[{xyxy[0]:.0f}, {xyxy[1]:.0f}, {xyxy[2]:.0f}, {xyxy[3]:.0f}]")

# 告知用户结果保存位置
# results[0].save_dir 是当前推理结果的保存目录
output_image_path = os.path.join(results[0].save_dir, os.path.basename(IMAGE_TO_PREDICT))
print(f"\n带有检测框的图片已保存到: {output_image_path}")
print("请检查该目录以查看推理结果。")

predict_image.py

这是一个用于安全帽检测的 YOLOv8 模型推理脚本，主要功能是使用训练好的模型对单张图片进行目标检测。

1. 配置部分

TRAINED_MODEL_PATH = 'runs/detect/hardhat_detection_v1/weights/best.pt'
IMAGE_TO_PREDICT = 'data/1.jpg'
CONF_THRESHOLD = 0.2
IOU_THRESHOLD = 0

模型路径：指向训练得到的最佳权重文件
测试图片：指定要检测的图片路径
置信度阈值：只显示置信度高于0.2的检测结果（降低误检）
IoU阈值：设置为0表示不过滤重叠框（可能显示多个检测结果）

2. 模型加载与验证

model = YOLO(TRAINED_MODEL_PATH)

加载训练好的模型权重
包含错误处理机制，确保模型文件存在且可用

3. 推理执行

results = model.predict(
    source=IMAGE_TO_PREDICT, 
    save=True, 
    show=False, 
    conf=CONF_THRESHOLD, 
    iou=IOU_THRESHOLD,
    name='single_image_prediction'
)

4. 结果解析与输出

for r in results:
    for box in r.boxes:
        cls_id = int(box.cls[0])           # 类别ID
        conf = float(box.conf[0])          # 置信度
        xyxy = box.xyxy[0].tolist()        # 边界框坐标
        class_name = model.names[cls_id]   # 类别名称

三、推理结果展示

(yolo_env) fanghong1997@LAPTOP-C3KDJQKL:~/csdn$ python predict_image.py
正在加载模型: runs/detect/hardhat_detection_v1/weights/best.pt
模型加载成功！
正在对图片 'data/1.jpg' 进行推理...

image 1/1 /home/fanghong1997/csdn/data/1.jpg: 448x640 4 safes, 54.6ms
Speed: 3.5ms preprocess, 54.6ms inference, 201.6ms postprocess per image at shape (1, 3, 448, 640)
Results saved to runs/detect/single_image_prediction

--- 推理结果详情 ---
图片路径: /home/fanghong1997/csdn/data/1.jpg
检测到 4 个目标。
  - 类别: safe (ID: 0), 置信度: 0.98, 坐标: [447, 93, 541, 330]
  - 类别: safe (ID: 0), 置信度: 0.96, 坐标: [247, 98, 328, 291]
  - 类别: safe (ID: 0), 置信度: 0.85, 坐标: [69, 95, 147, 288]
  - 类别: safe (ID: 0), 置信度: 0.24, 坐标: [188, 113, 232, 276]

带有检测框的图片已保存到: runs/detect/single_image_prediction/1.jpg
请检查该目录以查看推理结果。

推理日志

本次推理成功加载了训练好的安全帽检测模型，在单张图片中检测到4个安全帽目标，其中3个高置信度检测结果可靠，1个低置信度结果需要进一步验证。模型在GTX 1650显卡上实现了约4 FPS的推理速度，整体检测效果良好，结果已保存至指定目录。

推理示例

在推理验证环节，我们使用了训练集中的一张图片进行测试，并将检测置信度阈值设置为0.2以观察模型敏感度。从推理结果可以看出，当前模型性能仍较为粗糙，具体表现为：

(1) 首先，场景中共有6名工人，模型成功检测出其中4人的安全状态，但左二工人对应的检测框置信度仅为0.24。在实际应用场景中，这样的低置信度检测结果通常不可靠——例如，若将阈值提高到行业常用的0.5或0.6，该检测结果就会被过滤掉，从而导致漏检。

(2) 其次，另外两名工人未被检测到的原因可能包括：目标遮挡严重、光照条件不佳、目标尺寸过小、训练样本中类似场景不足，或者模型在复杂背景下的特征提取能力有限等。这些因素都可能导致模型对某些目标的响应值过低，无法达到即使是最低的置信度阈值。

至此，我们已经完成了从零搭建安全帽检测系统的整体框架，实现了包括数据准备、模型训练和推理验证在内的完整流程。在后续的章节中，我们将紧密结合目标检测领域的理论知识与YOLOv8的具体特性，从数据标注质量、模型结构优化、训练策略调整和推理参数优化等多个维度进行深入改进，最终构建一个准确、稳定且可靠的安全帽检测系统。