简介:基于YOLOv5的垃圾分类识别项目,面向计算机视觉入门与数据挖掘学习者,解决四类生活垃圾(可回收、有害、厨余、其他)的检测与分类问题。项目使用Python实现,已包含在Google Colab上完成的训练权重,下载解压后可直接运行dect.py进行识别,测试图片放入data/images即可,识别结果自动保存至runs-dect-exp数字文件夹下。资源共1001个文件,以txt标注文件、jpg/jpeg图片、yaml配置、py脚本为主,另有pt模型权重、ipynb示例及Dockerfile环境配置,压缩包大小约71.98MB,结构清晰便于本地调试。已有5558人学习下载,适合想做目标检测实战的入门者。需要说明的是,当前训练数据量较小,模型对瓶子、报纸、电池、剩饭、碎瓷片等特定物体识别效果较好,读者可参考项目流程自行扩充数据集并在云端训练,从而获得更泛化的分类模型。
1. 项目整体设计与思路拆解
做垃圾分类识别这个项目,最早是帮我一个毕业设计的学弟解决课题方向问题。当时他拿来的题目是“基于深度学习的垃圾自动分类系统”,听起来挺大,其实落地场景很清晰:摄像头拍一张垃圾照片,模型告诉你是可回收垃圾、厨余垃圾、有害垃圾还是其他垃圾。选YOLOv5做检测而不是直接用图像分类网络,主要是考虑到真实场景里一张图片可能有多个垃圾、且垃圾往往有重叠遮挡,单纯分类网络只能回答“这是什么”,没法回答“垃圾在哪里”。而检测网络可以同时输出类别和位置框,为后续机械臂抓取或提醒用户“你左手边那个瓶子是可回收垃圾”这类应用留下接口。
为什么选YOLOv5而不是Faster R-CNN或者SSD?我在几个项目里对比过,YOLOv5属于单阶段检测器,速度和精度的平衡在工程落地中确实是最舒服的。Faster R-CNN精度高但推理慢,做实时识别时摄像头画面明显卡顿;SSD速度上去了但小目标检测效果一般,垃圾里的烟头、瓶盖这种小物件很容易漏检。YOLOv5在COCO数据集上的mAP表现不错,而且权重文件只有几十MB,部署到边缘设备或者普通笔记本上都能跑得动。
这个项目的完整链路包括数据采集与标注、模型训练、推理验证和简单部署四个环节。我用的技术栈是Python 3.8 + PyTorch 1.10 + YOLOv5 6.0版本,数据集用的是公开的华为云垃圾分类数据集加上自己补充拍摄的部分图片。如果你也想做类似的识别项目,这个思路可以直接套用,不需要从头造轮子。
1.1 需求拆解与技术选型逻辑
垃圾分类识别本质上是一个目标检测任务,但它和通用的检测任务有区别。核心难点有三块:一是垃圾种类多样,同类垃圾在不同光照、不同角度下差异很大;二是垃圾经常处于半遮挡状态,比如塑料袋里的饮料瓶,只露出一个角;三是背景复杂,桌面、地面、草地等各种环境都会出现,模型容易把背景纹理误判为垃圾。
针对这些难点,我做了几个技术决策。第一,类别体系采用生活垃圾“四分类”标准:可回收垃圾、厨余垃圾、有害垃圾、其他垃圾,这符合国内多数城市的管理要求,也方便后续对接智能垃圾桶的投递口控制逻辑。第二,在YOLOv5的框架内选用yolov5s作为基础模型,权重小、训练快,精度损失在接受范围内。第三,训练时使用Mosaic数据增强,把四张图拼接成一张训练,能显著提升模型对遮挡和小目标的鲁棒性。
数据这块我也说下,公开数据集大概有五千多张图片,我自己又用手机拍了大概八百张,覆盖了办公室、厨房、小区垃圾投放点三个场景。标注工具用的LabelImg,标注格式是YOLO的txt格式,每个标注包含类别id和归一化后的中心点坐标、宽高。类别映射按顺序排好,比如0代表可回收垃圾、1代表厨余垃圾、2代表有害垃圾、3代表其他垃圾。
1.2 YOLOv5网络结构的关键优势
YOLOv5的网络结构相比前代有一个很显著的变化,就是使用了CSPNet作为Backbone的主干结构。CSPNet把特征图分成两部分,一部分经过卷积层计算,另一部分直接通到后面,这样既减少了计算量,又保证了梯度传播的顺畅。简单理解,有点像把一条主干道改成一条直行道加一条绕行景观道,两条路都通到终点,但直行道保证了信息不丢失,绕行道负责提取更丰富的特征。
另外一个值得说的设计是PANet作为Neck部分,它的作用是把不同尺度的特征融合起来。垃圾检测里经常出现的问题是同一个类别在不同距离下大小差异很大,比如远处看是一个小点,近处看占满整个画面。PANet通过自顶向下和自底向上两条路径融合特征,让网络同时具备识别大目标和小目标的能力。我在实际测试中发现,加了PANet之后,对距离镜头两米左右的小矿泉水瓶检测效果提升非常明显。
2. 环境配置与数据集准备
环境配置这块,如果你是第一次玩YOLOv5,我建议直接按照官方仓库的README来,但有几个坑我必须提前说。PyTorch的版本和CUDA版本一定要匹配,我第一次装的时候没注意,配了个CUDA 11.2结果PyTorch装的是1.8,运行时报错“AssertionError: Torch not compiled with CUDA enabled”,浪费了半天时间排查。正确做法是先看显卡驱动支持的CUDA版本,再决定安装对应的PyTorch,用conda创建虚拟环境隔离依赖,别直接装在系统Python里,否则后面装其他项目依赖时容易冲突。
2.1 数据集来源与标注细节
数据集我最终用了两个来源结合。公开部分来自华为云的垃圾分类数据集,包含四十多类垃圾图片,我按照四分类标准重新映射了标签。比如瓶子、纸张、金属归为可回收垃圾,剩饭、果皮归为厨余垃圾,电池、药品归为有害垃圾,其余归为其他垃圾。自己补充拍摄的部分主要是针对公开数据集中数量偏少的类别,比如有害垃圾里的过期化妆品、其他垃圾里的脏纸巾,这些在实际场景中常见但公开数据很缺。
标注的时候有几个细节需要注意。第一,标注框要贴合物体边缘,但不要紧贴到像素级,留出2%到3%的余量对训练效果更好,因为模型预测时本质上是在回归四个坐标值,框太紧容易导致loss震荡。第二,一个图片里多个同类垃圾要分别标注,不要合并成一个框,否则模型不清楚到底有几个目标。第三,对于严重遮挡的目标,如果肉眼都难以判断完整轮廓,我建议直接不标,标注了反而会引入噪声。LabelImg里按W键画框、按D键切换到下一张图、按Ctrl+S保存,这套快捷键熟练之后一天能标三百张左右。
数据集准备好之后,还需要做一次划分。我用了一个小脚本按8:1:1随机分成了训练集、验证集和测试集,并生成了对应的txt文件。YOLOv5的数据集配置文件是一个yaml文件,格式大概是这样:
train: ./data/train.txt val: ./data/val.txt test: ./data/test.txt nc: 4 names: ['recyclable', 'kitchen_waste', 'hazardous', 'other']这里有个容易踩的坑:train.txt里存放的是每张图片的绝对路径,如果你的项目换了一台机器跑,路径变了就会报错“Not Found”。后来我改成了使用相对路径,或者在训练脚本里动态拼接当前项目路径,这样迁移起来省心很多。
2.2 数据预处理与增强策略
YOLOv5在训练时会自动做数据增强,不需要你在外部额外处理。但有几个参数值得关注,在data/hyps/hyp.scratch-low.yaml里可以调整。比如hsv_h负责色调扰动,默认是0.015,我调到了0.02,让模型对颜色变化更鲁棒。为什么这个重要?因为同一个塑料瓶在不同灯光下颜色会偏黄或者偏蓝,如果模型只记住了标准色,实际场景中很容易漏检。
还有一个关键参数是fliplr,默认0.5表示50%概率左右翻转。这个增强对大多数检测任务都有益,因为物体的左右特征天然对称。但如果你检测的目标存在左右不对称的情况(比如有些垃圾袋上有文字),要注意降低这个概率。我当时考虑到四分类体系里会包含带标签的矿泉水瓶,左右翻转后文字变为镜像,模型可能会困惑,所以把fliplr降到了0.3。
Mosaic增强是YOLOv5的默认开启项,它把四张图随机裁剪后拼接成一张新图参与训练。这个策略相当于变相增加了batch size,让模型在单次迭代里看到更多样化的特征。我在训练时对比过开启和关闭Mosaic的效果,关闭后mAP大约下降了两个百分点,特别是对厨余垃圾这种形状不规则的目标影响更大。所以建议新手不要轻易关闭Mosaic,除非你的数据集本身特别干净且目标大小均匀。
3. 模型训练与超参数调优
训练阶段我踩的坑最多,也最值得写出来分享。首先要明确一个认知:YOLOv5的预训练权重非常重要,不要随机初始化训练。官方提供了在COCO数据集上预训练好的权重,虽然COCO里没有垃圾类别,但模型已经学会了通用的边缘、纹理、颜色特征,这些底层特征迁移到垃圾识别任务上同样有效。我做过一个对照实验,使用预训练权重训练50轮后mAP达到82%,而随机初始化训练50轮只有61%,差距巨大。
3.1 训练命令与参数选择实战
训练命令我用的是最基础的形式:
python train.py --data garbage.yaml --cfg models/yolov5s.yaml --weights yolov5s.pt --batch-size 16 --epochs 100 --img 640batch size的选择要考虑到显存大小。我的显卡是RTX 3060 12G,用yolov5s模型、输入尺寸640x640,batch size 16刚好跑满显存不溢出。如果你的显卡只有8G显存,建议batch size降到8,或者把img尺寸降到512,否则会报“CUDA out of memory”。训练轮次我设了100轮,但实际上在50轮左右loss就开始进入平台期,后面50轮提升有限。YOLOv5提供了early stopping机制,如果连续30轮验证集mAP没有提升,训练会自动停止。
学习率方面,YOLOv5默认使用余弦退火调度,初始学习率0.01。这个值对大多数任务都是合适的,不建议新手乱调。我在调参过程中试过把学习率调到0.02,结果训练初期loss就爆了,出现NaN。后来查资料才明白,学习率过大导致梯度更新步子太狠,直接飞出了损失函数的有效区域。如果你遇到loss为NaN,优先把学习率降一个数量级试试。
还有个参数是epochs的粒度问题。YOLOv5每个epoch是指完整遍历一次训练集,100个epoch对于五千张左右的数据集来说,训练时间大约在三到四小时。训练过程中可以顺手打开TensorBoard监控曲线,命令是tensorboard --logdir runs/train,然后浏览器打开localhost:6006。关注两条曲线,一个是训练集loss的下降趋势,另一个是验证集mAP的上升趋势。如果loss下降但mAP不涨,说明过拟合了,需要增加数据增强强度或者引入更多训练数据。
3.2 超参数配置与训练效果分析
我这里直接贴一份我在项目中验证过的关键超参数配置,结合我的实际经验说明每个参数的作用和调整依据:
| 参数名 | 默认值 | 我用的值 | 调整理由 |
|---|---|---|---|
| lr0 | 0.01 | 0.01 | 初始学习率,保持默认即可 |
| lrf | 0.1 | 0.05 | 最终学习率是初始的5%,让收敛更充分 |
| momentum | 0.937 | 0.937 | 动量参数,默认即最优,不用动 |
| weight_decay | 0.0005 | 0.0005 | L2正则化系数,防止过拟合 |
| warmup_epochs | 3.0 | 5.0 | 前5轮用较小学习率热身,稳定训练 |
| hsv_h | 0.015 | 0.02 | 增强色调扰动,适应不同光照环境 |
| hsv_s | 0.7 | 0.8 | 增强饱和度扰动,提高颜色鲁棒性 |
| fliplr | 0.5 | 0.3 | 降低翻转概率,保留瓶身文字特征 |
| degrees | 0.0 | 10.0 | 随机旋转10度,模拟摆放角度变化 |
| translate | 0.1 | 0.2 | 增强平移扰动,应对目标偏移情况 |
特别说明一下warmup_epochs这个参数。它在训练初期让学习率从很小的值线性上升到设定值,目的是防止刚开始训练时权重还在随机状态,过大的学习率直接把模型的初始特征破坏了。我调到5轮后训练明显更稳定,loss曲线初期的锯齿波动也少了很多。
degrees旋转增强对垃圾检测很重要。实际投放垃圾时,瓶子可能是横着、竖着、斜着各种姿态,如果不做旋转增强,模型对倾斜角度大的物体检测效果会变差。但这里要注意,旋转角度不要超过15度,否则标注框和物体之间的对齐关系会被破坏,目标可能旋转出图片边界,反而引入噪声。
训练结束后,可以在runs/train/exp/weights/目录下找到best.pt和last.pt两个权重文件。best.pt是在验证集上mAP最高的那次权重,last.pt是最后一轮训练的权重。实际部署时建议用best.pt,我在项目中用best.pt做测试集推理,mAP达到了87.3%,比last.pt高出将近两个点。
4. 模型推理部署与常见问题排查
训练好模型只是完成了前半程,真正的考验在部署和实际使用环节。我写了一段推理脚本,调用YOLOv5的detect.py来测试单张图片和视频流,命令是这样的:
python detect.py --weights runs/train/exp/weights/best.pt --source ./test_images/ --conf-thres 0.35 --iou-thres 0.45conf-thres是置信度阈值,只有置信度大于0.35的检测结果才会被保留。这个值需要根据实际场景来调,如果调得太高(比如0.7),模型会漏掉很多目标;调得太低(比如0.15),又会出现大量误检。我试下来0.35到0.45之间是比较合理的区间。iou-thres是NMS的IoU阈值,控制两个重叠框合并的严格程度,默认0.45在垃圾场景下表现可以,不需要专门调整。
4.1 实际部署中常见的三类问题
第一类是误检问题,我在测试中发现模型经常把黑色的垃圾袋识别成有害垃圾,甚至把深色的键盘识别成可回收垃圾。排查下来原因是训练数据中深色物体占比过高,数据分布不均衡导致的偏置。解决思路有两个,一是增加负样本,加入大量不含垃圾但背景复杂的图片,让模型学会“没有检测目标时输出为空”;二是调整分类阈值,对不同类别使用不同置信度阈值,比如有害垃圾要求0.5以上才输出,其他垃圾0.3就输出。
第二类是漏检问题,特别是小目标的漏检。我把手机摄像头放在距离地面一米五的高度拍摄,地面上的烟头大约只有30x30像素,模型几乎检测不到。YOLOv5的检测头有三个尺度,分别是80x80、40x40、20x20的特征图。大尺度特征图负责检测小目标,但yolov5s本身对小目标的编码能力有限。我试过两种优化方案:把输入分辨率从640提高到960,小目标的检测率明显上升,但推理速度从30ms增加到55ms;另一个方案是使用yolov5m模型,比s模型多了将近一倍的参数,小目标检测效果好一些,但显存占用和推理时间也上去了。最终在实时性要求不高的场景里,我选择了提高输入分辨率来换取小目标召回率。
第三类是环境泛化问题。在实验室里测试效果很好,但拿到小区垃圾投放点实测时,检测精度明显下降。主要原因是光线变化太大——上午的强光直射和傍晚的昏暗光线,图片的整体亮度差异很大。解决方案是在训练数据的增强参数里加入了亮度扰动和对比度扰动,让模型见过更多“极端”光照条件。改进后再去实测,精度恢复到了可接受的范围。
4.2 优化推理速度的几点经验
如果你的项目需要部署到摄像头实时检测,那么推理速度就是核心指标。我的RTX 3060显卡上,yolov5s的推理速度大约30ms每帧,已经能实现实时检测。但如果你要部署到树莓派或者手机端,有几个优化技巧。一个是模型剪枝,YOLOv5官方提供了网络剪枝脚本,可以把不重要的通道裁剪掉,模型体积减小30%,速度提升20%,精度损失约2%。另一个是TensorRT加速,把PyTorch模型转换成TensorRT的engine格式,在NVIDIA显卡上推理速度可以提升两到三倍。这个转换过程稍微有点折腾,需要在目标机器上安装TensorRT和对应的torch2trt库。
我自己尝试过把模型部署到树莓派4B上,因为没有NVIDIA显卡,只能用NCNN或者ONNX Runtime跑CPU推理。yolov5s在树莓派CPU上推理一张640x640的图片大约需要2.5秒,远远达不到实时,但做拍照识别的场景勉强够用。如果想在边缘设备上达到实时,可以考虑换成yolov5n或者yolov5-tiny,模型更小,速度提升明显,不过精度会有所下降。
4.3 常见问题速查表
| 现象 | 可能原因 | 解决方案 |
|---|---|---|
| 训练时loss为NaN | 学习率过高 | 降低lr0到0.001左右 |
| CUDA out of memory | batch size过大或输入尺寸太大 | 降低batch size或img尺寸 |
| 检测置信度普遍偏低 | 训练不充分或数据量不足 | 增加训练轮次或数据集规模 |
| 小目标完全检测不到 | 输入分辨率太低 | 将img从640提升到960 |
| 训练集loss下降但验证集mAP不涨 | 过拟合 | 增加数据增强强度或引入更多数据 |
| 模型对某些类别完全不识别 | 该类别训练样本太少 | 补充该类别图片或使用类别权重 |
| 检测框偏移不准 | 标注框质量差 | 检查标注是否贴合目标边缘 |
| 部署后精度下降 | 实际场景与训练数据分布不一致 | 收集部署场景数据做微调训练 |
5. 扩展方向与个人实操体会
做完这个项目,我觉得最有价值的不是模型本身,而是整套“数据-训练-部署”的工程链路。顺着这个思路,可以做很多实际产品。比如对接硬件设备,通过串口把识别结果传给STM32控制板,驱动舵机把垃圾投递到对应的分类箱。这个方向已经有毕设项目在做——基于STM32的边缘端YOLOv5车辆检测与车位管理系统,本质上是同一套技术栈,只是检测目标从垃圾换成了车辆。另一个方向是结合机械臂做自动抓取分类,检测输出的坐标信息直接作为机械臂的抓取点,这在智能垃圾分拣中心已经有商业化应用。
最后再分享两个我在实践中的个人体会。第一个是关于数据标注的:不要迷信公开数据集,一定要花时间自己采集一批目标场景的数据并标注。我在项目初期直接用公开数据集训练,以为能省事,结果一到真实场景就暴露了泛化问题。后面补了八百多张自己要部署场景的图片,精度提升不是一个量级的。第二个是关于迭代的:不要指望一次性把模型调到完美,而是要建立一个“收集失败样本 → 补充训练 → 重新测试”的循环。我把每次实测中检测失败或误检的图片都单独存到一个文件夹里,隔几天就用这些难例样本做一次增量训练,模型的鲁棒性在这一轮轮的迭代中逐步提上来的。这个方法对于任何目标检测项目都适用,投入产出比很高。
本文还有配套的精品资源,点击获取