简介这是一份基于YOLOv5的垃圾分类识别项目包针对可回收垃圾、有害垃圾、厨余垃圾、其他垃圾四类生活垃圾面向人工智能与计算机视觉方向的开发者和在校学生解决物品自动检测与分类问题。项目使用Python语言搭建完整训练与推理流程已在云端环境完成模型训练下载后可用PyCharm直接打开运行。压缩包共1001个文件大小约71.98MB。其中431个txt文件为标注数据231个jpg与209个jpeg为训练及测试图片39个yaml和7个yml为配置参数31个py为训练、推理等脚本另有2个pt格式的预训练权重文件类型覆盖从数据准备到模型部署的关键环节目录结构清晰便于按需查阅。项目已有5561人学习下载。包含可直接运行的推理脚本、测试图片、配置文件和权重将图片放入指定文件夹后运行脚本即可生成识别结果当前可识别瓶子、报纸、电池、剩饭、碎瓷片等典型物品。由于训练数据规模有限模型针对上述样本效果尚可同时保留扩展训练接口有兴趣的读者可自行补充数据集在云端继续训练用于进一步提升模型泛化能力也适合作为YOLOv5目标检测的入门实践参考。1. 为什么要拿 YOLOv5 做垃圾分类识别先认清你面对的是什么问题垃圾分类识别是典型的「类别粒度比目标检测默认设定更细」的场景。你说这是瓶矿泉水瓶检测网络只需要把它从背景里框出来但你要的是「可回收物-塑料瓶」和「可回收物-玻璃瓶」分开甚至同一张图里出现好几个瓶瓶罐罐还要按可回收、有害、厨余、其他四类去归类。这不是模型选型问题而是数据、类别设计、部署形态三者相互拉扯的问题。YOLOv5 到今天依然是这个任务最稳的起点生态成熟、单阶段结构快、模型权重小、社区资料多拿来做垃圾分类的基线几乎是零门槛。但正因为门槛低很多人三个月前跑通了 demo三个月后还卡在同一个地方——类别分不准、定制场景识别率掉得厉害、拿到新瓶型就漏检。这篇文章不打算给你复述 YOLOv5 的原理而是把我做垃圾分类识别时验证过的一整套方案讲清楚从类别怎么定、数据集怎么做到训练参数怎么设、推理阶段怎么调最后是四个最容易翻车的地方和对应的排查路径。2. 把「垃圾分类」翻译成检测任务类别体系、数据来源与标注方案的选型理由2.1 别急着标数据先定类别体系四分类、细分品类还是两阶段方案垃圾分类识别和通用物体检测最大的区别在于类别定义不是模型给的是政策和管理规则给的。国内多数场景要求的是四分类可回收、有害、厨余、其他但四分类直接作为检测类别会出现一个大问题——类别内差异极大。比如「可回收」这个类里既有矿泉水瓶细长又有纸箱方形还有易拉罐矮圆柱模型既要学它们的共性又要区分它们和「其他」里的陶瓷杯、卫生纸。特征维度被拉得太宽模型容量不够准确率会卡在 80% 上下很难再涨。我一般的做法是把任务拆成两个层次但模型还是单阶段。类别体系上做「细分品类」比如塑料瓶、玻璃瓶、易拉罐、纸箱、纸杯、果皮、剩饭、电池、灯泡、织物这些是 YOLOv5 能稳定学出来的物理类别分类逻辑在推理阶段用映射表做二次归并把细分品类映射到四分类。这样做有三个好处第一模型学的是具体物体的视觉特征而不是抽象管理概念学得动第二推理端可以做置信度过滤识别不了的就落到「其他」不会硬着头皮给错答案第三以后管理部门改规则你只需要改映射表不用重训。这里有个取舍要讲清楚。如果你想做的是「扫码式」的逐物品识别类别粒度越细越好甚至可以定到 30 个类以上如果你做的是社区垃圾桶站的整体画面识别需要把一批垃圾一次性归类那就必须用细分品类加映射的方案否则画面里的厨余垃圾和可回收物混在一起时单个类别检测框的置信度全都会下降。实际落地的团队大部分走的是后者我也建议新手先从这个方案入手。2.2 数据集从哪来公开数据集起步自采数据解决「自己的场景」垃圾分类领域有几套公开数据集包含了常见生活垃圾的图片和 VOC 格式标注类别覆盖从塑料、玻璃、纸张到电池、灯管。用它做预训练是划算的但我需要提醒你公开数据集的标注质量参差不齐有些图的边界框明显框大了有些类别标签标错了直接拿来做测试会是灾难。我的习惯是公开数据只做预训练和置信度基线测试部署前必须用自采数据微调。自采数据有几个原则。第一背景要杂垃圾桶站、小区地面、厨房台面、办公室桌面每个场景都要采第二光照要杂白天强光、黄昏逆光、晚上灯光下都要采否则夜间的漏检率会高到让你怀疑模型白训了第三同类物体要收集形态差异塑料瓶包括矿泉水瓶、饮料瓶、洗发水瓶形态差异比你想的大得多只采一种瓶型上去就翻车。采集量上每个细分品类最少 300 张真实图像边界框总数不少于 1500 个这个量级足够微调出一个能看的版本。标注工具方面能支持目标检测标注的常见开源工具都可以导出成 YOLO 格式的 txt 文件最省事也正好是 YOLOv5 原生读入的格式。如果你只有 VOC 格式的数据集也可以用 YOLOv5 仓库自带的格式转换脚本但那个脚本对某些字段的处理比较粗转完我会抽查 50 个文件确认归一化坐标没有越界。2.3 把图片和标注整理成 YOLOv5 目录结构一份可直接跑的准备工作YOLOv5 对数据目录的约定非常死板images 和 labels 必须分开放且训练集、验证集、测试集要独立。我把整理工作写成一个脚本作用是生成标准目录结构并按比例划分数据集。import os import random import shutil from pathlib import Path root Path(garbage_dataset) images_dir root / images labels_dir root / labels # 先扫描所有图片和同名 txt 标注 image_paths list((root / raw_images).glob(*.jpg)) all_items [] for img_path in image_paths: txt_path root / raw_labels / (img_path.stem .txt) if txt_path.exists(): all_items.append((img_path, txt_path)) random.seed(42) random.shuffle(all_items) train_ratio, val_ratio 0.8, 0.15 train_count int(len(all_items) * train_ratio) val_count int(len(all_items) * val_ratio) for split, items in [ (train, all_items[:train_count]), (val, all_items[train_count:train_count val_count]), (test, all_items[train_count val_count:]), ]: split_images images_dir / split split_labels labels_dir / split split_images.mkdir(parentsTrue, exist_okTrue) split_labels.mkdir(parentsTrue, exist_okTrue) for img_path, txt_path in items: shutil.copy(img_path, split_images / img_path.name) shutil.copy(txt_path, split_labels / txt_path.name)这段脚本的含义是把 raw_images 和 raw_labels 两个目录下的图片与同名标注配对随机打乱后按 8:1:1 划分成训练、验证、测试三份。参数方面train_ratio 和 val_ratio 是划分比例我推荐 0.8 和 0.15留 0.1 做最终测试random.seed(42) 固定随机种子确保每次复现划分结果一致排查问题时不会因为数据分布变化而分不清是模型问题还是数据问题。划分完之后还要确认一件事图片和 txt 文件同名txt 里每行是「类别序号 x_center y_center width height」全部是归一化坐标。偶尔会遇到标注工具导出的坐标是绝对值这种情况要在训练前做一次归一化否则模型训练时 loss 直接变成 NaN属于最典型的低级翻车。3. 训练垃圾分类模型的关键参数从过拟合到欠拟合之间找平衡点3.1 选择预训练权重与输入分辨率不要一上来就跑大模型YOLOv5 提供了从 n 到 x 的多个模型尺寸。垃圾分类识别这个任务检测目标通常不是极小目标但类间差异细比如纸杯和塑料杯我的经验是s 和 m 两个尺寸最适合做基线n 太小容易欠拟合l 和 x 在真实部署时推理速度明显下降NVIDIA Jetson 这类边缘设备上根本跑不动。预训练权重建议用 COCO 的 yolov5s.pt。虽然 COCO 没有垃圾分类类别但它学到了通用物体的轮廓和纹理特征迁移到垃圾上比从随机权重开始训练收敛快很多泛化能力也好不少。有人觉得 COCO 预训练和垃圾分类数据分布差异大干脆不用实际做对比之后会发现从随机初始化出发训练同样的 epochmAP 低了大概 8 个百分点速度还慢了一倍。预训练权重从来没有「必须完全匹配任务」的要求只要底层特征复用得上就值得用。输入分辨率这块我推荐 640。理由是垃圾场景里常有多个物体堆叠边界框尺寸分布广640 既能保证小物体特征不丢失又能在常见的消费级显卡上以 batch size 16 顺利训练。如果你只关注小区垃圾桶站这类近距离场景1280 会带来明显的精度提升但训练时间和显存占用几乎翻倍需要在资源和效果之间做取舍。3.2 训练命令与每个参数的实际含义假设你的数据文件是 garbage.yaml内容指定了训练集和验证集路径、类别数量和类别名清单训练命令长这样python train.py \ --img 640 \ --batch 16 \ --epochs 150 \ --data garbage.yaml \ --weights yolov5s.pt \ --hyp data/hyps/hyp.scratch-low.yaml \ --name garbage_s_640逐项说明参数--img 640 是训练时的输入尺寸训练时 YOLOv5 会在这个尺寸基础上做随机缩放增强范围是 0.5 到 1.5 倍--batch 16 是一次迭代喂给模型的图片数这个值受显存限制如果训练时报 CUDA out of memory优先把 batch 降到 8而不是换小模型--epochs 150 是训练轮数垃圾分类这种中等规模数据集150 轮足够如果 100 轮时验证集 mAP 还在明显上升可以接着训到 200 轮--hyp 指向超参数文件YOLOv5 自带的 hyp.scratch-low.yaml 适合小模型微调它把学习率、数据增强幅度都调得比较保守不容易跑飞。训练过程中要盯两个指标训练 loss 和验证 mAP。训练 loss 持续下降但验证 mAP 从某个 epoch 开始不再涨说明模型开始过拟合垃圾分类场景里过拟合的典型表现是「训练集上分得很准换一批新图就漏检」本质上是模型记住了训练图的背景而不是垃圾本体。解决方法是提前终止训练或者把 hyp 文件里的 hsv_h、hsv_s、flipud 等增强参数调大让模型见过更多变化。另一个指标是各类别的单独 APYOLOv5 训练日志会输出每个类的 AP 值你要找出 AP 最低的那几个类它们就是之后需要补充数据的方向。3.3 类别不平衡与难例挖掘为什么你的易拉罐总是被认成塑料瓶生活垃圾数据天然不平衡果皮、塑料瓶、纸箱数量大电池、灯管数量少。如果不做处理模型会把所有注意力放在高频类上低频类的 AP 会惨不忍睹。YOLOv5 本身没有内置类别频率加权机制所以要在数据处理层面处理。我会先统计每个类别的边界框数量找出边界框数小于 500 的类别单独做数据增强具体是三件事对这类图片做随机旋转 90 度和水平翻转增加样本量用小批量过采样让每个 batch 里至少出现一张包含低频类别的图如果低频类还是不够就去真实场景补拍。补拍的时候不要只在理想角度拍要在垃圾混放、物体半遮挡、阴影遮挡的条件下拍这些难例比那些摆拍的数据对模型提升大得多。难例挖掘是另一个容易被忽视的点。第一轮训练完把验证集里所有被漏检、误检的图片收集起来人工确认后单独放在 hard_examples 目录里混入训练集再训一轮。这个操作一般能把低频类的 AP 提升 5 到 10 个点。不要指望数据增强能完全替代真实难例增强只是让模型在已有数据上做插值真实难例才能教会模型处理未见的组合方式。4. 用训练好的模型做识别推理置信度阈值、NMS 与类别映射的落地逻辑4.1 推理脚本的几个关键输出与踩坑点训练完成后YOLOv5 的 detect.py 可以直接用来跑推理但直接它的默认参数不一定适配垃圾分类场景。python detect.py \ --weights runs/train/garbage_s_640/weights/best.pt \ --source ./test_images \ --img 640 \ --conf 0.35 \ --iou 0.45 \ --save-txt \ --save-conf--conf 0.35 是置信度阈值低于这个值的检测框会被丢弃。垃圾分类场景里这个值我一般设成 0.35 到 0.4太高会漏掉被遮挡的物体太低会产生大量误检。--iou 0.45 是 NMS 的 IoU 阈值多个重叠框如果 IoU 超过这个值就只保留分数最高的那个。--save-txt 会把检测结果存成 txt--save-conf 会在 txt 里附上每个框的置信度分析 badcase 时这个文件非常有用。检测结果里每个物体的类别序号是训练时 data.yaml 里定义的下标。整合成垃圾分类四分类时需要一个映射表。比如训练时定了 0 塑料瓶、1 玻璃瓶、2 果皮推理时就要把 0 和 1 归到「可回收」2 归到「厨余」。这个映射逻辑单独写一个函数处理干净不要散落在不同脚本里。CLASS_TO_FOUR { 0: 可回收, # 塑料瓶 1: 可回收, # 玻璃瓶 2: 可回收, # 易拉罐 3: 可回收, # 纸箱 4: 厨余, # 果皮 5: 厨余, # 剩饭 6: 有害, # 电池 7: 有害, # 灯管 8: 其他, # 织物 } def map_to_four(result_detections): four_class_votes {可回收: 0, 有害: 0, 厨余: 0, 其他: 0} for det in result_detections: class_id int(det[class_id]) conf det[confidence] if conf 0.35: four_class_votes[CLASS_TO_FOUR[class_id]] 1 return max(four_class_votes, keyfour_class_votes.get)这段代码看起来简单但里面有两个容易被忽略的设计。第一它在做「投票」而不是「取最高置信度」——因为垃圾分类识别往往面对一堆垃圾不是只有单个物体投票比单框置信度更能反映整体类别第二阈值 0.35 在映射函数里又做了一次过滤避免低置信度检测框对投票结果产生干扰。如果你只取置信度最高的那个框作为最终分类结果遇到画面里同时有果皮和矿泉水瓶时结果通常是被果皮主导但按四分类的桶站管理语义这袋垃圾到底是厨余还是可回收需要重新考虑。4.2 部署时把模型导出为 ONNX分类速度与精度之间的可调空间训练好的 PyTorch 模型直接部署不是不行但推理效率上差一截。导出成 ONNX 再做前向推理速度提升明显尤其对 CPU 部署或边缘设备。YOLOv5 仓库自带导出脚本python export.py \ --weights runs/train/garbage_s_640/weights/best.pt \ --img 640 \ --include onnx \ --simplify--simplify 会调用 onnx-simplifier 对计算图做简化去掉一些冗余算子我建议默认带上。导出完之后可以用 onnxruntime 做一次推理验证确保输出形状和 PyTorch 原模型一致。这里有一个常见的坑导出的 ONNX 模型输出格式和 YOLOv5 原版的 detect.py 解析方式不完全一致需要按输出矩阵的前 4 列是框坐标、第 5 列是物体置信度、后面是各类别置信度的方式解析。如果你直接用 onnxruntime 的输出做后处理大概率会得到一堆乱码坐标。部署阶段最关键的一个参数是置信度阈值。训练阶段你用的是 0.35部署时我一般会根据实际运行环境的误检率调高到 0.45 或 0.5。因为部署场景里背景更杂垃圾桶站的水泥纹理、反光地面、路人身上的图案都会触发误检调高阈值是性价比最高的误检抑制手段。但阈值不是越高越好调太高会把真实垃圾漏掉具体调多少要在目标场景的实拍样本上做验证记录准确率和召回率的平衡点。5. 垃圾分类识别避坑指南四个真实踩过的坑与排查路径5.1 一类物体分裂成多个边界框同一个瓶子被识别成三个框现象模型把同一个矿泉水瓶同时识别出塑料瓶、玻璃瓶、易拉罐三个不同类别的框置信度都不低。原因这种情况下最有可能的是 NMS 的 IoU 阈值过高导致两个类别框因为 IoU 低于阈值而没有被合并而不是模型真的把同一个物体识别成了多个类别。另一种情况是你的训练数据里同一个物体在不同图上被标注成了不同类别标注噪声被模型学进去了。解决先把 --iou 从 0.45 调低到 0.3看多框现象是否减少。如果减少了说明是 NMS 合并问题。如果没减少去检查标注数据把所有训练图里同一个物体的标注类别抽出来看是否存在不同标注员标了不同类别的情况。标注一致性检查在垃圾分类这种细粒度任务里非常关键建议标注前定一份带参考图的标注规范标注后随机抽查 5% 的文件。5.2 夜晚灯光下识别率骤降同样的模型白天 90 分晚上 60 分现象模型在白天光照良好的图片上 mAP 接近 0.9但换上夜间灯光场景后漏检率显著增加尤其是深色物体完全消失。原因这是光照域差问题。训练集里白天的图片占比过高模型学到的颜色分布整体偏亮夜间灯光下色温偏黄且局部过曝模型的特征分布发生偏移。这类问题不是调参能解决的。解决数据采集阶段就要把夜间场景补进去这是最省事的手段。如果已经训练完了补救方案是采集夜间图片做标注用这些图片对模型做增量训练学习率设成正常训练的 0.1 倍epochs 设成 30 到 50。注意增量训练只需要输入夜间数据不需要把白天数据再喂一遍否则白天效果会被稀释。5.3 瓶盖、标签等小物体整个漏检小目标问题在垃圾场景的具体表现现象大件的纸箱、瓶子都能检测出来但瓶盖、烟头、电池这类小物体基本漏检或置信度极低。原因YOLOv5 默认下采样到 80x80、40x40 的特征图小物体在低分辨率特征图上信息基本丢失。训练时的输入尺寸如果偏小小物体特征就更加留不住。解决有两个手段。第一是把训练尺寸从 640 升到 1280直接增加小物体的像素数量代价是训练时间和显存增加第二是调整 anchor 尺寸在模型配置文件里针对小物体增加小 anchor。实际操作中我遇到最多的情况是两者都要做1280 输入加小 anchor小物体的 AP 能从 0.3 提升到 0.65 左右。如果你受限于硬件资源不想改训练尺寸那可以尝试把图片裁剪成几个不重叠的区域分别做检测再把结果合并这是不重训模型的替代方案。5.4 训练 loss 正常但验证 mAP 震荡剧烈学习率与 batch size 的配合问题现象训练 loss 曲线平滑下降但每次验证的 mAP 波动非常大上一轮 0.8下一轮就掉到 0.6再下一轮又涨回 0.78。原因最常见的是学习率设置偏大加上验证集图片数量少导致 mAP 对模型权重微小变化非常敏感。另一个原因是 batch size 太小BN 层的统计量不稳定波动被放大了。解决把 hyp 文件里的 lr0 从默认值调低比如从 0.01 降到 0.005观察 mAP 曲线是否变平滑同时把验证集扩充到 500 张以上减少评估时的随机性。还有一个不太直观但很有效的操作在训练最后 20 个 epoch 启用余弦退火学习率衰减这会让模型在收敛末期更稳定mAP 曲线的震荡幅度会明显变小。6. 让垃圾分类识别模型「越用越准」混淆矩阵与增量训练的实战习惯模型跑通只是第一步真正要投入使用需要建立起「评估、排查、迭代」的闭环。我每次做完一个版本的模型会在测试集上统计出混淆矩阵重点看那些成对出现的错误模式。比如易拉罐和塑料瓶互相混淆说明两者在颜色和形状特征上太接近需要补充更多带包装、带标签的样本纸杯和纸箱混淆大概率是模型学习了纹理而没有学习结构需要补充角度不同的样本。混淆矩阵就是判断模型缺陷在哪里的地图不要只盯整体 mAP。拿到混淆矩阵后我把错误样本分成三类漏检、误检、类别分错。漏检是框没出来看置信度阈值的设置和该类别样本量误检是背景被当成垃圾看阈值和背景数据类别分错是框出来了但标签错看类别间的特征差异和标注一致性。每类错误都有对应的解法不是无脑加数据。数据补完不是重新从头训练而是做增量训练。增量训练的操作是把旧权重作为预训练权重只喂新采集的数据学习率设成正常训练的十分之一训练 30 到 50 个 epoch。这个做法能保留旧模型已经学到的所有能力同时把新样本的信息融合进去不会出现「旧数据全忘」的灾难。我有一次图省事在新旧数据合并后从头训练了 200 轮结果新版模型对旧场景的识别效果反而下滑了后来分析是采样分布变化导致模型过拟合了新场景。从那以后我一直用增量训练的模式稳很多。关于置信度阈值我会给每个类别设独立的置信度阈值而不是全局统一值。查混淆矩阵时经常发现塑料瓶置信度普遍 0.8 以上电池普遍只有 0.5 左右如果统一用 0.5 做阈值误检主要来自塑料瓶漏检主要来自电池。分开设阈值之后两类问题同时缓解。这个操作实现成本很低在推理脚本里加一个 per-class 的阈值字典就行建议你一上来就按这个思路设计。最后再提醒一句垃圾分类识别的难点不在模型结构在你对场景和数据的管理。把类别体系设计好、数据分布看清、迭代闭环跑起来YOLOv5 这个方案能撑住绝大多数真实场景。希望这些踩坑经验能帮你少走一段弯路。本文还有配套的精品资源点击获取