⛑️ 矿井安全智能哨兵：YOLOv5煤矿工人安全帽检测系统全栈实战| 煤矿工人安全帽检测数据集 煤矿安全帽检测数据集 井下工人安全帽检测数据集

煤矿井下环境复杂，工人安全帽佩戴不规范是重大安全隐患。传统人工巡检效率低、成本高。本文基于7242张实地标注图像，构建了高精度安全帽检测数据集，并提供了从数据加载、YOLOv5模型训练到推理部署的完整PyTorch代码，为智慧矿井建设提供可落地的视觉AI解决方案。

📊 数据集核心指标与构成

数据集特点

• 场景专注：所有图像均采集自煤矿井下或矿区环境，包含不同光照、角度和遮挡情况。
• 格式标准：提供YOLO格式（TXT）和VOC格式（XML）标注，兼容主流目标检测框架。
• 开箱即用：数据集已按8:1:1比例划分，并附带classes.txt类别文件，可直接用于模型训练。
  
  
  

🏗️ 项目目录结构建议

mine_safety_helmet_detection/
├── dataset/                          # 数据集
│   ├── images/
│   │   ├── train/   (5793张)
│   │   ├── val/     (724张)
│   │   └── test/    (725张)
│   └── labels/
│       ├── train/   (对应YOLO .txt)
│       ├── val/
│       └── test/
├── config/
│   └── helmet_data.yaml              # 数据集配置文件
├── src/
│   ├── dataset.py                    # 自定义数据集加载类
│   ├── train_yolov5.py              # 训练脚本
│   ├── predict.py                    # 推理与可视化脚本
│   └── utils.py                      # 工具函数
├── weights/
│   └── best_model.pt                 # 训练好的模型权重
├── requirements.txt                  # 项目依赖
└── README.md

💻 核心代码实现与场景注释

1. 自定义数据集加载类 src/dataset.py

# 对应主题场景：加载煤矿安全帽图像与YOLO格式标签，并进行预处理
import os
import cv2
import torch
from torch.utils.data import Dataset
from torchvision import transforms

class MineSafetyHelmetDataset(Dataset):
    """煤矿安全帽检测数据集，支持训练/验证/测试集加载"""
    def __init__(self, image_dir, label_dir, transform=None):
        self.image_dir = image_dir
        self.label_dir = label_dir
        self.transform = transform
        # 获取所有图像文件 (假设为.jpg格式)
        self.image_files = [f for f in os.listdir(image_dir) if f.endswith('.jpg')]

    def __len__(self):
        return len(self.image_files)

    def __getitem__(self, idx):
        img_name = self.image_files[idx]
        img_path = os.path.join(self.image_dir, img_name)
        label_path = os.path.join(self.label_dir, img_name.replace('.jpg', '.txt'))

        # 读取图像并转为RGB
        image = cv2.imread(img_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        height, width = image.shape[:2]

        # 解析YOLO格式标签: class_id x_center y_center width height (归一化)
        boxes = []
        with open(label_path, 'r') as f:
            for line in f.readlines():
                parts = line.strip().split()
                if len(parts) != 5: continue
                class_id, x_c, y_c, w, h = map(float, parts)
                # 转换为 (x_min, y_min, x_max, y_max, class_id) 便于可视化
                x_min = int((x_c - w / 2) * width)
                y_min = int((y_c - h / 2) * height)
                x_max = int((x_c + w / 2) * width)
                y_max = int((y_c + h / 2) * height)
                boxes.append([x_min, y_min, x_max, y_max, int(class_id)])

        # 应用数据增强 (如果有)
        if self.transform:
            # 注意：albumentations等库支持同时变换图像和边界框
            transformed = self.transform(image=image, bboxes=boxes)
            image = transformed['image']
            boxes = transformed['bboxes']

        # 转换为张量 (图像: CxHxW, 边界框: Nx5)
        image = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0
        boxes = torch.tensor(boxes, dtype=torch.float32)
        return image, boxes

2. 数据集配置文件 config/helmet_data.yaml

# 对应主题场景：为YOLOv5训练指定数据路径和类别
train: ./dataset/images/train
val: ./dataset/images/val
test: ./dataset/images/test

nc: 1  # 类别数量 (仅安全帽)
names: ['Mine_Safety_Helmet']  # 类别名称

3. YOLOv5 模型训练脚本 src/train_yolov5.py

# 对应主题场景：使用YOLOv5s预训练模型在安全帽数据集上进行微调
import torch
import yaml
from yolov5.models.common import DetectMultiBackend
from yolov5.utils.general import check_img_size, non_max_suppression
from yolov5.utils.torch_utils import select_device
from torch.utils.data import DataLoader
from src.dataset import MineSafetyHelmetDataset
from torchvision import transforms
import albumentations as A
from albumentations.pytorch import ToTensorV2

def get_transform(train=True):
    """定义训练和验证时的数据增强策略"""
    if train:
        return A.Compose([
            A.HorizontalFlip(p=0.5),                     # 水平翻转
            A.RandomBrightnessContrast(p=0.2),           # 亮度对比度调整
            A.Resize(height=640, width=640),             # 统一尺寸
            A.Normalize(mean=[0,0,0], std=[1,1,1]),      # 归一化
            ToTensorV2()
        ], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['class_labels']))
    else:
        return A.Compose([
            A.Resize(height=640, width=640),
            A.Normalize(mean=[0,0,0], std=[1,1,1]),
            ToTensorV2()
        ], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['class_labels']))

def train_yolov5(data_dir, epochs=100, batch_size=16, device='cuda'):
    """核心训练函数"""
    device = select_device(device)
    # 加载YOLOv5s预训练模型 (结构)
    model = DetectMultiBackend('yolov5s.pt', device=device)
    model.nc = 1  # 更新类别数
    model.names = ['Mine_Safety_Helmet']
    model.to(device)

    # 创建数据集和数据加载器
    train_dataset = MineSafetyHelmetDataset(
        os.path.join(data_dir, 'images/train'),
        os.path.join(data_dir, 'labels/train'),
        transform=get_transform(train=True)
    )
    val_dataset = MineSafetyHelmetDataset(
        os.path.join(data_dir, 'images/val'),
        os.path.join(data_dir, 'labels/val'),
        transform=get_transform(train=False)
    )
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4, collate_fn=lambda x: tuple(zip(*x)))
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4, collate_fn=lambda x: tuple(zip(*x)))

    # 优化器与损失 (简化示例，实际YOLOv5训练推荐使用官方train.py)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.937, weight_decay=0.0005)
    # 此处省略完整训练循环，建议直接调用YOLOv5的train.py以获得最佳实践:
    # !python yolov5/train.py --img 640 --batch {batch_size} --epochs {epochs} \
    #        --data config/helmet_data.yaml --weights yolov5s.pt --project helmet_exp
    print("训练启动... (请参考YOLOv5官方train.py进行完整训练)")
    # 保存模型
    torch.save(model.state_dict(), 'weights/best_model.pt')

if __name__ == '__main__':
    train_yolov5(data_dir='./dataset')

建议使用YOLOv5官方训练命令（更稳定、功能完整）：

python yolov5/train.py --img 640 --batch 16 --epochs 100 \
                       --data config/helmet_data.yaml \
                       --weights yolov5s.pt \
                       --project helmet_training_results

4. 模型推理与可视化 src/predict.py

# 对应主题场景：加载训练好的模型，对井下监控图像进行安全帽检测并标注
import torch
import cv2
import matplotlib.pyplot as plt
from yolov5.models.common import DetectMultiBackend
from yolov5.utils.general import non_max_suppression, scale_coords
from yolov5.utils.torch_utils import select_device

def detect_helmets(image_path, model_path='weights/best_model.pt', conf_thres=0.25):
    """检测图像中的安全帽并返回标注后的图像"""
    device = select_device('cuda' if torch.cuda.is_available() else 'cpu')
    model = DetectMultiBackend(model_path, device=device)
    model.eval()

    # 读取并预处理图像
    img0 = cv2.imread(image_path)  # BGR
    img = cv2.cvtColor(img0, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, (640, 640))
    img_tensor = torch.from_numpy(img).permute(2, 0, 1).float().div(255.0).unsqueeze(0).to(device)

    # 推理
    with torch.no_grad():
        pred = model(img_tensor)
        detections = non_max_suppression(pred, conf_thres, iou_thres=0.45)[0]

    # 绘制结果
    if detections is not None and len(detections):
        detections[:, :4] = scale_coords(img_tensor.shape[2:], detections[:, :4], img0.shape).round()
        for *xyxy, conf, cls in reversed(detections):
            label = f'Mine_Safety_Helmet {conf:.2f}'
            # 绘制边界框和标签
            cv2.rectangle(img0, (int(xyxy[0]), int(xyxy[1])), (int(xyxy[2]), int(xyxy[3])), (0, 255, 0), 2)
            cv2.putText(img0, label, (int(xyxy[0]), int(xyxy[1]) - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

    return img0

if __name__ == '__main__':
    result_img = detect_helmets('path/to/test_image.jpg')
    cv2.imshow('Safety Helmet Detection', result_img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

🔧 模型优化与部署建议

部署命令示例：

# 导出为ONNX
python yolov5/export.py --weights helmet_training_results/exp/weights/best.pt --include onnx

# 使用TensorRT (需安装对应工具)
python yolov5/export.py --weights best.pt --include engine

🔮 项目价值与扩展方向

核心创新点

• 高价值数据集：7242张矿区实拍图像，为安全帽检测提供稀缺数据资源。
• 标准化流程：提供从数据准备、模型训练到推理的端到端YOLOv5代码。
• 直接应用场景：可集成至矿井监控系统，实现24/7自动安全巡检。

扩展方向

1. 多类别检测：增加“未戴安全帽”、“佩戴不规范”、“异物”等类别。
2. 行为分析：结合姿态估计，检测工人是否在危险区域取下安全帽。
3. 预警系统联动：检测到违规时，通过声光报警或即时消息通知安全员。
4. 历史数据分析：统计各时段、各区域违规佩戴频率，辅助安全决策。

🏷️ 技术标签

#智慧矿井 #安全帽检测 #YOLOv5 #目标检测 #PyTorch #煤矿安全 #计算机视觉 #工业AI #实时监控 #深度学习