简介目标检测是计算机视觉中应用最广泛的技术之一其核心在于同时完成物体定位与分类。以YOLOv8为代表的一阶段检测模型凭借端到端的训练方式和高效的推理速度在工业场景中备受青睐。然而当目标不仅需要检测“有无”还需区分颜色、型号等细粒度特征时任务就变成了检测与细粒度分类的结合。工地安全帽的颜色识别正是典型场景不同颜色代表不同人员角色自动识别帽色能显著提升巡检效率。本文从数据准备、标注规范、模型训练到ONNX/TensorRT部署详细展示了使用YOLOv8训练多色安全帽检测模型的完整链路。针对光照变化、视角差异、类别不平衡等工程实际问题给出了可复用的解决方案。对于手头有数千张图片、希望落地目标检测项目的工程师尤其值得参考。 管过工地安全的人应该都有这种体会每天巡检最费精力的不是“有没有戴安全帽”而是“戴的是什么颜色”。施工现场的帽子颜色是有管理语义的——黄帽一般是一线作业人员白帽多为管理人员或监理红帽可能是访客或专职安全员蓝帽通常是电工、焊工这类特殊工种。每个区域能不能进、人员构成合不合理、有没有违规串岗都要靠帽色快速判断。以前全靠人的眼睛盯现在摄像头铺得到处都是自然想用目标检测把这活儿自动化。这也是我这次用YOLOv8做不同颜色安全帽检测的出发点数据集一共3000张最终训练好的模型可以直接拿去做实时推理。这篇博文就把整个流程完整拆开讲讲数据集怎么规划、标注要注意什么、训练参数怎么调、模型怎么部署以及实战里最容易踩的坑。适合手里有个几千张图片、想把目标检测真正落地的工程师尤其是第一次接触YOLOv8的读者。1.1 从“二分类”到“细分色”的定位转变普通安全帽检测任务本质上是一个二分类目标检测画个框判断框里是“戴了”还是“没戴”。但加上颜色维度之后任务性质就变了。它不再只是检测安全帽的有无而是检测细粒度分类同时完成——既要保证框能稳定框住目标还要让模型在同类目标里区分出黄、白、红、蓝几种颜色。YOLOv8的多类别输出机制天然支持这个需求它的检测头每个锚点都会输出一组类别概率类别之间的差异学习完全交给损失函数去拉大所以把不同颜色定义成不同类别模型训练起来并不会增加架构复杂度。但不要因为架构没变就掉以轻心。颜色分类在实际场景里比想象中难得多。同一个黄色安全帽在正午阳光下和建筑物阴影里的RGB值差距非常大不同厂商摄像头的白平衡处理也会让同一种颜色偏暖或偏冷安全帽用久了会褪色工地上沾满灰尘后黄色可能看起来像灰白色。更麻烦的是背景干扰挖机机身、反光背心、脚手架的涂料都可能和某一种帽色高度接近。我在标注阶段就遇到过一张图黄色帽子靠在黄色围挡旁边人眼都得盯两秒才分得清边界这种样本对模型来说就是高压题。1.2 为什么选YOLOv8而不是分割方案拿到需求后我第一反应先排除了实例分割方案。安全帽在画面里经常只有几十个像素大小分割模型在这种小目标上优势不明显反而推理开销大。检测方案用矩形框对遮挡、小目标、密集场景的鲁棒性更好后处理也更简单。YOLOv8相比v5和v7最大的提升在于anchor-free检测头简化了正负样本匹配逻辑训练更稳定而且ultralytics官方仓库把训练、验证、导出全链路包装好了对做工程的来说省了很多事。另外一个实际考虑是部署生态。工地上不可能都配一台顶配GPU服务器边缘盒子、NVR、嵌入式板卡才是主流硬件YOLOv8导出ONNX、TensorRT都很顺畅OpenCV DNN和ONNX Runtime都能直接推理RK3588这类国产板子也有现成的NPU适配工具链。可以说从训练到落地整个链路是通的这是它相比一些学术模型最大的优势。2. 3000张数据集的来源规划与标注规范2.1 数据从哪来公开数据集打底现场图片补量先说结论这3000张图我建议不要全部从网上扒也不要全部自己拍。最好的方式是“公开数据集打底现场图片补量”。目前能搜到的安全帽公开数据集以二分类为主真正带有颜色标注的不多数据质量也参差不齐很多标注框贴得不严或者类别命名混乱。我的做法是用公开数据里画质较好的一部分作为底子大约1500张再结合现场巡检相机、手机实拍大约1500张凑成3000张。有人会觉得3000张是不是太少了但实测下来对于安全帽这种类别外观高度统一的目标3000张是够用的。安全帽不像猫狗那样有千奇百怪的姿态它的形状、外观相对固定模型需要学的主要是不同场景下的颜色差异和尺度变化。我建议遵循6:2:2的划分也就是训练集1800张、验证集600张、测试集600张。这里要特别注意划分时必须以“场景”为单位而不是以“单张图”为单位。同一段视频序列里抽出来、画面几乎相同的图片不能一部分进训练集一部分进测试集否则验证结果虚高部署到新场景就现原形。图片收集时还要注意覆盖不同天气和时间段。我当时专门分了几类晴天上午、晴天下午、阴天、黄昏、夜间灯光、室内灯光每类至少占5%的比例。颜色检测最怕环境光单一如果3000张全是晴天正午拍的模型到了阴天环境基本就废了。2.2 类别设计四色加未戴人头一共五类类别定义是这类项目里最容易被忽略、却最影响结果的地方。安全帽颜色在实际工程中并没有统一标准不同工地甚至不同地域的管理惯例都不一样所以我这里给出一个通用的设计思路具体颜色可以按自己项目替换。我最终用的是这五类类别ID类别名称说明0yellow_helmet黄色安全帽一线作业人员1white_helmet白色安全帽管理人员/监理2red_helmet红色安全帽访客/专职安全员3blue_helmet蓝色安全帽特殊工种4head未佩戴安全帽的人头/头部为什么要单独加一个head类这是很多人做安全帽检测时漏掉的一步。如果只检测帽子模型其实学的是“帽子的外观”而很难学会“没戴帽子”这个负例。加上head类后模型被迫去学习头部和帽子之间的形态差异实际推理时就能区分“没戴帽子”和“帽子被遮挡”这两种情况。后处理逻辑也可以简化如果某个head检测框与任何helmet检测框重叠率低于阈值就判定为未戴安全帽。颜色类别数量不建议超过五到六类。工地现场常见的帽色就是黄、白、红、蓝再加一个橙色多为救援或外来参观人员已经顶天了。类别越多颜色相近的类之间就越容易互相混淆标注成本也会成倍上涨。如果你拿到的数据里某种颜色特别少比如红色只有几十张我建议先把红色合并到“其他颜色”里等数据补够了再单列而不是硬上。2.3 标注工具与YOLO格式要点标注工具我用的是LabelImg和X-AnyLabeling两个都是免费的。LabelImg是老牌工具轻量适合快速框选X-AnyLabeling带了SAM模型辅助分割半自动标注效率更高适合那种帽子目标很多、一张图要框十几二十个框的场景。至于具体选哪个看个人习惯导出格式只要是YOLO txt格式就行label标注软件内部会帮你转换成归一化的坐标。一个典型的YOLO标注文件长这样图片img_001.jpg对应img_001.txt0 0.482031 0.351562 0.186719 0.183594 4 0.629688 0.398438 0.072656 0.145312每一行分别是类别ID、归一化后的中心点x坐标、中心点y坐标、归一化后的框宽、框高。所有坐标都用图片宽高做了归一化取值范围在0到1之间。这里有个最容易犯的低级错误标注软件导出时会自动归一化但手工改文件或者用脚本处理时经常忘记除以图片宽高训练结果就是loss炸裂、bounding box全乱。标注时的框贴合度也要注意。我的经验是帽子目标框要刚好包住实体部分上下左右留1%到2%的边距即可不要把帽子下方的人脸或肩膀也包进来。因为类别学习的重点是帽子颜色包进来太多无关像素模型就会把脸的颜色也当作特征之一换个人种或换种肤色就可能误判。遮挡情况下如果安全帽被遮挡超过50%我建议直接不标避免给模型传递错误信息如果遮挡小于30%按完整框标注让模型学着去处理部分遮挡。2.4 标注中最容易翻车的三个坑第一是颜色标注不一致。这是最致命的。如果两张图中同样是黄色安全帽一张标成yellow_helmet另一张因为光线暗标成了white_helmet模型就会学到错误的颜色边界损失函数怎么调都降不下来。所以在标注规范里我要求一律按“肉眼在自然光下看到的真实颜色”来标严格禁止按照像素颜色值判断。第二是远距离小目标漏标。画面角落里比拇指还小的安全帽最容易漏。漏标等于给模型提供了错误负样本训练后模型会倾向于漏检小目标。第三是类别不平衡。我一开始五类的分布大致是黄色约40%白色约25%蓝色约20%红色约10%未戴人头约5%。红色样本太少导致训练后红帽的recall只有0.7左右明显低于其他颜色。后期我给红色样本单独做了水平翻转、亮度扰动、HSV色相微调把有效样本补到了800张指标才追上。3. 训练前配置GTX 1660 Ti也能跑的硬件与软件清单3.1 硬件选型与显存估算先说结论训练这套模型不需要多么夸张的硬件GTX 1660 Ti 6GB这种级别的显卡完全够用。我实际训练用的就是一台1660 Ti的机器选的模型是YOLOv8s输入分辨率640x640batch size设为16显存占用大约5.2GB刚好卡在边界上跑完300个epoch。如果你手里的显存只有6GB那就别硬上YOLOv8m或者更大的模型训练速度慢batch size还只能压到8以下反而影响BN层统计稳定性。推理阶段就更宽松了。CPU用ONNX Runtime跑单张640x640图片大概100到300毫秒做一两路视频流没问题要是上了GPU或者RK3588的NPU单帧推理时间可以压到30毫秒以内完全满足实时视频分析的需求。如果你用的是Jetson Nano或树莓派这类板子也可以用YOLOv8n模型参数更少帧率能再翻一倍。显存和batch size的关系可以用一个粗略公式估算YOLOv8s在640分辨率下单张图的训练显存占用大约在300到400MB之间再加上模型权重、梯度、优化器状态实际开销会翻倍。6GB显存跑batch 16已经是极限了batch 20很可能直接OOM。如果你发现显存不够优先调小batch size不要急着调小图片分辨率——分辨率太低小尺寸安全帽会直接丢失特征颜色更分不清。3.2 环境安装别在版本上浪费时间新手的卡点七成在环境配置上这里直接给一套经过验证的版本组合conda create -n yolov8 python3.10 -y conda activate yolov8 pip install torch2.1.2 torchvision0.16.2 --index-url https://download.pytorch.org/whl/cu118 pip install ultralytics8.2.0PyTorch版本用2.1.x或2.2.x都行不需要追求最新。网上有人问PyTorch 2.13支持不支持YOLOv8先别管这个ultralytics库对PyTorch版本的要求其实很宽只要不是太老的PyTorch 1.x都能正常跑重点是CUDA版本和显卡驱动要匹配。装完之后跑一句验证命令yolo predict modelyolov8s.pt sourcehttps://ultralytics.com/images/bus.jpg能输出推理结果就说明环境没问题。注意第一次运行会自动下载预训练权重如果下载比较慢可以提前手动把yolov8s.pt下载好放到当前目录下ultralytics会自动识别。3.3 数据目录组织与YAML配置数据目录按下面的结构组织这是ultralytics默认约定的格式照做就好helmet_dataset/ ├── images/ │ ├── train/ # 1800张 │ ├── val/ # 600张 │ └── test/ # 600张 ├── labels/ │ ├── train/ │ ├── val/ │ └── test/ └── helmet.yamlhelmet.yaml文件内容如下path是数据集根目录的绝对路径或相对路径train/val/test指向图片子目录names里的类别名称必须和标注文件里的类别ID一一对应顺序错一个整个模型就废了path: ./helmet_dataset train: images/train val: images/val test: images/test nc: 5 names: 0: yellow_helmet 1: white_helmet 2: red_helmet 3: blue_helmet 4: head一个很容易踩的坑val目录里的每张图片必须在labels/val里有对应的txt文件哪怕这张图一个目标都没有也要放一个空的txt文件没有对应标注文件的图片会被ultralytics跳过但不会报错最终导致实际参与验证的图片数量比预期少很多。4. 训练全流程模型选择、超参设置与收敛判断4.1 训练命令与关键参数解释我最终的训练命令长这样yolo detect train \ datahelmet.yaml \ modelyolov8s.pt \ epochs200 \ imgsz640 \ batch16 \ patience30 \ device0 \ workers4 \ optimizerauto \ seed42逐项解释一下。modelyolov8s.pt表示从官方预训练权重开始微调这是几乎所有场景的默认最佳选择。预训练模型已经在COCO数据集上学过通用的特征提取能力底层的边缘、纹理、形状特征直接迁移过来比自己从零初始化训练少花三倍时间最终精度也更高。epochs200不是越多越好我这里设置了patience30意思是验证集指标连续30个epoch没有提升就自动早停。optimizerauto让ultralytics根据模型自动选择优化器实测下来它自己选SGD或AdamW的效果通常比手动指定更好。imgsz640是速度和精度的平衡点。安全帽这种目标在监控画面里经常很小如果你有大量远距离目标可以考虑把imgsz提到960小目标recall会有明显提升但训练时间会增加一倍。对3000张图的数据量来说我用640训练后测试mAP已经足够高就没必要加分辨率。4.2 训练过程的损失曲线怎么判断训练开始后会在runs/detect/train目录下生成本轮的results.png这是判断训练是否正常的最重要依据。这张图里包含三类损失曲线box_loss边界框回归损失、cls_loss分类损失、dfl_loss分布焦点损失。以我的训练过程为例三者的曲线趋势非常典型前30个epoch快速下降30到80个epoch慢速下降80到150个epoch进入平台期只有微小的波动。判断训练是否正常的几条经验第一train_loss和val_loss应该同步下降。如果train_loss还在下降但val_loss开始回升说明已经过拟合了这时候调低epochs或者加数据增强才有意义继续跑只会让测试集指标变差。第二cls_loss曲线如果在某个epoch出现突然跳升几乎可以肯定是数据集里有标注错误的样本模型正在努力拟合一个错误标签损耗就炸了。第三mAP50曲线通常是阶梯式上升的不要因为中间出现几个epoch的平台期就急着重启训练给它一点耐心。4.3 收敛后主要看哪些指标训练结束或早停后ultralytics会在runs/detect/train/weights/下保存两个权重文件best.pt和last.pt。best.pt是验证集mAP最高的那一个epoch的权重last.pt是最后一个epoch的权重。我最终用best.pt做后续验证。自己训练的模型我建议重点看这几个指标指标含义参考合格值mAP50IoU阈值0.5下的平均精度均值0.90以上为优秀mAP50-95不同IoU阈值下平均精度的均值更严格0.70以上可接受Precision所有预测框里预测正确的比例误检越少越高0.90左右Recall所有真实目标里被找出来的比例漏检越少越高0.88以上F1-score精度和召回率的调和平均0.90左右我的最终结果是mAP50约0.938mAP50-95约0.746红色帽子的recall从0.7提升到0.87。对于安全帽检测这个场景mAP50达到0.93已经具备实际部署的条件了。5. 验证、混淆矩阵与嵌入式部署实践5.1 验证集跑一遍重点看混淆矩阵训练结束后先用下面的命令在验证集上正式评估一次yolo detect val datahelmet.yaml modelruns/detect/train/weights/best.ptultralytics会在结果目录里生成confusion_matrix.png这个图是判断颜色分类是否做对的关键。对于本任务混淆矩阵的主要观察点有两个一是yellow_helmet类别是否和white_helmet类别互相混淆。这两个颜色一个偏亮黄一个偏白在过曝的图片里确实容易混我在第一版模型里就看到了大约5%的黄帽被预测成白帽。解决方法是补充一批过曝和强烈阳光下的图片同时把标注规范里“以自然光下真实颜色为准”的标准再强调一遍。二是head类是否频繁被预测为某一种帽子。如果一张没戴帽子的头被预测成yellow_helmet说明模型学到的帽子特征还不够聚焦需要检查标注框是否包进了人脸。验证集上确认没问题后再用test目录里的600张图测试一次确认mAP与验证集相差不超过2个百分点差的多了说明数据划分有问题或者某些场景在验证集里过拟合了。5.2 导出ONNX并部署到嵌入式设备训练好的模型最终要落到硬件上跑。ultralytics官方的导出命令很简洁yolo export modelruns/detect/train/weights/best.pt formatonnx opset12 imgsz640导出后会得到一个best.onnx文件opset选12是为了兼容更多推理框架如果目标设备是最新的TensorRT或者OpenCVopset可以适当选新版。用ONNX Runtime做CPU推理的示例代码import cv2 import numpy as np import onnxruntime as ort session ort.InferenceSession(best.onnx) input_name session.get_inputs()[0].name img cv2.imread(test.jpg) img_rgb cv2.cvtColor(img, cv2.COLOR_BGR2RGB) resized cv2.resize(img_rgb, (640, 640)) input_tensor resized.astype(np.float16 if float16 in session.get_inputs()[0].type else np.float32) / 255.0 input_tensor np.transpose(input_tensor, (2, 0, 1))[None] outputs session.run(None, {input_name: input_tensor}) # outputs[0] 形状为 (1, 84, 8400)需要做NMS后处理推理输出的原始数据是一个维度为[1, 4类别数1, 8400]的特征图8400是YOLOv8在不同尺度特征图上的预测框总数需要自己写非极大值抑制后处理把重叠的框去掉保留置信度最高的那一个。这个后处理逻辑如果不想自己写可以直接用ultralytics的YOLO类加载ONNX模型它会自动完成前处理和后处理from ultralytics import YOLO model YOLO(best.onnx) results model(test.jpg)但如果是部署到RK3588、Jetson这类嵌入式设备我强烈建议用厂商自带的NPU工具链做量化导出。RK3588对应的工具链支持把ONNX模型转成RKNN格式Jetson系列可以直接用TensorRT加速。需要注意两个点一是量化时要用一批真实场景图片作为校准集不能只用训练集图片否则量化损失会比较大二是安全帽检测部署到板子上之后输入分辨率不要盲目调小我实测过降到416x416后远距离小目标的漏检率明显上升实时性和精度最终选了512x512作为折中方案。5.3 部署中遇到的硬坑最典型的坑是推理结果和训练时看到的可视化效果不一致。训练时有NMS后处理有置信度阈值过滤导出的ONNX裸输出是不过滤的如果你直接拿模型输出全部框去画图会看到一大堆置信度只有0.1的垃圾框。部署时要自己设定置信度阈值一般设为0.25到0.35之间NMS的IoU阈值设为0.5左右。另一个坑是颜色偏移。嵌入式设备摄像头通常有明显的色偏用PC上训练时用的图片做测试没问题一接到现场摄像头就发现黄色帽子全变成橙色了。解决思路是在采集训练数据时把现场摄像头的截图也放一部分进去让模型提前适应设备色彩特性。如果做不到就在摄像头端先做白平衡校正再喂给模型。这种数据域不一致的问题靠调模型参数是解决不了的只能在数据环节下功夫。6. 踩坑记录与下一步可选的改进方案6.1 我的实测踩坑清单回看这轮从数据到部署的全过程有几个坑花了很长时间才定位到根因值得记录一下。第一个是过曝图片导致黄白混色。有一批来自现场摄像头的图片安全帽区域严重过曝黄色帽子在画面上看起来几乎就是白色被模型全部预测成了white_helmet。我当时以为是模型没训好反复调损失函数和超参最后才发现是数据分布的问题。解决方法是把过曝图片挑出来做了曝光补偿处理和亮度增强后重新加入训练集同时让模型看到更多“看似白色但其实是黄色”的难例。这件事让我明白模型的错误往往是数据问题的镜像先查数据再调参数是处理一切检测异常的优先顺序。第二个是顶视摄像头视角下的安全帽几乎只有帽顶一个圆面和正常视角的安全帽外观差异巨大。第一版模型在这类视角下recall直接掉到0.6以下。后来补了一批从高处俯拍的工地图片模型才学会跨视角识别。第三个是类别不平衡对recall的影响。红色帽子样本太少cls_loss虽然总体下降但红色类的召回率始终上不去。补充增强后的红色样本后问题解决。这提醒我在设计数据分布时要让每个类别至少占15%否则模型很容易学会“偷懒”把所有不确定的框都分给高频类别。6.2 模型结构层面的改进选项如果你的数据集比3000张更大或者项目有更高的精度要求光靠YOLOv8s本身的空间可能不够可以在模型结构上动手。网上讨论比较多的几个方向我按性价比排序说一下看法。第一是把注意力机制嵌入到C2F模块里。EMA或ECA这类轻量注意力模块参数量增加很少但能明显提升颜色特征的区分度。实现方式是在ultralytics/nn/modules.py里定义新的C2F变体再把yolov8.yaml中对应层的编号替换进去。这个改法对颜色分类尤其管用因为注意力机制能帮模型更聚焦于帽子区域的颜色特征而不是背景干扰。EMA结合进C2F的代码实现在ultralytics改进教程里已经有现成方案直接搜索“EMA注意力机制融入C2F”就能找到对应的配置说明。第二是替换下采样模块。YOLOv8默认的下采样是简单的步长为2的卷积有研究者提出ADown模块先用平均池化和最大池化分支降低特征丢失再融合。对小目标的改善比较明显代价是推理耗时增加一点点。如果摄像头画面距离远、帽子很小这个方向值得试。第三是增加小目标检测头。YOLOv8的检测头通常输出三个尺度如果你发现大量安全帽只有几个像素大小可以考虑在P2层新增一个更浅尺度的检测头专门负责小目标。代价是显存占用明显增加6GB显卡基本跑不动需要上12GB以上显存的卡。6.3 数据层面的后续优化方向结构改进做完了再回头看数据。当前3000张数据覆盖了晴天、阴天、黄昏、夜间灯光等基本光照条件但还缺少雨天、雾天、逆光直射这三种恶劣场景。如果你的项目要部署的工地恰好处于多雨地区这几种缺失场景会严重影响实际表现。补数据的思路不是盲目增加数量而是每类场景补200到300张关键样本让模型见过足够多的“坏情况”。另一个被低估的方向是针对性测试和标注校验。在训练完第一版模型后我特意把测试图片按“容易识别”和“难以识别”分成两组让标注员重新审视难例组的标注质量结果发现不少漏标和错标。每轮重新训练前先修正旧数据比添加新数据更重要因为干净数据中等数量效果往往强过脏数据大数量。最后再分享一下我这次做完的整体体会安全帽颜色检测这个项目真正的瓶颈从来不在模型结构而在数据组织和设备适配。3000张图配合YOLOv8s已经能训练出具备落地能力的模型关键是每个环节都要按规范做。如果你手头恰好有几千张工地图片想做类似的事我的建议是不要纠结模型选yolov8s还是yolov8m先把数据整理干净、把颜色一致性统一好训练出来的模型一定不会让你失望。这个流程跑通一次之后后面再加新颜色、新场景都是流水线式的活会顺手很多。本文还有配套的精品资源点击获取