首页> 新闻动态> 产品技术> 工地智能监控：安全帽识别 AI算法从训练到落地

工地智能监控：安全帽识别 AI算法从训练到落地

作者：万物纵横

发布时间：2026-03-03 09:55

阅读量：

一、整体方案思路

安全帽识别本质是目标检测任务（识别图片/视频中人头区域是否佩戴安全帽），完整流程分为：

1. 数据准备（采集、标注、预处理）

2. 模型选择与训练（基于成熟框架快速调参）

3. 模型优化（轻量化、精度调优）

4. 落地部署（边缘设备/云端）

5. 系统集成（对接监控硬件、告警系统）

二、各阶段实操落地

1. 数据准备（核心基础）

（1）数据采集

数据源：工地实拍视频帧（不同光线、角度、施工场景）、公开安全帽数据集（如HelmetV2、SafetyHelmet-Dataset）

数据量要求：至少1000张标注样本（建议2000+，覆盖不同场景）

数据格式：JPG/PNG图片，分辨率建议1080P以内（平衡精度和速度）

（2）数据标注

标注工具：LabelImg（轻量免费）、LabelMe、LabelStudio（支持团队协作）

标注规则：

类别定义：helmet（佩戴安全帽）、no_helmet（未佩戴安全帽）

标注框：精准框住人头区域（而非全身）

标注格式：PASCAL VOC（XML）或YOLO（TXT）（推荐YOLO格式，训练效率更高）

（3）数据预处理

import os
import cv2
import numpy as np
from sklearn.model_selection import train_test_split

# 1. 数据增强（提升模型泛化能力）
def augment_image(image, bboxes):
    """
    简单数据增强：随机翻转、亮度调整
    image: 输入图片
    bboxes: 标注框 [x1, y1, x2, y2, class]
    """
    # 随机水平翻转
    if np.random.rand() > 0.5:
        image = cv2.flip(image, 1)
        h, w = image.shape[:2]
        # 调整标注框坐标
        for bbox in bboxes:
            x1, y1, x2, y2, cls = bbox
            new_x1 = w - x2
            new_x2 = w - x1
            bbox[0], bbox[2] = new_x1, new_x2
    
    # 随机亮度调整
    if np.random.rand() > 0.5:
        alpha = np.random.uniform(0.7, 1.3)  # 亮度系数
        image = cv2.convertScaleAbs(image, alpha=alpha, beta=0)
    
    return image, bboxes

# 2. 数据集划分（训练集/验证集 9:1）
def split_dataset(img_dir, label_dir, save_dir):
    os.makedirs(os.path.join(save_dir, 'train/images'), exist_ok=True)
    os.makedirs(os.path.join(save_dir, 'train/labels'), exist_ok=True)
    os.makedirs(os.path.join(save_dir, 'val/images'), exist_ok=True)
    os.makedirs(os.path.join(save_dir, 'val/labels'), exist_ok=True)
    
    # 获取所有图片文件名
    img_files = [f for f in os.listdir(img_dir) if f.endswith(('.jpg', '.png'))]
    train_files, val_files = train_test_split(img_files, test_size=0.1, random_state=42)
    
    # 复制文件到对应目录
    def copy_files(files, split):
        for f in files:
            # 复制图片
            src_img = os.path.join(img_dir, f)
            dst_img = os.path.join(save_dir, split, 'images', f)
            cv2.imwrite(dst_img, cv2.imread(src_img))
            
            # 复制标注文件
            label_file = f.ｒｅｐｌａｃｅ('.jpg', '.txt').ｒｅｐｌａｃｅ('.png', '.txt')
            src_label = os.path.join(label_dir, label_file)
            dst_label = os.path.join(save_dir, split, 'labels', label_file)
            with open(src_label, 'r') as fr, open(dst_label, 'w') as fw:
                fw.write(fr.read())
    
    copy_files(train_files, 'train')
    copy_files(val_files, 'val')
    print(f"数据集划分完成：训练集{len(train_files)}张，验证集{len(val_files)}张")

# 调用示例
# split_dataset('data/raw/images', 'data/raw/labels', 'data/processed')

2. 模型选择与训练

（1）模型选型（兼顾精度和速度）

模型	特点	适用场景
YOLOv8n/s	轻量、速度快、精度适中	边缘设备（摄像头/NVR）
YOLOv8m/l	精度高、速度稍慢	云端/高性能服务器
YOLOv5	部署生态成熟、资料多	快速落地验证

推荐优先使用YOLOv8（Ultralytics官方维护，开箱即用）。

（2）模型训练（基于YOLOv8）

前置条件：

# 安装依赖
pip install ultralytics opencv-python pillow numpy

训练代码：

from ultralytics import YOLO

# 1. 加载预训练模型（基于COCO数据集的预训练权重，加速收敛）
model = YOLO('yolov8n.pt')  # 轻量版，边缘部署优先；yolov8s.pt为标准版

# 2. 配置数据集yaml文件（示例：helmet.yaml）
"""
# helmet.yaml内容
path: ./data/processed  # 数据集根目录
train: train/images     # 训练集图片路径
val: val/images         # 验证集图片路径

# 类别
names:
  0: helmet
  1: no_helmet
"""

# 3. 开始训练
results = model.train(
    data='helmet.yaml',  # 数据集配置文件
    epochs=100,          # 训练轮数（根据验证集精度调整）
    batch=16,            # 批次大小（根据GPU显存调整，1080Ti建议16-32）
    imgsz=640,           # 输入图片尺寸
    lr0=0.01,            # 初始学习率
    device=0,            # GPU编号（CPU用cpu）
    patience=20,         # 早停策略（20轮精度不提升则停止）
    save=True,           # 保存最佳模型
    project='runs/helmet',# 结果保存目录
    name='helmet_detect' # 实验名称
)

# 4. 模型验证
metrics = model.val()  # 输出mAP、Precision、Recall等指标
print(f"验证集mAP50: {metrics.box.map50:.2f}")  # 目标：mAP50>0.9

# 5. 模型导出（保存为pt格式，后续用于部署）
model.export(format='onnx')  # 导出ONNX格式，适配边缘推理框架

3. 模型优化（落地关键）

工地场景多为边缘设备（如摄像头、工控机），需优化模型速度和体积：

1. 模型轻量化：

使用YOLOv8n（体积仅6MB，推理速度<10ms/帧）

量化：将FP32模型转为INT8（体积减少75%，速度提升2-3倍）

# YOLOv8量化示例
model = YOLO('runs/helmet/helmet_detect/weights/best.pt')
model.export(format='onnx', int8=True)  # 导出INT8量化的ONNX模型

2. 推理优化：

使用OpenVINO/TensorRT加速推理（适配Intel/英伟达边缘设备）

降低输入分辨率（如640→480，速度提升，精度损失<5%）

4. 落地部署（两种主流方案）

方案1：边缘部署（推荐，低延迟）

部署到工地摄像头/工控机，实时分析视频流：

import cv2
from ultralytics import YOLO

# 加载训练好的模型
model = YOLO('runs/helmet/helmet_detect/weights/best.pt')

# 打开摄像头/视频流（工地摄像头通常为RTSP协议）
cap = cv2.VideoCapture("rtsp://admin:123456@192.168.1.100:554/stream1")  # 替换为实际RTSP地址

# 告警阈值（连续n帧检测到未佩戴安全帽则触发告警）
ALARM_THRESHOLD = 3
no_helmet_count = 0

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break
    
    # 模型推理（设置conf=0.5过滤低置信度结果）
    results = model(frame, conf=0.5, imgsz=640)
    
    # 解析结果并绘制标注
    annotated_frame = results[0].plot()
    
    # 检测是否有未佩戴安全帽的情况
    detections = results[0].boxes
    no_helmet_detected = any([model.names[int(cls)] == 'no_helmet' for cls in detections.cls])
    
    # 告警逻辑
    if no_helmet_detected:
        no_helmet_count += 1
        if no_helmet_count >= ALARM_THRESHOLD:
            # 触发告警（可对接声光报警器、短信通知、平台推送等）
            cv2.putText(annotated_frame, "ALARM: NO HELMET!", (50, 50), 
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
            # 这里可添加告警接口调用代码（如调用短信API、MQTT推送等）
    else:
        no_helmet_count = 0
    
    # 显示实时画面
    cv2.imshow('Helmet Detection', annotated_frame)
    
    # 按q退出
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

方案2：云端部署（适合集中管理）

1. 边缘摄像头将视频帧推送到云端服务器；

2. 云端部署模型推理服务（如FastAPI+YOLOv8）；

3. 推理结果存储到数据库，前端可视化展示。

# 云端推理API示例（FastAPI）
from fastapi import FastAPI, UploadFile, File
from ultralytics import YOLO
import cv2
import numpy as np

app = FastAPI(title="Helmet Detection API")
model = YOLO('runs/helmet/helmet_detect/weights/best.pt')

@app.post("/detect")
async def detect_helmet(file: UploadFile = File(...)):
    # 读取上传的图片
    contents = await file.read()
    nparr = np.frombuffer(contents, np.uint8)
    frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    
    # 模型推理
    results = model(frame, conf=0.5)
    
    # 解析结果
    detections = []
    for box in results[0].boxes:
        cls = model.names[int(box.cls)]
        conf = float(box.conf)
        x1, y1, x2, y2 = map(int, box.xyxy[0])
        detections.append({
            "class": cls,
            "confidence": conf,
            "bbox": [x1, y1, x2, y2]
        })
    
    # 判断是否有未佩戴安全帽的情况
    has_no_helmet = any([d['class'] == 'no_helmet' for d in detections])
    
    return {
        "detections": detections,
        "alarm": has_no_helmet,
        "message": "No helmet detected!" if has_no_helmet else "All helmets are worn correctly."
    }

# 启动命令：uvicorn app:app --host 0.0.0.0 --port 8000