简介:基于YOLOv5的垃圾分类识别项目,面向人工智能与机器学习入门及实战开发者,解决生活场景中可回收垃圾、有害垃圾、厨余垃圾和其他垃圾的图像检测问题。压缩包约71.98MB,共1001个文件,包含31个Python脚本(如dect.py)、39个yaml配置、231张jpg与209张jpeg样本图片、431个txt标注文件,以及训练好的pt权重;文件覆盖模型配置、数据标注、推理脚本与权重,解压后用PyCharm打开即可运行。使用时将测试图片放入data/images目录,运行dect.py即可完成识别,结果自动保存至runs/dect/exp+数字文件夹,便于逐一查看。由于训练数据量有限,当前模型主要识别瓶子、报纸、电池、剩饭、碎瓷片等典型物品,适合入门YOLOv5训练与推理流程,也可作为垃圾分类演示或课程设计参考。目前已有5558人学习下载,对于希望快速上手目标检测项目的开发者,这份完整代码与配置能带来直接借鉴价值。
1. 项目概述与整体思路
做垃圾分类识别这事儿,说起来其实已经不算新鲜了,但真正动手把它完整落地,从数据集标注一直走到边缘设备部署,这一整条链路能踩通的并不多。这次的项目我选的是YOLOv5,不是因为它最新,恰恰是因为它足够成熟、生态完善、资料多、坑也基本都被前人踩平了,特别适合做毕设或者个人实战练手。
垃圾分类识别本质上就是一个目标检测任务:给定一张图像,模型要判断出图像里的垃圾属于哪一类,并且用边界框标出它的位置。与图像分类不同,目标检测不仅要回答"这是什么",还要回答"它在哪儿",因此YOLO这类单阶段检测器在速度和精度之间取得了很好的平衡,非常适合实时性要求高的场景。
我这次做垃圾分类,最初的设想并不仅仅是在电脑上跑个demo——那太没挑战了。我更想做的是一个能真正部署到边缘设备上的完整方案。所以项目拆分成三个阶段:第一阶段,整理并标注垃圾分类数据集;第二阶段,用YOLOv5训练自己的检测模型并调优;第三阶段,把训练好的模型部署到边缘设备上,验证它在真实硬件上的推理表现。从实际效果来看,这套路对毕设、对想系统入门目标检测的开发者都非常适用。
先说明一下,我这里用的技术栈是基于纯PyTorch的YOLOv5框架,但项目思维和流程可以迁移到YOLOv8、YOLOv11或者任何基于anchor的检测器上。模型选型时我对比过当时各版本的差异,最终选择YOLOv5s作为主力训练模型:它的体积大概14MB左右,在GPU上能有毫秒级的推理速度,在树莓派这类低功耗设备上做优化后也能跑到可用的帧率,刚好卡在精度和性能的甜点区。
数据层面,主流公开数据集包括华为云的垃圾分类数据集、TrashNet等,但公开数据集普遍存在类别不均衡、背景单一、样本量不足的问题。因此我采用"公开数据集+自采图片扩充"的策略,最终构建了一个包含6大类(可回收垃圾、厨余垃圾、有害垃圾、其他垃圾,以及更细分的塑料瓶、纸箱等子类)的自有数据集。这里要特别提醒一句:不要盲目追求类别数量,如果你的目标是毕设展示或功能验证,先把4大类的核心类别做好做精,远比凑一堆垃圾类别但每个类别只有几十张图要强得多。
2. 数据集准备与标注细节
2.1 数据来源与类别设计
垃圾分类这个任务和通用目标检测有个显著区别:垃圾种类极其庞杂,而且同类垃圾在不同光照、角度、破损程度下的外观差异很大。比如一个被压扁的易拉罐和一个完整的易拉罐,虽然同属一个类别,但视觉特征差异极大。这就对数据集的多样性和标注质量提出了更高的要求。
我在构建数据集时,将类别设计为六大类,兼顾了实际场景和标注成本:
- 可回收垃圾(纸板、塑料瓶、玻璃瓶、易拉罐、衣物)
- 厨余垃圾(剩饭、果皮、菜叶)
- 有害垃圾(电池、灯管、药品)
- 其他垃圾(烟蒂、陶瓷碎片、卫生纸)
- 塑料瓶(单独作为子类,因为它在实际场景中出现频率最高)
- 纸箱(单独作为子类,常出现的快递包装)
类别数量这里有个权衡:类别太少,模型学不到区分能力;类别太多,标注工作量成倍增加,而且某些类别容易混淆(比如玻璃瓶和陶瓷碎片,外观上确实接近)。我最终的经验值是:初期控制在6-8个类别以内最合适,后续可以增量添加。
2.2 标注工具与标注规范
数据集使用的是YOLO格式,每个图像对应一个同名的txt文件,每行格式为:类别ID、归一化中心点x坐标、归一化中心点y坐标、归一化宽度、归一化高度。框坐标是相对于图像宽高的比例值,范围在0到1之间。
标注工具我推荐使用LabelImg,它对YOLO格式的支持很成熟,操作简单,不需要什么学习成本。LabelImg虽然界面朴素,但胜在稳定、导出格式标准。如果你在MacOS或Linux上遇到兼容问题,也可以换用labelme或X-AnyLabeling这类工具,导出时记得统一转为YOLO格式。
标注有一个行业通用规范:框要尽量贴紧目标边缘,但也不需要精确到像素级。目标检测对标注框的宽容度其实还不错,框稍微大1-2个像素不会对训练产生明显影响,真正致命的是漏标和错标。漏标会导致模型漏检,错标会污染训练数据,这些才是需要花大力气避免的。
我踩过的一个坑是:一开始把"压扁的易拉罐"和"完整的易拉罐"当成两类分别标注,结果训练时这两类反复互相误检。后来想明白了,模型本来就难以区分同一物体的不同形变状态,强制让模型学这种细粒度区分,数据量又不够,只会导致掉点。合并类别后,问题立刻消失。所以在这里给大家一个判断标准:如果你自己看图片都要犹豫三秒才知道属于哪一类,那这个类别的定义就是有问题的。
2.3 数据清洗与增强策略
数据清洗这一步非常关键,不过也是最容易被新手忽略的。我在收集完原始图片后,做了三件事:去重(用MD5校验值,防止网络图片重复下载导致数据冗余)、过滤(去掉严重模糊、过度曝光、目标占比过小的图片)、质量检查(随机抽检图片,确认标注框和类别标签没有明显错误)。
清洗完成后,我用脚本统计了每个类别的样本数量。统计结果中,厨余垃圾类样本有一千两百多张,而有害垃圾只有三百多张,其他垃圾也只有五百张左右。这种不均衡会直接导致模型对少数类学习不足、漏检率上升。
解决类别不均衡,我的方案是分层抽样加数据增强。先用离线增强工具(如imgaug、Albumentations)对样本较少的类别做增强,包括:水平翻转、随机旋转(±30度)、亮度对比度扰动、HSV色域变换、随机裁剪(注意裁剪的比例不要让目标太小)、高斯噪声。增强后每个类别都补充到八百张左右,整体数据集大约五千两百张图片。
这里有个要领:增强不是越多越好,而是要将操作控制在"不改变目标语义"的范围内。比如水平翻转是安全的,但垂直翻转对某些类别会反直觉(比如正常场景下垃圾不会从天而降);随机旋转角度太大也会让模型学到错误的空间关系。
3. YOLOv5环境搭建与训练参数调优
3.1 环境配置与常见坑位
YOLOv5的环境配置一直有个口碑:相比YOLOv8和后来的版本,它更依赖"恰好正确的依赖版本组合"。这也成了不少新手被劝退的重灾区。我实际搭建的环境如下,这套组合在Ubuntu 20.04和Windows 11上都验证过:
- Python 3.8(不要用3.10及以上,部分依赖会出现兼容性问题)
- PyTorch 1.12.1 + CUDA 11.3 + cuDNN 8.2
- 其他依赖:opencv-python、numpy、pandas、tqdm、pyyaml、matplotlib
安装YOLOv5的代码相对简单,克隆仓库后安装requirements.txt就好:
git clone https://github.com/ultralytics/yolov5 cd yolov5 pip install -r requirements.txt但我必须要指出requirements.txt里最大的隐患:它会安装最新版的opencv-python和numpy,这常常与已有的环境产生依赖冲突。稳妥的做法是先不装requirements.txt,而是手动安装核心依赖:torch、opencv、numpy、tqdm、pyyaml、matplotlib,其他等运行时缺什么再补什么。
还有一个非常实际的坑:如果你只有CPU环境(比如只有MacBook或普通笔记本),训练YOLOv5虽然慢但并非不可行。问题是不要天真地直接用默认的YOLOv5s模型,而是应该用小模型、降低输入尺寸、调小batch size。我在纯CPU的机器上用YOLOv5s训练了几次,一次一百轮大约要跑大半天,这个时间成本对快速迭代是很不利的。
3.2 训练参数配置与超参数选择
YOLOv5的训练参数大多数可以在训练命令中直接指定,不需要修改代码。我实际使用的训练命令如下:
python train.py \ --data data/garbage.yaml \ --weights yolov5s.pt \ --img 640 \ --batch 16 \ --epochs 100 \ --workers 4 \ --device 0这里逐个解释参数的意义以及我的调优心得:
- --data:数据集配置文件,YAML格式,里面指定训练集/验证集图片路径和类别名称。这是最容易配错的地方,路径一定要用绝对路径或者相对repo根目录的路径,我见过好几人在这里折腾一晚上。
- --weights:预训练权重。YOLOv5最强大的优势之一是它的迁移学习能力,用COCO预训练权重在自定义数据集上微调,效果远好于从零训练。即便你的类别和COCO完全不同也没关系,模型照样能迁移底层的纹理、边缘、形状特征。如果没有GPU可以把--weights设为空,从零训练,但效果会明显差一些。
- --img:训练输入尺寸。默认640,如果你的目标物体比较小,可以考虑768或896,但要付出更多显存和训练时间。垃圾分类的目标一般体积不小,640完全够用。
- --batch:批大小。显存不够时优先调这个参数。8G显存跑YOLOv5s,batch 16是上限,再大会OOM(显存溢出)。
- --epochs:训练轮数。不是越多越好——我在一百轮验证时,看到模型在前50轮就已经收敛,后面50轮几乎是在震荡或者过拟合边缘徘徊。建议一开始用100轮观察训练曲线,如果早停权值(best.pt)出现在前一半,可以果断把epochs降到60-80,节省时间。
超参数层面,YOLOv5的hyperparameters.yml里定义了一堆参数,新手完全不用动,用默认值就好。如果你想微调,我验证过的直觉是:重点看lr0(初始学习率)和mosaic(马赛克增强概率)。默认lr0=0.01,如果你的数据集较小,可以适当降到0.005,防止前期发散。mosaic默认1.0,对小数据集效果很好,但如果你发现某些小目标漏检,可以考虑适当降低mosaic,因为马赛克增强会让小目标更难被捕捉。
3.3 训练过程中的关键观察
训练日志里的loss曲线:训练loss会持续下降,val loss一开始下降,后面可能上升,这是过拟合的信号。这时需要做的是:增加数据增强、增加数据量、或者引入早停。YOLOv5内置了早停机制,patience默认100,意思是一百轮内验证集指标没有提升就自动停。
我这次训练大概用时情况:单张NVIDIA RTX 3060(12G显存)训练一百轮,每轮约30-40秒,总时长约一小时。对于只有一块入门显卡或者云端租赁GPU的朋友,这个成本完全可以接受。
训练完成后会在runs/train/expN目录下生成多个权重文件:last.pt(最后一轮的权重,一般不用)、best.pt(验证集指标最优的权重,用于后续部署)、以及一系列曲线图和混淆矩阵。我做模型选型时,重点看的是混淆矩阵和PR曲线:混淆矩阵的每个格子表示真实类别与预测类别的关系,如果某些类别互相混淆严重,说明类别定义或者数据标注需要调整;PR曲线下的面积(AP)综合衡量精度和召回率,每个类别的AP值差异较大时,需要针对性地补充数据。
4. 模型评估与常见问题排查
4.1 评估指标与实际效果
训练结束后,用验证集做最终评估。我这次得到的整体mAP@0.5大约在0.923,mAP@0.5:0.95大约0.685,单张图片平均推理时间约6毫秒(RTX 3060)。这个水平的模型在垃圾分类场景已经具备实用价值,但要直接放到生产环境中还差一些,主要是对未见过的复杂背景泛化能力仍需提升。
模型测试时,我还做了一件很多人不做的事情——用手机拍了几张真实垃圾图片来测试,而不是只依赖测试集。这一步非常推荐,因为测试集与训练集同源,模型见过类似的背景,测试分数会偏高。真实手机拍摄的图片背景更杂、光照更乱,是检验模型鲁棒性的硬指标。实测下来,模型在独立拍摄图片上的表现比验证集mAP低了不少,尤其是厨余垃圾类,因为真实场景中厨余垃圾往往形状不规则,与背景对比度低。
评估结果出来后,我做的优化工作有三块:第一,针对漏检较多的类别增加样本量,重新标注了一批真实拍摄的图片;第二,调整了图像分辨率从640提升到768,在可接受的性能代价内显著提升了小目标的检出率;第三,使用了测试时增强(TTA)——推理时对图像做多尺度缩放、水平翻转等增强,然后融合结果,虽然推理速度慢了约两倍,但精度提升显著,适合对实时性要求不高的离线场景。一个实用经验是:如果你的设备算力有限,优先用多尺度推理,而不是盲目加大模型体积。
4.2 模型推理与可视化验证
完成训练后,用下面的命令做单张图片或者视频流的推理验证:
python detect.py \ --weights runs/train/expN/weights/best.pt \ --source data/images/test.jpg \ --conf-thres 0.25 \ --iou-thres 0.45--conf-thres是置信度阈值,默认0.25,如果模型误检多就调高,漏检多就调低;--iou-thres是NMS时的IoU阈值,控制重叠框的抑制程度。实际调参时,我发现对垃圾分类场景,conf-thres设0.3比较平衡,太低的话模型会把一些纹理复杂的背景误认为垃圾,太高的话又会漏掉一些远处的小目标。
4.3 常见问题排查清单
我在这个项目以及帮朋友调试的过程中,整理了下面这份排查清单,基本覆盖了YOLOv5训练和推理中的绝大部分痛点:
| 问题现象 | 可能原因 | 解决思路 |
|---|---|---|
| 训练时loss不下降 | 学习率过高/数据集标注错误 | 调低lr0、随机抽检标注文件 |
| 模型预测结果全部为同一类 | 类别标签索引错位(标签和yaml的顺序不一致) | 检查labels的类别ID和data yaml的names顺序 |
| mAP很高但实际图片检测效果差 | 过拟合,模型记住了训练集背景 | 加强数据增强、增加真实场景图片、降低epochs |
| 显存不足(OOM) | batch size过大/输入尺寸大/模型过大 | 调小batch、调小--img、换yolov5s为yolov5n |
| 训练时CPU占用率爆满 | --workers设置过高 | 调低workers数,如--workers 0或2 |
| 推理时检测框大量重叠 | NMS阈值过低 | 调高--iou-thres,如0.5 |
| 小目标距离远但漏检 | 输入分辨率低/mosaic增强过强 | 调大--img、降低mosaic的超参数 |
其中"labels类别ID与yaml的names顺序不一致"是我见过最多的低级错误。标注工具默认从0开始计数,但如果你在yaml里写的names顺序和标注时的类别顺序不一致(比如标注时0对应可回收垃圾,yaml里0却写了厨余垃圾),模型训练不会报错,但推理结果会全部对不上。这个坑一定要在训练前预先检查,我的办法是训练后立刻拿一两张标注过的图片做一次可视化推理,人工比对类别是否正确。
5. 边缘端部署与模型优化
5.1 模型轻量化处理
训练好的PyTorch模型不能直接放到边缘设备上跑,需要经过转换和优化。整个部署链路我整理为四步:先把PyTorch模型导出为ONNX,再用ONNX Runtime或TensorRT做推理加速,同时可以配合INT8量化进一步压缩模型体积,最后在目标设备上做实测调优。
# 导出ONNX格式 python export.py \ --weights runs/train/expN/weights/best.pt \ --include onnx \ --opset 12ONNX是个中间格式,相当于把PyTorch模型"翻译"成一种硬件厂商都认识的通用语言。但翻译后的模型还不是最优状态,在NVIDIA平台上一般还会进一步编译成TensorRT引擎,推理速度能提升2-4倍;在树莓派这类ARM平台上,则可以使用NCNN或RKNN以达到类似的效果。
量化这块,YOLOv5的代码里已经集成了INT8量化的入口。用少量校准图片做动态范围统计,可以将模型权重从FP32压到INT8,体积缩小为原来的四分之一,推理速度也相应提升。但量化不是没有代价的,本项目中量化后的模型mAP大概掉了3.5个百分点,在垃圾类别外观差异大的类别上(尤其是可回收垃圾里形态各异的物件)掉点更明显。所以在实际项目中,我的经验是优先尝试TensorRT的FP16精度,速度和INT8差不多,但精度损失几乎可以忽略。
5.2 树莓派部署实测
作为部署验证,我在树莓派5上做了完整的实测。树莓派5的算力比4代强了不少,在CPU上运行优化后的ONNX模型,输入分辨率640,FP32精度下推理速度大概在2-3 FPS;换用NCNN的FP16优化后能跑到5-6 FPS。对于垃圾分类这种低频场景(人把垃圾拿到摄像头前,等一两秒出结果是完全可接受的),这个速度已经具备可用性。
树莓派部署的具体流程是:在PC上导出为ONNX格式,拷贝到树莓派,安装onnxruntime(ARM版本),然后用Python调用推理接口。代码不算复杂,核心是一个封装了预处理、推理、后处理的类。预处理部分注意把图片的BGR转RGB、归一化到0-1、再转换成NCHW格式;后处理则是把输出的三个特征层解码成边界框和类别。
这个部署还带了一个小功能:把树莓派接上摄像头模块,用OpenCV捕获视频流,当检测盒置信度超过阈值时,控制一个舵机转向对应的垃圾桶格口。这就从一个单纯的识别模型升级成了有实际交互能力的系统。
如果你想把端侧算力再压低一档,STM32这类MCU上也可以运行极轻量化的模型,比如YOLOv5n配合更小的输入尺寸(如320x320),以及MCU上的推理框架(如STM32Cube.AI)。这方向一般用于车辆检测、车位管理等对实时性要求较高的场景,但垃圾分类场景如果要用MCU,通常还需要配合外部摄像头模块,常规开发板加上摄像头模块也是可行的路径。优化思路是相通的:压缩输入尺寸、降低模型通道数、使用INT8量化。
5.3 部署中的若干注意细节
部署过程中有几个特别容易踩的坑值得单独提示。第一个是颜色通道顺序:OpenCV默认是BGR,而训练时YOLOv5用的是RGB格式,如果不做转换,部署后推理结果会严重劣化——这个坑的特征是模型看起来没坏,但所有检测框和置信度都变得十分离谱。第二个是归一化方式:YOLOv5训练时将像素值除以255做归一化,并且没有减均值除方差的操作,部署时一定要复现这个过程。第三个是输入尺寸:训练时YOLOv5会自动做letterbox(保持宽高比的缩放加灰边填充),部署时如果不做letterbox而是直接resize,会降低检测精度。
我在树莓派上调试时还发现一个现象:同一份ONNX模型,用onnxruntime直接跑的结果和用OpenCV的dnn模块跑的结果有细微差异。这不是bug,而是两个框架对某些算子在浮点运算顺序或实现细节上不一致导致的。如果你的项目对精度一致性要求很高,建议部署和验证都用同一个推理框架。
6. 最后的经验小结
绕了一圈,从数据集标注、模型训练、参数调优再到边缘部署,这套流程走下来最深的感受是:目标检测项目真正的工程量不在模型结构,而在于数据质量、工程细节和部署调试。YOLOv5作为一套成熟框架把模型部分封装得很好,但数据和部署这两块硬骨头,没有任何框架能替你啃。
如果你打算以这个方向做毕设,我建议在完成基本识别功能后,往两个角度加分:一是加一个GUI界面,让识别结果直观展示;二是把边缘部署的实际录屏和性能数据作为亮点呈现给老师看——这比只交一个训练好的模型要打动人多得多。
准备收尾了,分享一个我在这个项目里最得意的小优化:推理时将多个摄像头采集的待检测图片攒成batch一次跑。树莓派上单张推理耗时约200毫秒,但四张一起推理只需要约500毫秒,平均每张125毫秒,吞吐量提升了将近四成。这类小技巧不像换大模型那样引人注意,但在资源受限的嵌入式设备上,往往才是决定系统实战可用性的关键一步。
本文还有配套的精品资源,点击获取