ARTICLE DETAIL

资讯详情

深耕商务建站与企业官网运营的一线实战洞察。

基于YOLOv5 4.0的垃圾分拣目标检测实战:从环境配置到部署

基于YOLOv5 4.0的垃圾分拣目标检测实战:从环境配置到部署 垃圾分拣项目找上门的时候我心里其实是有点打鼓的——场景看着简单不就是几个垃圾桶里的瓶瓶罐罐嘛真做起来才知道要在传送带上一秒抓一个目标还要兼顾准确率和部署成本光靠通用检测模型是行不通的。当时团队里有同事提了一嘴用YOLOv5 4.0版本来做目标检测我抱着试一试的心态跑通了全套流程从环境配置到数据集制作再到模型训练最后落地成一条可用的分拣业务流。这篇博客就把这整条链路摊开讲环境怎么配、数据怎么标、模型怎么调、坑在哪里全部按实操顺序来适合刚接手检测项目、尤其是想在工业场景里做视觉分拣的同学参考。我用的硬件是普通工作站加一块中端显卡没有上集群也没有买高配服务器YOLOv5 4.0这套东西对硬件的要求其实比想象中低但环境搭不对照样跑不起来。先别急着下载代码下面从环境配置开始把每一步拆清楚。1. 环境配置裸机到能跑通的完整清单以及为什么4.0版本这么选1.1 Python环境与系统版本选择YOLOv5 4.0版本发布于2021年中段它的依赖锁得不算死但也不是随便拿个Python版本就能跑。我实测比较稳的组合是Python 3.8.10或3.9.7这两个版本在Windows和Linux下都能正常编译torch扩展。Python版本别选太高3.10及以上在部分机器上装torch1.9会碰到二进制兼容问题得不偿失。Windows下我建议用Anaconda创建独立虚拟环境Linux下用conda或venv都行。我的习惯是创建虚拟环境conda create -n yolov5 python3.8激活环境conda activate yolov5确认pip版本pip install --upgrade pip这里有个很多人会忽略的细节YOLOv5 4.0的requirements.txt里torch和torchvision的版本是带“”的比如torch1.7.0。如果直接无脑安装最新版torch碰上CUDA版本不匹配的概率很高。我当时踩过一个坑——装完最新torch后模型训练时loss直接nan查了半天是显卡驱动和CUDA库不匹配。后来老老实实按官网推荐的匹配表来一次性跑通。1.2 CUDA、cudnn与torch版本匹配这一步是环境配置里最核心、也最容易炸的一环。YOLOv5 4.0本身不做CUDA强绑定但GPU训练必须用正确的驱动。我的经验是显卡驱动版本决定你能支持到哪一版CUDA Toolkit命令nvidia-smi看右上角的CUDA Version这个数字是驱动支持的最高版本不代表你的环境已经装了对应CUDA。实际训练时用的CUDA是torch自带的不需要单独安装完整的CUDA Toolkit但你需要确保驱动够新。cudnn也建议装一下虽然torch可以自动加载部分算子但我遇到过推理速度忽高忽低的情况后来发现是cudnn未配置导致卷积算法未被优化。我这边的稳定组合是驱动470.xx CUDA 11.1 cudnn 8.0.5 torch 1.9.0 torchvision 0.10.0。装torch用官方匹配索引pip install torch1.9.0cu111 torchvision0.10.0cu111 -f https://download.pytorch.org/whl/torch_stable.html装完后验证一下python -c import torch; print(torch.cuda.is_available())显示True就说明GPU路线通了。这一步值得多花半小时确认因为后面所有训练时间都依赖这个基础。1.3 下载YOLOv5 4.0源码与环境依赖YOLOv5 4.0的代码结构相比后来的5.0、6.0版本要简洁不少特别适合学习或者二次改造。官方仓库的tag切换成v4.0也可以直接下载zip包解压。解压后目录里最显眼的就是train.py、detect.py、models、data、utils这些模块。接下来安装依赖pip install -r requirements.txtrequirements里包含numpy、opencv-python、matplotlib、pyyaml、tqdm这些常用库还有pycocotools用于COCO指标计算。如果是在Windows上pycocotools安装失败的概率比较大官方推荐装pycocotools-windows。我试过两次装Windows版可以避免编译报错安装命令是pip install pycocotools-windows依赖装完后可以立刻用预训练权重跑一次demo验证全链路是否通。做法是下载yolov5s.pt放进weights目录然后执行python detect.py --source data/images/bus.jpg --weights weights/yolov5s.pt如果能看到图像窗口弹出并绘制出检测框就说明环境彻底通了。这一步我建议别跳过因为很多问题会在第一次推理时就暴露。1.4 垃圾分拣场景为什么选用4.0版本这个选择值得多说一句。4.0版本的网络结构以CSPDarknet53为主配合PANet特征融合层它的anchor设置和损失函数已经相当成熟训练小目标比如矿泉水瓶的瓶盖效果也能接受。相对于更新的6.0版本4.0的代码结构更直白修改anchor或自定义dataloader的难度更低对于垃圾分拣这种类别不算多、背景相对固定的场景完全够用。垃圾分拣业务的特殊性在于一是类别少但形态差异大比如塑料瓶和纸箱的宽高比就完全不同二是场景里会出现密集堆叠的情况目标之间互相遮挡对NMS阈值和anchor尺寸特别敏感。4.0版本的自适应anchor计算功能和多尺度训练恰好能处理这部分问题。而且4.0版本推理权重文件更小后续导出ONNX或TensorRT的兼容性也更好。2. 数据集制作垃圾图片的采集、标注与划分2.1 采集什么样的图片才有用训练垃圾分拣模型最忌讳的是图片太“干净”。很多团队图省事从网上下载了一堆商品图背景是纯白或纯灰结果模型在实验室里AP高到0.99一到传送带上就拉胯。为什么会这样因为模型的泛化能力来自背景多样性、光照差异、角度变化和信息噪音。我自己整理数据时定了三条硬标准背景必须多样化覆盖传送带、垃圾桶、地板、户外、不同光照条件。目标必须有多角度每个类别至少要拍到360度不同姿态同时包含俯视、侧视、倾斜视角。必须有遮挡和堆叠样本分拣业务里瓶子会叠在一起纸箱会半开口如果训练集里全是单目标完美摆放模型永远学不会处理实际工况。采集方法上我推荐用视频抽帧而不是连拍照片。拿手机或工业相机对着垃圾堆横向扫一遍录10分钟视频然后每隔3到5帧抽一张这样能极大丰富目标姿态。抽帧脚本用OpenCV写就行几十行代码搞定。2.2 标注工具与标注规范YOLOv5 4.0的标注格式是每个图片对应一个txt文件每行是class_id x_center y_center width height坐标值都需要归一化到0到1之间。很多同学第一次做标注会用labelimg生成VOC格式的XML然后再转YOLO格式其实labelimg本身就直接支持YOLO格式输出不用多这一步转换。我用的工具是labelimg操作要点如下打开labelimg点击“PascalVOC”切换成“YOLO”模式。用Open Dir打开图片目录Change Save Dir指定标签保存目录。快捷键W新建框D下一张A上一张CtrlS保存。标注规范这块垃圾分拣很容易犯“框得太大”的毛病——把瓶子周围的背景也包进去了。正确的做法是框要紧贴目标可见部分的外轮廓如果目标被部分遮挡就标可见部分不要臆想完整的形状。置信度阈值在训练时会因为框不准确而出现严重波动宁可框小一点也不要框大。类别定义我建议用英文做标签名比如plastic_bottle、cardboard_box、aluminum_can、other_waste后面在yaml文件里再映射中文显示名。这个习惯可以避免很多编码问题。2.3 数据集划分与目录结构YOLOv5 4.0期待的目录结构有两种一种是官方COCO式的images和labels两个大目录然后在yaml里指定train和val路径另一种是每个数据集子集单独建目录比如dataset/ ├── images/ │ ├── train/ │ └── val/ └── labels/ ├── train/ └── val/两种都行我推荐第二种因为后续扩展test集更清晰。划分比例我按8:1:1也就是80%训练、10%验证、10%测试。注意划分的时候要保证同一个物体的不同图片不要被割裂得太严重但也不要刻意把同一视频帧的数据全放一组否则验证集会虚高。一个细节垃圾分拣数据集的类别数量不建议超过10类除非你确实需要区分很多细分材质。类别越多每个类的训练样本数就会摊薄模型在传送带高速场景下的误检率会明显上升。我们项目初期定了12类训练后mAP只有0.78后来合并成8类mAP直接涨到0.92所以宁可在类别定义上做文章也别硬扛分类细度。2.4 数据增强策略与脏数据清洗YOLOv5 4.0的train.py里自带数据增强包括mosaic、mixup、随机透视、HSV扰动、随机翻转这些。对于垃圾分拣场景最有用的是mosaic和HSV扰动。因为垃圾材质的颜色受光照影响特别大塑料瓶在不同灯光下反光效果完全不同HSV扰动会让模型忽略颜色绝对数值、更多学习形状和纹理特征。但增强也不是越狠越好。我在训练过程中发现mixup开启过强会导致模型对小目标的定位框偏移尤其是瓶盖这种小物件。4.0版本里通过超参数hsv_h、hsv_s、hsv_v可以调节增强幅度我习惯把hsv_s和hsv_v调低到0.4左右保留一点颜色信息又不至于让材质失真。脏数据清洗是很多人忽略的一步。标注完以后我建议跑一次分析脚本查看每类目标的数量、目标面积分布、图片像素宽度比。如果发现某张图片的标签特别多比如20多个框或者框面积小于图片面积的0.1%这种样本在训练时往往会算进太多背景负样本必要时可以直接删除。另外如果两张图片内容几乎一样也会导致过拟合需要做去重。3. 模型训练参数调整、训练过程监控与过拟合排查3.1 从预训练权重开始训练的底层逻辑垃圾分拣数据集规模通常只有几千张从头训练一个检测模型非常亏。YOLOv5 4.0提供了一组在COCO上预训练好的权重这些权重已经学会了“什么是物体边缘”“什么是纹理”“什么是遮挡”这些通用特征。我们只需要在它的基础上微调最后几层让它更懂“瓶子”“纸箱”“易拉罐”这几个具体类别的模样。我在项目里使用yolov5s.pt作为初始权重考虑到部署设备是嵌入式级别的板卡没有直接上yolov5l或者yolov5x。有个常见误区是模型越大效果越好但垃圾分拣的实时性要求很高yolov5s在1080Ti上推理能达到6到8毫秒一帧完全够用yolov5l虽然mAP高一点但推理时间翻倍工业现场根本用不了。训练命令python train.py --data garbage.yaml --cfg models/yolov5s.yaml --weights weights/yolov5s.pt --batch-size 16 --epochs 200 --img 640--img 640的含义是训练输入尺寸640x640垃圾分拣目标一般中等偏小这个尺寸比较平衡。如果追求更快的训练能把尺寸降到416但小目标的召回率会下降不太推荐。3.2 数据配置yaml文件写法与anchor调整在写yaml文件前我建议先按官方数据集格式组织目录。garbage.yaml的写法如下train: dataset/images/train val: dataset/images/val test: dataset/images/test nc: 8 names: [plastic_bottle, cardboard_box, aluminum_can, glass_bottle, plastic_bag, wood_block, metal_parts, other_waste]这里面最容易出问题的是路径写法。YOLOv5 4.0中路径可以写相对路径也可以写绝对路径我在Windows上碰到过一次路径分隔符反斜杠导致的读取失败改成正斜杠就正常了。Anchor是YOLOv5一个比较特别的设计。4.0版本在训练时会自动根据标注数据重新计算anchor尺寸所以如果数据集里目标比COCO普遍小很多anchor会自动变小不需要手工调整。但自动anchor计算也有个前提——目标标注必须足够准确。如果标注框普遍松弛计算出的anchor也会偏大。训练日志里会有这样一句Autoanchor: 4.6 anchors, 8.2, 9.1, etc, Better than old anchors看到这句话就说明anchor已经重算过了。如果显示Unchanged anchors说明预设anchor已经够用不用担心。3.3 训练过程中的指标监控与loss曲线判断训练时我最关心的不是总loss而是box_loss、obj_loss和cls_loss这三个分量的变化。YOLOv5 4.0训练会实时打印类似这样的内容Epoch 0: box_loss0.068, obj_loss0.035, cls_loss0.021, precision0.52, recall0.65, mAP0.50.48如果box_loss一直下降说明定位框学习正常如果obj_loss降不下去说明目标和背景的区分度不够这时候需要考虑增加难负样本比如空垃圾袋、纸屑这类容易被当成目标的“疑似物”cls_loss抖动剧烈往往和某个类别的样本过少有关。我更习惯用tensorboard实时监控曲线YOLOv5 4.0自带了tensorboard --logdir runs的启动命令。训练了大概50轮后能明显看到曲线斜率变缓这是正常现象。到150轮左右mAP基本稳定再往后训练收益很低反而有过拟合风险。3.4 过拟合问题排查一个典型的分拣业务案例我们项目第一次训练时验证集mAP达到了0.94很漂亮但现场实测识别率只有0.6。这种现象几乎可以断定是过拟合。怎么排查先看训练集和验证集的loss曲线。如果训练集loss持续下降而验证集loss在第120轮后开始反弹说明模型开始死记硬背训练集特征了。我做了三件事关掉部分数据增强将--cache的缓存大图勾选去掉同时调低mosaic概率。增加验证集数量原验证集只占总数据5%后来重新划分成15%。在模型最后一层加dropoutYOLOv5没有简单的dropout开关但可以通过降低模型深度系数、增加正则化来缓解。我把yolov5s.yaml的depth_multiple从0.33降到0.3相当于让模型变窄一点强制它学更抽象的特征。这三步做完现场识别率恢复到0.88。经验就是垃圾分拣场景的数据集规模不会太大与其追高训练集指标不如把验证集指标压稳。3.5 训练完成后的指标评估与选择best.pt还是last.pt训练结束会在runs/train/exp下生成两个权重文件best.pt和last.pt。名字已经说明了一切——best.pt是训练过程中验证集mAP最高的权重last.pt是最后一轮的权重。工业落地上我强烈建议用best.pt但要注意best.pt的“best”是基于验证集算出来的验证集和实际分拣场景可能还有分布差异。稳妥的做法是挑三到四个不同训练阶段的best.pt用现场录像截图做成一个独立测试集跑val.py对比出真正的“最优”。这个额外验证步骤能过滤掉很多“纸面mAP大师”。另外验证集mAP0.5和mAP0.5:0.95的区别也要关注。垃圾分拣业务一般只看mAP0.5因为分拣机械臂抓取时并不严格要求框和真实目标的高度贴合IOU阈值0.5已经足够。如果mAP0.5:0.95很低但mAP0.5很高说明定位精度一般但对业务够用。4. 部署到分拣业务流模型导出、推理加速与现场调优4.1 从PyTorch到ONNX再到TensorRT的导出路径训练完后的PyTorch权重没法直接在生产环境里跑需要一个导出过程。YOLOv5 4.0自带了export.py能导出ONNX、TorchScript、CoreML等格式。我这边最终部署用的是TensorRT因为分拣工位的算力卡非常怕CPU推理而TensorRT能最大程度利用GPU的并行能力。导出ONNX的命令python export.py --weights runs/train/exp/best.pt --include onnx --opset 11导出时会有一个注意点YOLOv5 4.0的检测头里有部分动态shape操作直接用onnxruntime跑会报错需要在export时加--dynamic参数。但我更推荐的是一步到位导出TensorRT enginepython export.py --weights best.pt --include engine --device 0TensorRT引擎是序列化过的优化计算图推理速度比ONNX还快。前提是TensorRT版本和你用的CUDA版本兼容我这边用了TensorRT 8.0配合CUDA 11.1非常顺。4.2 推理端代码改造与输入预处理YOLOv5 4.0官方detect.py是面向演示的直接放到生产环境会有问题一是每帧都打印日志和画框影响吞吐二是一次只能处理单张图无法应对连续帧流水。我改造推理端时做了三件事移除可视化部分只返回检测框数组。使用torch.cuda.synchronize()控制帧率。把输入图像resize为640x640时使用letterbox方式保留长宽比而不是强制拉伸。这个细节极其重要——垃圾瓶子细长如果拉伸变形宽高比信息丢失检测框会严重偏移。letterbox的算法本身很简单计算缩放比例补边为114像素的灰色。YOLOv5 4.0的utils里已经有letterbox函数直接复用即可。如果你自己写一定要记得在推理前对图像执行同样的预处理、推理后对坐标做反向映射。4.3 NMS阈值和置信度阈值的现场调优模型训练时的阈值和部署时的阈值是两回事。训练阶段的conf_thres和iou_thres影响的是mAP计算方式部署阶段的这两个参数直接决定现场检测行为。我刚开始部署时用的默认阈值conf_thres0.25, iou_thres0.45结果是传送带上偶尔漏捡。后来把conf_thres降到0.15把iou_thres调到0.5漏检率好了不少但误检略有增加。于是我在中间加了一个按类别设定阈值的逻辑对“塑料瓶”这种形态稳定的类别用0.3对“其他垃圾”这种外形不规则的类别用0.15。这样整体识别效果更符合分拣业务的需求。还有一个坑当目标密集堆叠时NMS的max_det参数默认是300如果一帧图像里检测超过300个框后面的目标会被丢弃导致漏检。垃圾分拣场景单帧一般不超过50个目标所以我把max_det设为100减少内存压力。4.4 现场现场光照变化与鲁棒性验证分拣传送带的现场环境远不是实验室那种稳定光源。我在部署时发现上午9点阳光照进厂房时模型对“铝罐”的识别率跌了20%。原因是铝罐高光反射蛋白白色区域在图片里和透明塑料袋特别像。解决方案不是重新采集大量视频数据而是先用opencv对接入摄像头的曝光参数做了硬约束把曝光时间固定、自动白平衡关闭。另外在内网里部署了夜间红外补光灯进一步杜绝环境光波动。这种工程手段比加数据更直接。但注意如果换了摄像头型号曝光参数需要重新标定不能直接套用。4.5 与PLC和机械臂的联动逻辑分拣业务的最终目标是把检测结果转成机械臂抓取指令。这里有个检测速度之外的瓶颈YOLOv5推理返回的是图像坐标系里的框坐标而机械臂需要的是相机坐标系或世界坐标系下的抓取点。我的做法是先用张正友标定法对相机做内参标定再通过手眼标定建立相机到机械臂底座的变换矩阵。之后推理出的目标中心点坐标通过矩阵变换映射到机械臂空间。由于传送带一直在动我还加了一个简单的速度补偿根据传送带线速度和检测到目标到机械臂末端的运动时间差预测目标到达抓取点时的新位置。这部分一开始用固定速度仿真效果还行但实际传输带速度波动时会抓偏后来在编码器上加了一个滑动平均滤波误差控制在0.5厘米以内。这一整套逻辑下来从图像采集到最后机械臂抓取端到端延迟大概在120毫秒左右满足低中速传送带0.3米/秒以下的分拣需求。5. 训练之外必须知道的事数据迭代闭环和成本控制5.1 建立“脏样本回流”机制模型上线后绝对不是一劳永逸。分拣现场每周都会出现新的“意外垃圾”比如湿纸巾、果皮、捆扎带这些在训练集里根本没有对应类别模型会强行把它们预测成最相近的类别结果机械臂抓了不该抓的东西。解决方法是建立回流机制现场IPC把置信度低于0.3的检测框和对应截图自动保存到一个“待标注”目录。每隔两周标注人员只处理这些新增难例补充到训练集重新训练一轮。这个机制第一周就帮我们发现了70多个误检案例第二周模型识别率又涨了4个百分点。对于垃圾分拣这种长尾分布极其明显的场景这个闭环比一开始堆数据集更经济、更有效。5.2 硬件选型与训练时间成本估算在训练成本上YOLOv5s用单张1080Ti训练200轮、每轮约1000张图大概需要2.5到3小时。如果是用CPU训练同样的工作量可能要40个小时以上。所以哪怕只是短期项目我也建议租一块云GPU来跑训练而不是用本机CPU硬抗。训练完成后把best.pt拉回本地部署即可。数据集标注成本也要算清楚每张图如果有2到3个目标熟练标注员每小时能标注80到120张图。一个3000张图的垃圾数据集纯标注工作在25小时左右。看起来不多但加上质量检查和返工实际大概需要2人日。所以标注前定严格规范非常有必要否则返工成本会翻倍。5.3 类别不均衡和背景负样本的取舍垃圾分拣数据集里塑料瓶往往占60%以上而金属零件可能只有5%。直接按原始分布训练模型会对金属零件类别的召回率特别低。我的做法是对样本量少的类别做离线复制增强不是简单的复制粘贴而是对该类图片做小幅旋转、平移、缩放后重新保存为新的样本。同时在下采样阶段限制每张图每类的实例数不超过30防止一个场景里的同类别目标刷屏导致训练失衡。背景负样本是另一个容易被忽略的坑。分拣场景里空无一物的传送带图片在训练集里一张都不能少。这些负样本能让模型的obj_loss快速收敛减少误报。我每批次会混入5%到10%的空背景图效果很明显。5.4 模型更新时的权重管理每个版本迭代出的best.pt我都会在权重文件名里加上日期和数据量后缀比如best_20240601_3200img.pt。这样在比对现场效果时能快速定位到底是哪套数据训出来的模型。有一次现场反馈识别率下降排查了半天发现是某个同事用了旧版模型文件没有同步更新。权重管理听着啰嗦但在工业项目里非常重要。另外最好保留训练当时的opt.yaml文件YOLOv5 4.0每次训练都会自动保存。这个文件记录了所有超参数想要复现结果或排查差异时直接载入即可不用猜。6. 一段供直接参考的完整训练流程脚本这里把我自己在项目里常用的训练流程串起来给你一份可以照抄的清单不包含花哨操作只保留稳定跑通的部分。# 1. 激活环境 conda activate yolov5 # 2. 数据目录准备好后先验证标注格式 python -c import os img_dir dataset/images/train label_dir dataset/labels/train count 0 for img in os.listdir(img_dir): label os.path.join(label_dir, img.replace(.jpg, .txt)) if not os.path.exists(label): print(缺少label:, img) count 1 print(缺失标签数量:, count) 运行这个验证脚本能第一时间发现漏标问题比直接开训练再报错要高效得多。我曾经靠这个脚本一次性找出37张漏标图片。# 2. 修改garbage.yaml路径确认nc和names正确 # 3. 开始训练 python train.py --data garbage.yaml --cfg models/yolov5s.yaml --weights weights/yolov5s.pt --batch-size 16 --epochs 200 --img 640 --project garbage_project --name v1 # 4. 训练完后在runs/train/garbage_project/v1/weights下找到best.pt # 5. 用val.py评估测试集 python val.py --data garbage.yaml --weights runs/train/garbage_project/v1/weights/best.pt --img 640 # 6. 导出TensorRT engine python export.py --weights runs/train/garbage_project/v1/weights/best.pt --include engine --device 0我训练时加过很多额外参数比如--cos-lr、--label-smoothing但实测下来对垃圾分拣数据集提升不大反而增加了训练时间。如果你是第一次跑建议先不加这些花哨参数跑通整条链路再说。关于批次大小如果显存只有8G--batch-size 16跑640分辨率可能会OOM可以把batch降到8或者把--img降到544。不要为了凑batch疯狂调低分辨率垃圾分拣的小目标会因此丢失得不偿失。7. 最后再分享一些实际部署中的小心得这个项目前后折腾了将近两个月最深的体会是YOLOv5 4.0本身只是一个工具真正的难点全都在工具之外。环境配置看起来繁琐但按部就班来半天时间绝对能搞定数据集制作才是整个项目里最花时间的部分占掉了我六成精力而模型训练反而最简单因为官方代码已经足够成熟剩下的事情就是盯曲线、调阈值。再提一个很多人没注意到的细节4.0版本的--cache参数官方默认表示缓存图片到内存如果内存不足会有警告。但很多教程会把--cache当成默认照抄下来我建议显存紧张时不要开缓存否则训练过程中可能因为内存交换导致整机卡死。我踩过的另一个坑是Windows平台上PyTorch的DataLoader多线程问题。如果num_workers设得过高在Windows下会出现DataLoader worker超时导致训练中断。建议Windows用户把--workers设为0或2Linux下可以设到8。分拣业务流真正落地时视觉检测只是其中一环机械臂的轨迹规划、传送带的启停逻辑、PLC信号的时序配合每一个环节都需要跨团队沟通。作为算法工程师别只顾着调模型多去现场看看机械臂抓空或者抓漏的瞬间很多模型上的问题在现场看一次就明白了。代码和数据可以随时更新但现场的经验不会自动积累。希望这篇内容能让你少踩一些我踩过的坑也欢迎交流你们在垃圾分拣或者类似检测场景里遇到的具体问题。
返回列表
PREV
查看更多资讯
NEXT
返回资讯列表