简介这是一套面向图像分类初学者的智能垃圾分类完整工程基于TensorFlow与OpenCV提供自定义DIY网络模型覆盖数据集制作、训练、预测全流程。资源共1046个文件包含1041张分类图像、两个Python脚本train.py负责训练、predict.py负责预测、预训练模型my_model.h5及演示视频mp4压缩包大小824.68MB数据与代码结构清晰。已有1330人学习使用。读者可下载后直接运行训练或预测脚本图像上以中文显示干垃圾、湿垃圾、可回收垃圾、有害垃圾四类结果也可将其作为通用图像分类模板快速迁移至其他分类任务非常适合毕业设计、课程实验及算法入门。1. 从一堆湿垃圾里分清干湿先动手的还是图像处理垃圾分类这个题目搁十年前是环境工程的事搁今天它首先是个图像分类问题。社区垃圾桶前要分清楚塑料瓶、纸箱、玻璃罐食堂后厨要把剩饭剩菜和餐盒分开产线传送带上的物料要按材质归位——这些场景的输入都是摄像头拍下来的图像输出是类别标签。但真正做过这类项目的人都知道卡住准确率的往往不是分类网络本身而是最前面那一段图像处理做得够不够扎实。光照不均匀、反光、遮挡、物体堆叠、传送带速度不一这些干扰不处理掉模型参数调得再好也白搭。这篇博客按我实际写代码的顺序来讲先把图像处理当成一个独立的工程环节来设计再解决数据集和标注然后落到分类模型的选型与训练最后给出部署和排错时的关键参数。适合已经能跑通一个分类 Demo、想做成能顶住现场环境的完整方案的工程师阅读。本文不追求刷榜式精度目标是让系统在现场换了个光线条件之后准确率不至于断崖下跌。2. 图像处理在垃圾分类里的三个落点增强、去干扰、定位分割图中分类的输入质量直接决定上限图像处理是把质量下限抬起来的唯一手段。在垃圾分类这个具体问题上图像处理要解决三件事一是把不同光照、不同角度下拍出来的图拉回到一个相对统一的分布上二是把传送带、垃圾桶边缘、人手这些背景干扰尽可能地压下去三是把堆叠的物体从空间上切开让分类网络每次只面对一个目标。2.1 垃圾分类场景下的图像采集制约帧率、曝光与运动模糊先看输入源头。现场摄像头大多数是海康或者大华的 RTSP 流帧率 25fps分辨率 1080p这个参数组合在垃圾分类现场有两个天然问题。第一产线传送带速度到 0.5m/s 之后单帧曝光时间超过 5ms 就会产生明显的运动模糊塑料瓶身上的纹理和标签文字会糊成一团分类网络拿到的特征本身就是脏的。第二垃圾桶投放口的光照在一天内变化剧烈早晨和傍晚的色温差能达到 3000K 以上同一个纸箱在两种色温下拍出来的颜色特征会漂移。所以我一般会在编码端就做一次参数锁定而不是把原始帧丢给后续的预处理。下面这组参数是海康 SDK 里的常用设置适用于多数 USB 摄像头和网络摄像机# camera_config.py import cv2 cap cv2.VideoCapture(rtsp://192.168.1.64:554/stream1) cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1920) cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080) cap.set(cv2.CAP_PROP_FPS, 25) cap.set(cv2.CAP_PROP_AUTO_EXPOSURE, 0) # 关闭自动曝光 cap.set(cv2.CAP_PROP_EXPOSURE, 150) # 手动曝光单位随驱动变化 cap.set(cv2.CAP_PROP_WHITE_BALANCE_BLUE_U, 120) # 锁定白平衡避免色温漂移 ret, frame cap.read()代码的逻辑是关闭相机的自动曝光和自动白平衡用一个固定曝光值 固定白平衡值去采集。这样做的好处是训练集和推理时的图像分布一致坏处是环境光照变化后需要人工重调。折中方案是每 5 分钟取一次画面统计亮度均值超过阈值就触发一次曝光参数重设。2.2 预处理 Pipeline降噪、色彩归一化与增强的推荐顺序图像处理的操作顺序是有讲究的。我见过不少项目先把图缩到 224×224 再做增强这个顺序在垃圾分类场景下是错的。缩图本身会引入高频信息的丢失小物体比如瓶盖上的纹理在缩放过程中就可能被直接抹掉再增强也补不回来。正确的顺序是先做几何校正和色彩处理再做归一化最后才缩放。具体 pipeline 我一般写成下面这样# preprocess.py import cv2 import numpy as np def preprocess_frame(frame, target_size(256, 256)): # 1. 轻度高斯滤波只去掉传感器噪点保留边缘 frame cv2.GaussianBlur(frame, (3, 3), 0) # 2. 转 LAB 空间对 L 通道做 CLAHE避免色彩失真 lab cv2.cvtColor(frame, cv2.COLOR_BGR2LAB) l, a, b cv2.split(lab) clahe cv2.createCLAHE(clipLimit2.0, tileGridSize(8, 8)) l clahe.apply(l) lab cv2.merge([l, a, b]) frame cv2.cvtColor(lab, cv2.COLOR_LAB2BGR) # 3. 等比缩放后居中填充避免物体拉伸变形 h, w frame.shape[:2] scale target_size[0] / max(h, w) new_w, new_h int(w * scale), int(h * scale) resized cv2.resize(frame, (new_w, new_h)) canvas np.zeros((target_size[0], target_size[1], 3), dtypenp.uint8) y_off (target_size[0] - new_h) // 2 x_off (target_size[1] - new_w) // 2 canvas[y_off:y_offnew_h, x_off:x_offnew_w] resized return canvasCLAHE 在垃圾分类里的作用被很多人低估。它把图像的局部对比度拉均匀对塑料瓶的反光区域和纸箱的暗角区域都有明显的修正效果。clipLimit 默认是 2.0但塑料类物体表面高光重的情况下可以调到 1.5参数调低之后增强力度降下来高光区域不容易出现过曝。2.3 对分类任务直接有帮助的分割与定位图像处理如果场景是桶口投放物体进桶之后是堆叠状态不做分割直接送分类网络网络会把两个物体的特征混合在一起。常见做法是用形态学处理做前景分离先对背景差分后的二值图做开运算去粘连再做连通域分析取每个物体的外接框。这个处理对硬件几乎没有额外负担纯 CPU 也能跑实时。以下是针对传送带场景的前景分割最小实现# segmentation.py import cv2 import numpy as np def extract_foreground(frame_bg_diff, kernel_size(5, 5)): # 开运算先腐蚀再膨胀切断细小的粘连 kernel cv2.getStructuringElement(cv2.MORPH_RECT, kernel_size) opened cv2.morphologyEx(frame_bg_diff, cv2.MORPH_OPEN, kernel) # 用连通域过滤掉面积过小的噪声区域 num_labels, labels, stats, centroids cv2.connectedComponentsWithStats(opened, 8) boxes [] for i in range(1, num_labels): area stats[i, cv2.CC_STAT_AREA] if area 1500: # 面积阈值随摄像头距离调整 continue x, y, w, h stats[i, cv2.CC_STAT_LEFT], stats[i, cv2.CC_STAT_TOP], \ stats[i, cv2.CC_STAT_WIDTH], stats[i, cv2.CC_STAT_HEIGHT] boxes.append((x, y, w, h)) return boxes形态学处理的核心是核大小和面积阈值。核太小切不断物体之间的黏连核太大会把小的物体比如瓶盖直接腐蚀没了。先从小核开始如果检测出的框把两个物体框在一起逐步加大核的尺寸。这套流程用到的是 OpenCV 的基础图像处理算法本身不涉及任何训练但能让后面分类任务的数据干净很多。3. 垃圾分类数据集的获取与类目体系设计直接决定模型上限图像处理把图修好了接下来面对的问题是拿什么数据训练分类模型。垃圾分类项目里数据集的坑比模型选型更大。3.1 可用的公开数据集与自采数据的比例公开数据集方面华为云 2020 年开放的垃圾分类数据集是最常被用作起点的包含 40 余类常见生活垃圾每类 200-500 张不等。但它的局限是每一类都是干净背景下的单物体图没有堆叠场景、没有现场光照变化。这意味着用它训出来的模型在实验室测试集上能到 95% 以上的准确率一搬到现场就可能掉到 80% 以下。我的做法是公开数据只用来做预训练或者冷启动必须自采至少 3000 张现场图像再微调。自采图像要覆盖不同的投放高度、不同的垃圾桶填充度、不同时段的光照。分类模型的性能对训练分布和推理分布的差异极其敏感自采数据的价值是公开数据无法替代的。3.2 类目怎么定按材质分还是按投放物分类目体系是垃圾分类图像分类里最容易出错的设计决策。很多项目把类目定为“塑料瓶、易拉罐、纸箱、玻璃瓶、果皮、电池”这看起来合理但落到标注环节就会出现大量模糊样本。比如“饮料瓶”和“塑料瓶”是不是一类“外卖盒”算塑料还是算其他垃圾经验是按材质分大类按物品分小类但小类之间必须有肉眼可辨的视觉差异。下面是一套经过验证的类目表适合产线或者桶口识别场景大类小类举例视觉特征要点容易混淆的类可回收塑料类塑料瓶、塑料盆、泡沫箱反光、半透明、质地硬玻璃瓶可回收纸质类纸箱、报纸、纸杯表面粗糙、颜色偏黄褐布料可回收金属类易拉罐、金属罐高反光、有环状纹理玻璃瓶厨余垃圾剩饭、果皮、菜叶颜色湿润、无规则形状被打湿的纸有害垃圾电池、灯管、药品形状规则、常有印刷图标普通圆柱物其他垃圾用过的纸巾、尘土特征弱化、边缘不清所有类别类目定完后建议对每一类写一份标注规范文档明确边界案例怎么标。比如“沾了油的纸箱”应该算可回收还是其他垃圾如果现场答案是按回收价值来那沾油的纸箱在图像上确实很难和普通纸箱区分需要给模型一个优先级策略这在训练时就要通过在损失函数里加权体现。3.3 标注工具选择与清洗策略先判错再训练标注工具我常用的是 LabelImg 和 X-AnyLabeling前者适合做矩形框标注后者适合做多边形分割标注。分类任务只需要给整张图打一个类别标签用 LabelImg 就够了。但要注意垃圾分类场景下同一张图里多个物体属于不同类别的情况很常见这种情况下整图一个标签会让模型学到一个混合特征。处理办法是先用第 2 章的前景分割把每个物体裁成独立的小图再对每个小图单独分类标注。这一步可以半自动完成分割出来的图先用预训练模型打上预测标签人工只需要修正错标标注速度能提升 3 倍左右。标注完之后必须做一轮清洗——把所有标注置信度低或者不同标注员意见不一致的样本集中检查这些样本往往就是后面训练时 loss 降不下去的来源。我一般还会检查是否有“背景污染”样本就是某类物体的图里背景区域占了 60% 以上这类图会让模型学到背景而非物体本身直接删掉或者裁掉背景再标。这个动作看起来简单但在垃圾分类项目里对准确率的提升比换一个更大的模型更明显。4. 选 CNN 还是 Transformer垃圾分类分类模型的训练参数与损失函数设计数据准备好了模型选型就是下一个绕不开的问题。图像分类算法发展到今天主流两派是 CNN 和 Transformer。对垃圾分类这种类别间差异不算特别精细、但对推理延迟有要求的场景选型有它自己的逻辑。4.1 ResNet 作为基线为什么 18 或 34 层够用对于现场部署推理来说ResNet18 通常是我的首选基线。垃圾分类的类别间差异主要是颜色、形状、纹理、反光特性这些表层特征不要求模型具备像医学影像那样精细的语义理解能力。ResNet18 参数量约 11M单张 224×224 图像在 Jetson Nano 上推理约 20ms准确率已经能到 90% 以上。ResNet34 参数量 21M准确率比 18 提升 1-1.5 个百分点延迟增加约 8ms。是否上 34 取决于现场是 CPU 还是 GPUCPU 环境下我基本只用 18。ResNet 的残差结构价值在于它能让梯度跨层传播意味着即使是小模型也能训得深而不会梯度消失。实际微调时加载 ImageNet 预训练权重比从头训练节省至少 4 倍时间而且小数据集的收敛稳定性好得多。垃圾分类数据集往往只有几千张从头训容易过拟合预训练权重是必须的。4.2 训练脚本冻结骨干、微调全连接层最优学习率区间以下是基于 PyTorch 的微调训练最小脚本覆盖了从数据加载到训练循环的关键参数# train.py import torch import torch.nn as nn from torchvision import models, transforms, datasets # 数据增强垃圾分类场景轻量增强即可重增强会破坏颜色特征 transform_train transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.7, 1.0)), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness0.3, contrast0.3, saturation0.2), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) num_classes 40 model.fc nn.Linear(model.fc.in_features, num_classes) # 冻结骨干只训练全连接层 for param in model.parameters(): param.requires_grad False for param in model.fc.parameters(): param.requires_grad True optimizer torch.optim.Adam(model.fc.parameters(), lr1e-3, weight_decay1e-4) criterion nn.CrossEntropyLoss() # 训练 10 个 epoch 后解冻骨干用更低学习率微调 # 注意解冻后学习率必须降到 1e-5 以下否则预训练特征会被破坏参数上有几个关键点。第一学习率。冻结骨干阶段 1e-3 是安全的但解冻骨干之后必须降到 1e-5 以下否则骨干庞大的梯度会直接覆盖掉预训练学到的通用特征。第二数据增强里 ColorJitter 的 saturation 不能超过 0.3塑料瓶的颜色饱和度是识别关键词增强太狠会让模型分不清不同颜色的瓶子。第三weight_decay 用 1e-4 还是 5e-4 看类别数量类别少就小一点。4.3 Transformer 分支ViT 和 Swin 在垃圾分类中的实际表现最新的图像分类模型形态基本都转向了 Transformer 架构在垃圾分类上表现如何需要谨慎判断。ViT 在没有大规模数据预训练时是不如有同量级 CNN 的所以垃圾分类项目里直接用 ViT 的意义有限除非用上 CLIP 这类预训练模型来做零样本或者少样本分类。Swin Transformer 因为引入了窗口注意力计算效率比 ViT 高更适合作为分类骨干。我在实验中对比过 ResNet18 和 Swin-TTiny 版在同一个垃圾分类数据集上的表现结论是 Swin-T 的准确率比 ResNet18 高约 2%但推理延迟增加了近 3 倍。下表是一个典型的对比结果指标ResNet18ResNet34Swin-TViT-B/16Top-1 准确率91.2%93.1%94.0%93.5%参数量11M21M28M86MCPU 推理延迟224×22435ms55ms120ms180msGPU 推理延迟Jetson Nano18ms26ms58ms80ms结论很直接如果算力宽裕Swin-T 值得上如果是边缘设备ResNet18 的性价比最高。Transformer 序列化建模对长距离依赖的捕捉能力在垃圾分类场景里不是刚需垃圾个体并不大局部特征已经足够判别。4.4 类别不均衡与易混淆类的损失函数调整垃圾分类数据天然不均衡有害垃圾在经济生活中的占比远低于可回收垃圾导致数据量可能差 10 倍以上。处理手段有两个一是数据层面做过采样或者用数据增强补足少数类二是损失函数层面用 Focal Loss 或者类别权重加权。Focal Loss 的核心参数是 gamma 和 alpha。gamma 控制对困难样本的关注程度设 2.0 时模型会把注意力集中到那些被混淆的类别上alpha 用来平衡类别频率比如有害垃圾占比 5% 时alpha 设 10.0 左右。一个常见的误用是把这两个参数同时设置得过高模型会开始过拟合到少数类的个别样本上实际上 gamma 取 1.5 到 2.0 之间alpha 取最大类/最小类的频率比就足够。面对易混淆的类别对比如打湿的纸巾和厨余垃圾除了调损失函数还可以考虑在模型输出层之后加一个手动规则预测概率低于 0.6 的样本强制改判为“其他垃圾”。这个兜底逻辑在垃圾分类场景下非常实用因为误判回收物比误判其他垃圾更影响回收链流程。5. 从 PyTorch 到边缘部署ONNX 导出、TensorRT 加速与硬件选型训练完的模型只是项目的一半垃圾分类系统的另一半是部署。部署环节决定了两件事能不能跑在算力有限的边缘设备上以及推理时延能不能满足传送带节奏。部署上常见的问题是 PyTorch 模型直接上生产环境GPU 内存占用高、推理延迟大还得装全套 PyTorch 运行时。5.1 ONNX 导出与动态轴配置的注意点模型部署的第一个动作是导出 ONNX但直接导出会有两个坑。第一PyTorch 模型默认的 forward 入参包含可变的 batch 维度导出的 ONNX 模型如果锁死了 batch1会导致推理时固定只能一次处理一张图需要在导出时显式声明动态轴。第二模型里的某些算子比如 adaptive avg pool 在转 ONNX 时可能不被支持需要用固定尺寸的 avg pool 替代。# export_onnx.py import torch import onnx from models import get_model model get_model(resnet18, num_classes40) model.load_state_dict(torch.load(best_model.pth)) model.eval() dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, garbage_cls.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}}, opset_version13 ) # 校验导出后的模型 onnx_model onnx.load(garbage_cls.onnx) onnx.checker.check_model(onnx_model) print(ONNX export OK, ir_version:, onnx_model.ir_version)opset 版本建议设 13 以上低版本对 ResNet 里的算子是够的但涉及 Swin Transformer 的窗口注意力时需要更高的算子集支持。导出后先用 onnxruntime 跑一遍验证输出和 PyTorch 原始模型的输出误差在 1e-4 量级内才能继续往下一步走。5.2 TensorRT 与 INT8 量化的精度与延迟实测如果目标硬件是 NVIDIA 的 Jetson 系列或者带 TensorCore 的 GPUONNX Runtime 的上限没被充分释放。因为 ONNX Runtime 走的是 CUDA 通用算子TensorRT 会对整个推理图做层融合和算子选择延迟可以再降低一半以上。TensorRT 部署流程是把 ONNX 模型转成 TensorRT engine过程中可以指定精度类型# 转成 FP16 engine trtexec --onnxgarbage_cls.onnx --saveEnginegarbage_fp16.engine --fp16 # 转成 INT8 engine需要校准数据 trtexec --onnxgarbage_cls.onnx --saveEnginegarbage_int8.engine \ --int8 --calib/path/to/calibration_data校准数据在 INT8 量化里是关键环节。校准集要选 500 到 1000 张有代表性的图像覆盖各类别和不同的光照条件。量化后准确率通常会掉 1-2%但对延迟敏感的场景值得。INT8 的另一个好处是显存占用只有 FP16 的一半对只有 4GB 显存的 Jetson Nano 这类设备来说非常关键。在硬件选型上不同场景有不同的最合适选择部署场景推荐硬件实测延迟ResNet18参考算力成本实验室 DemoCPUIntel i535ms低社区智能垃圾桶树莓派 4B120ms低小区集中投放点Jetson Nano18ms中产线快速分拣工控机 RTX 30605ms较高树莓派跑 120ms 意味着每秒只能处理 8 张图如果传送带速度比较高就顶不住了。现场应基于物体通过视野的时间来反推需要的推理延迟一般要求推理耗时不超过物体在视野中出现时长的 1/3。5.3 数据预处理在部署端的对齐问题部署时最容易出现的问题反而是最不起眼的预处理不一致。PyTorch 训练时的预处理用 PIL 读取和 Resize部署时如果用 OpenCV 读取 BGR 图像直接输入通道顺序和归一化方式全部错了模型准确率能直接掉 20 多个点。解决方法是把归一化和通道变换全部统一在推理代码里先cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)转通道再归一化到 [0,1]再用均值标准差做标准化最后转成 NCHW 布局。这一步看起来基础但我在实际项目里至少遇到三次以上因为忘记转通道而出现的“模型部署后失效”问题。正确做法是写一段小程序同时跑 PyTorch 和 ONNX/TensorRT输入同一张图对比输出 softmax 概率是否一致。6. 模型准确率不达标时最值得调的 3 个开关垃圾分类模型在现场跑着跑着准确率掉下来或者实验室里就达不到目标多数情况不是模型架构有问题而是几个可以系统排查的环节出了问题。以下 3 个检查项是我每到一个新项目必查的它们能覆盖 80% 的精度问题。6.1 检查背景过拟合中间层特征图可视化背景过拟合的典型表现是测试集上某一类准确率特别高但换一个背景后准确率崩掉。用 Grad-CAM 输出热力图叠加在原图上可以看到模型到底在看什么。# grad_cam.py import cv2 import torch from pytorch_grad_cam import GradCAM from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget from pytorch_grad_cam.utils.image import show_cam_on_image cam GradCAM(modelmodel, target_layers[model.layer4[-1]]) input_tensor preprocess_frame(cv2.imread(test.jpg)) targets [ClassifierOutputTarget(predicted_class_id)] grayscale_cam cam(input_tensorinput_tensor, targetstargets) visualization show_cam_on_image(input_tensor_normalized, grayscale_cam[0], use_rgbTrue)如果热力图的注意力集中在背景桌面上而不是物体本身说明模型被背景干扰了。解决办法是回到标注清洗阶段确保每类图像的背景有足够的多样性开启更多旋转、裁剪增强也能缓解。6.2 检查类间混淆的具体样本混淆矩阵诊断只看总准确率无法定位问题根源。输出一张混淆矩阵找出被混淆最多的类别对。垃圾场景里最容易混淆的通常是“塑料瓶 vs 玻璃瓶”“打湿的纸 vs 厨余垃圾”。如果混淆矩阵显示这两个类互相错判率超过 15%不要去加网络深度应该回看数据集中这两个类的样本区分度。对塑料瓶和玻璃瓶模型的区分能力受限是因为两者反光特征相似。此时可以放弃 RGB 图像改用 HSV 空间提取饱和度特征再加一个并行分支两个分支的特征拼接后分类这个方法对透明/半透明物体的区分效果立竿见影。6.3 一个典型的易犯错误预测置信度过滤机制模型输出的 softmax 概率并不总是可信的一个不分青红皂白把所有物体都判成“其他垃圾”的模型容易出现高置信度的错误预测。现场系统必须给预测结果加一道置信度闸门。probs, idx torch.max(softmax_output, dim1) if probs.item() 0.45: # 置信度阈值按场景调节 final_label unknown # 转人工复核或重新拍照 else: final_label class_names[idx.item()]阈值 0.45 是一个经验值。设置太高会增加“不识别”的比例造成物体重复拍照、产线卡顿设置太低则失去过滤意义。我的做法是先统计验证集上所有正确预测样本的置信度分布取分布第 10 百分位的值作为初始阈值再在现场根据误报率微调。这才是真正值得写在项目文档里的部署经验。本文还有配套的精品资源点击获取