⛑️ 矿井安全智能哨兵YOLOv5煤矿工人安全帽检测系统全栈实战| 煤矿工人安全帽检测数据集 煤矿安全帽检测数据集 井下工人安全帽检测数据集煤矿井下环境复杂工人安全帽佩戴不规范是重大安全隐患。传统人工巡检效率低、成本高。本文基于7242张实地标注图像构建了高精度安全帽检测数据集并提供了从数据加载、YOLOv5模型训练到推理部署的完整PyTorch代码为智慧矿井建设提供可落地的视觉AI解决方案。 数据集核心指标与构成数据集划分图像数量占比用途训练集 (Train)579380%模型参数学习验证集 (Val)72410%超参数调优与早停测试集 (Test)72510%最终性能评估总计7242100%单一类别煤矿安全帽数据集特点场景专注所有图像均采集自煤矿井下或矿区环境包含不同光照、角度和遮挡情况。格式标准提供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格式标签并进行预处理importosimportcv2importtorchfromtorch.utils.dataimportDatasetfromtorchvisionimporttransformsclassMineSafetyHelmetDataset(Dataset):煤矿安全帽检测数据集支持训练/验证/测试集加载def__init__(self,image_dir,label_dir,transformNone):self.image_dirimage_dir self.label_dirlabel_dir self.transformtransform# 获取所有图像文件 (假设为.jpg格式)self.image_files[fforfinos.listdir(image_dir)iff.endswith(.jpg)]def__len__(self):returnlen(self.image_files)def__getitem__(self,idx):img_nameself.image_files[idx]img_pathos.path.join(self.image_dir,img_name)label_pathos.path.join(self.label_dir,img_name.replace(.jpg,.txt))# 读取图像并转为RGBimagecv2.imread(img_path)imagecv2.cvtColor(image,cv2.COLOR_BGR2RGB)height,widthimage.shape[:2]# 解析YOLO格式标签: class_id x_center y_center width height (归一化)boxes[]withopen(label_path,r)asf:forlineinf.readlines():partsline.strip().split()iflen(parts)!5:continueclass_id,x_c,y_c,w,hmap(float,parts)# 转换为 (x_min, y_min, x_max, y_max, class_id) 便于可视化x_minint((x_c-w/2)*width)y_minint((y_c-h/2)*height)x_maxint((x_cw/2)*width)y_maxint((y_ch/2)*height)boxes.append([x_min,y_min,x_max,y_max,int(class_id)])# 应用数据增强 (如果有)ifself.transform:# 注意albumentations等库支持同时变换图像和边界框transformedself.transform(imageimage,bboxesboxes)imagetransformed[image]boxestransformed[bboxes]# 转换为张量 (图像: CxHxW, 边界框: Nx5)imagetorch.from_numpy(image).permute(2,0,1).float()/255.0boxestorch.tensor(boxes,dtypetorch.float32)returnimage,boxes2. 数据集配置文件config/helmet_data.yaml# 对应主题场景为YOLOv5训练指定数据路径和类别train:./dataset/images/trainval:./dataset/images/valtest:./dataset/images/testnc:1# 类别数量 (仅安全帽)names:[Mine_Safety_Helmet]# 类别名称3. YOLOv5 模型训练脚本src/train_yolov5.py# 对应主题场景使用YOLOv5s预训练模型在安全帽数据集上进行微调importtorchimportyamlfromyolov5.models.commonimportDetectMultiBackendfromyolov5.utils.generalimportcheck_img_size,non_max_suppressionfromyolov5.utils.torch_utilsimportselect_devicefromtorch.utils.dataimportDataLoaderfromsrc.datasetimportMineSafetyHelmetDatasetfromtorchvisionimporttransformsimportalbumentationsasAfromalbumentations.pytorchimportToTensorV2defget_transform(trainTrue):定义训练和验证时的数据增强策略iftrain:returnA.Compose([A.HorizontalFlip(p0.5),# 水平翻转A.RandomBrightnessContrast(p0.2),# 亮度对比度调整A.Resize(height640,width640),# 统一尺寸A.Normalize(mean[0,0,0],std[1,1,1]),# 归一化ToTensorV2()],bbox_paramsA.BboxParams(formatpascal_voc,label_fields[class_labels]))else:returnA.Compose([A.Resize(height640,width640),A.Normalize(mean[0,0,0],std[1,1,1]),ToTensorV2()],bbox_paramsA.BboxParams(formatpascal_voc,label_fields[class_labels]))deftrain_yolov5(data_dir,epochs100,batch_size16,devicecuda):核心训练函数deviceselect_device(device)# 加载YOLOv5s预训练模型 (结构)modelDetectMultiBackend(yolov5s.pt,devicedevice)model.nc1# 更新类别数model.names[Mine_Safety_Helmet]model.to(device)# 创建数据集和数据加载器train_datasetMineSafetyHelmetDataset(os.path.join(data_dir,images/train),os.path.join(data_dir,labels/train),transformget_transform(trainTrue))val_datasetMineSafetyHelmetDataset(os.path.join(data_dir,images/val),os.path.join(data_dir,labels/val),transformget_transform(trainFalse))train_loaderDataLoader(train_dataset,batch_sizebatch_size,shuffleTrue,num_workers4,collate_fnlambdax:tuple(zip(*x)))val_loaderDataLoader(val_dataset,batch_sizebatch_size,shuffleFalse,num_workers4,collate_fnlambdax:tuple(zip(*x)))# 优化器与损失 (简化示例实际YOLOv5训练推荐使用官方train.py)optimizertorch.optim.SGD(model.parameters(),lr0.01,momentum0.937,weight_decay0.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_expprint(训练启动... (请参考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--img640--batch16--epochs100\--dataconfig/helmet_data.yaml\--weightsyolov5s.pt\--projecthelmet_training_results4. 模型推理与可视化src/predict.py# 对应主题场景加载训练好的模型对井下监控图像进行安全帽检测并标注importtorchimportcv2importmatplotlib.pyplotaspltfromyolov5.models.commonimportDetectMultiBackendfromyolov5.utils.generalimportnon_max_suppression,scale_coordsfromyolov5.utils.torch_utilsimportselect_devicedefdetect_helmets(image_path,model_pathweights/best_model.pt,conf_thres0.25):检测图像中的安全帽并返回标注后的图像deviceselect_device(cudaiftorch.cuda.is_available()elsecpu)modelDetectMultiBackend(model_path,devicedevice)model.eval()# 读取并预处理图像img0cv2.imread(image_path)# BGRimgcv2.cvtColor(img0,cv2.COLOR_BGR2RGB)imgcv2.resize(img,(640,640))img_tensortorch.from_numpy(img).permute(2,0,1).float().div(255.0).unsqueeze(0).to(device)# 推理withtorch.no_grad():predmodel(img_tensor)detectionsnon_max_suppression(pred,conf_thres,iou_thres0.45)[0]# 绘制结果ifdetectionsisnotNoneandlen(detections):detections[:,:4]scale_coords(img_tensor.shape[2:],detections[:,:4],img0.shape).round()for*xyxy,conf,clsinreversed(detections):labelfMine_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)returnimg0if__name____main__:result_imgdetect_helmets(path/to/test_image.jpg)cv2.imshow(Safety Helmet Detection,result_img)cv2.waitKey(0)cv2.destroyAllWindows() 模型优化与部署建议优化方向具体方法数据增强增加Mosaic、MixUp等高级增强提升模型对遮挡、光照变化的鲁棒性。类别不平衡本数据集为单类别若后续增加“未戴安全帽”等类别需关注正负样本平衡。模型轻量化使用YOLOv5s或更小的nano版本便于在边缘设备如Jetson Nano上实时推理。部署加速导出为ONNX或TensorRT格式获得2-3倍推理速度提升。部署命令示例# 导出为ONNXpython yolov5/export.py--weightshelmet_training_results/exp/weights/best.pt--includeonnx# 使用TensorRT (需安装对应工具)python yolov5/export.py--weightsbest.pt--includeengine 项目价值与扩展方向核心创新点高价值数据集7242张矿区实拍图像为安全帽检测提供稀缺数据资源。标准化流程提供从数据准备、模型训练到推理的端到端YOLOv5代码。直接应用场景可集成至矿井监控系统实现24/7自动安全巡检。扩展方向多类别检测增加“未戴安全帽”、“佩戴不规范”、“异物”等类别。行为分析结合姿态估计检测工人是否在危险区域取下安全帽。预警系统联动检测到违规时通过声光报警或即时消息通知安全员。历史数据分析统计各时段、各区域违规佩戴频率辅助安全决策。️ 技术标签#智慧矿井#安全帽检测#YOLOv5#目标检测#PyTorch#煤矿安全#计算机视觉#工业AI#实时监控#深度学习
