
简介面向图像分类与计算机视觉入门者的一套垃圾分类实战项目基于TensorFlow和OpenCV构建自定义神经网络模型覆盖从数据集准备、模型训练到预测推理的完整流程。项目实现了干垃圾、湿垃圾、可回收垃圾、有害垃圾四类识别提供train.py与predict.py两个Python脚本前者用于训练自定义网络后者可加载模型并对图片预测结果以中文直接显示在图像上。资源共包含1046个文件其中1041张JPG图片构成分类数据集另有2个Python脚本、1个已训练好的H5模型文件、1个操作演示MP4视频及1个说明文档压缩包整体大小约824.68MB目录结构清晰便于按需查阅。目前已有1332人学习下载非常适合深度学习新手模仿改造既可对照完整代码理解图像分类原理也可将本项目作为通用模板迁移到其他多类别图像识别任务中。1. 一个能跑通的图像分类模板垃圾分类四分类从训练到中文标注图像分类入门最尴尬的时刻不是看不懂卷积原理而是拿着原理搭工程时卡在数据集怎么组织、训练完怎么保存、预测时怎么把结果画回图上这三段衔接上。这套基于 TensorFlow 和 OpenCV 的垃圾分类项目恰好把三段全串在了一起干垃圾、湿垃圾、可回收垃圾、有害垃圾四分类附带 DIY 卷积网络、train.py 训练脚本、predict.py 预测脚本训练好的 my_model.h5 和一批测试图都在包里解压改下路径就能跑通。更关键的是它没用预训练模型网络完全手动搭建既是垃圾分类 demo也是能套到任意图像分类任务上的通用模板。2. 数据集与目录约定四分类任务的样本组织与预检一个图像分类项目能不能跑通第一步往往不是模型结构而是数据目录。Keras 的 flow_from_directory 有一套严格的目录约定把每个子文件夹名当类别标签自动读取图片并生成 one-hot 编码。这个机制对垃圾分类这种天然按类别归档的场景非常适合但前提是你必须按它约定的规则组织文件否则脚本会静默跳过图或者直接报 Found 0 images。2.1 四种垃圾类别的目录组织方式推荐目录结构如下garbage_dataset/ ├── train/ │ ├── dry/ # 干垃圾 │ ├── wet/ # 湿垃圾 │ ├── recyclable/ # 可回收垃圾 │ └── harmful/ # 有害垃圾 ├── val/ │ ├── dry/ │ ├── wet/ │ ├── recyclable/ │ └── harmful/ └── test/每个子文件夹内部的图片命名用什么前缀不重要train.py 只认最底层文件夹名。类别名为什么用英文因为 flow_from_directory 按 ASCII 顺序生成类别索引中文文件夹名会引入字符编码上的不确定性常见做法是用拼音或英文存文件夹中文显示交给 predict.py 在画图时做映射。这里有个容易踩的隐性点类别的索引顺序不取决于你创建文件夹的顺序而是按文件夹名的字母序排。上面四个文件夹按字母序是 dry、harmful、recyclable、wet对应索引 0、1、2、3和你直觉里的「干、有、可、湿」顺序完全对不上文件夹名字母序索引中文标签one-hot 编码dry0干垃圾[1,0,0,0]harmful1有害垃圾[0,1,0,0]recyclable2可回收垃圾[0,0,1,0]wet3湿垃圾[0,0,0,1]后面 predict.py 里手动映射中文标签时必须按照这个 class_indices 的顺序来写这是整个项目里最容易翻车的点。我一般会在训练脚本里先打印 train_generator.class_indices确认顺序再往下走。注意类别索引的顺序由文件夹名的字母序决定而不是创建顺序。train.py 打印的 class_indices 是 predict.py 映射中文标签的唯一依据。样本数量方面这类小项目每类 20~50 张图就能跑通流程但类别之间要尽量均衡。四分类如果一类 50 张、另一类只有 10 张模型会倾向于把不确定样本判给样本多的类后期提升准确率很被动。从原始照片整理成这个结构也不复杂几行脚本就能完成import os import shutil # raw_photos 下先手工按类别分好raw_photos/dry、raw_photos/wet ... raw_root raw_photos target_root garbage_dataset for cls in [dry, wet, recyclable, harmful]: for subset in [train, val]: os.makedirs(f{target_root}/{subset}/{cls}, exist_okTrue) files os.listdir(f{raw_root}/{cls}) for i, name in enumerate(files): # 每 5 张抽 1 张进 val其余进 train示例用取模足够 subset val if i % 5 0 else train shutil.copyfile( f{raw_root}/{cls}/{name}, f{target_root}/{subset}/{cls}/{name} )这段脚本做的事是按类别建立 train/val 双层目录然后用取模的方式做随机划分。i % 5 0意味着每个类约 20% 的图进验证集。真实项目里建议换成random.shuffle固定随机种子再做 split避免每次跑结果不一致。2.2 图像预检用 OpenCV 扫出损坏文件与异常通道数据集的坑都不在明面上最典型的是网上爬下来或手机导出的图里混入了损坏文件、灰度图或带透明通道的 PNG训练跑到一半报错错误信息却指向不明。我一般先跑一段预检脚本import cv2 import imghdr from pathlib import Path data_root Path(garbage_dataset) for subset in [train, val, test]: subset_dir data_root / subset if not subset_dir.exists(): continue for cls_dir in subset_dir.iterdir(): if not cls_dir.is_dir(): continue for img_path in cls_dir.iterdir(): # imghdr 按内容检测真实格式扩展名造假骗不过它 if imghdr.what(img_path) is None: print(损坏文件:, img_path) continue # OpenCV 解不出来的一律删掉或重新导出 img cv2.imread(str(img_path)) if img is None: print(OpenCV 无法读取:, img_path) continue h, w, c img.shape if c ! 3: print(非三通道图像:, img_path, shape, (h, w, c)) if min(h, w) 64: print(分辨率过小:, img_path, shape, (h, w, c))脚本分三层检查imghdr 先按文件内容判断是不是真图像能挡住改名换后缀的假文件OpenCV 再尝试解码解不出来的直接标记最后检查通道数和最小边长。最小边长设成 64比训练尺寸 224 小得多只用来排除严重异常的图片避免训练时读入一张 20x20 的缩略图干扰模型。灰度图c1和 RGBA PNGc4建议统一转成三通道 JPG 再进目录因为网络输入层固定要求 shape(224, 224, 3)通道对不上时训练报错会非常隐晦。3. DIY 网络设计为什么手动搭卷积网络比迁移学习更适合当模板做图像分类不少人的第一反应是上 VGG、ResNet 预训练模型或者直接尝鲜最新的图像分类模型结构精度确实高。但这类方案对入门者有个致命问题整个网络像个黑匣子改一个参数你不知道影响什么出了问题也不知道从哪排查。这份资源刻意选择手动搭建网络我认为这是它最大的价值——你要学习的是图像分类的完整链路不是调包技巧。3.1 选型理由小数据、四分类、可解释性优先垃圾四分类本身不是难题数据集体量也小手动搭一个三层卷积的小网络完全够用。参数量控制在百万以内普通 CPU 上一个 epoch 只需几分钟单张图的预测在毫秒级对新手实验来说这个反馈速度很重要。手动网络训练稳定不依赖预训练权重下载断网环境下也能完整复现这是很多跑实验的从业者很在意的一点。对比之下迁移学习要加载几十上百 MB 的预训练权重微调逻辑也更绕哪些层冻结、哪些层解冻、学习率怎么分层设每一步都在增加认知负担。作为一份用来理解和改造的模板可解释性比极限精度重要得多。你改一个卷积核数量、加一层池化能立刻看到参数量和训练速度的变化这种反馈是预训练模型给不了的。网络尺寸上我的经验是滤波器通道数按 16 → 32 → 64 递增层输出尺寸通道变化作用Conv2D 3x3224x224x163→16提取边缘、颜色等底层特征MaxPooling 2x2112x112x16—下采样扩大感受野Conv2D 3x3112x112x3216→32组合成部件级特征MaxPooling 2x256x56x32—下采样Conv2D 3x356x56x6432→64提取语义特征MaxPooling 2x228x28x64—下采样Flatten50176—展平接全连接Dense Dropout64—特征聚合防过拟合Dense softmax4—输出四类概率通道递增有讲究浅层学边缘和颜色块深层需要学「像塑料袋还是像电池」这种语义概念特征容量必须逐层放大。反过来第一层就用 128 通道小数据集很快就过拟合训练集准确率很高、验证集一塌糊涂。3.2 Keras 网络结构代码与参数说明手动搭网络我用 Keras Sequential每一层就是一行改起来直观from tensorflow.keras import Sequential from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout model Sequential([ Conv2D(16, (3, 3), activationrelu, input_shape(224, 224, 3)), MaxPooling2D((2, 2)), Conv2D(32, (3, 3), activationrelu), MaxPooling2D((2, 2)), Conv2D(64, (3, 3), activationrelu), MaxPooling2D((2, 2)), Flatten(), Dense(64, activationrelu), Dropout(0.5), Dense(4, activationsoftmax) ]) model.summary()这套结构有三个参数需要重点说明。input_shape(224, 224, 3) 里的 224 必须和 train.py 里 flow_from_directory 的 target_size 完全一致否则模型输入和实际喂进来的数据对不上预测阶段 resize 也要按同一尺寸。Dropout(0.5) 插在第一个全连接层之后训练时随机扔掉一半神经元这是小数据集防过拟合性价比最高的一招它只在训练时生效预测时自动关闭predict.py 里不需要额外处理。最后 Dense(4, activationsoftmax) 的 4 是类别数之后改成自己的任务时这个数字必须跟着标签数量变改错会在训练时报 shape 不匹配。4. train.py 详解数据增强、训练配置与 my_model.h5 的导出train.py 是这份资源的主训练入口做的事情无非三件把第 2 章的目录结构转成数据流、把第 3 章的网络真正跑起来、把结果存成 my_model.h5。环境上只需要 tensorflow、opencv-python、Pillow 三个库装好就能跑。三个环节在脚本里对应三段代码拆开看不复杂。4.1 ImageDataGenerator 与数据流构造数据读取用 Keras 的 ImageDataGenerator 是最省事的方案。它把读图、缩放、增强、打包 batch、生成标签全封装了避免手写数据管道时在文件路径和 one-hot 编码上反复出错。from tensorflow.keras.preprocessing.image import ImageDataGenerator # 训练集带数据增强验证集只做缩放不做扰动 train_datagen ImageDataGenerator( rescale1.0 / 255, rotation_range10, width_shift_range0.1, height_shift_range0.1, zoom_range0.1, horizontal_flipTrue, fill_modenearest ) val_datagen ImageDataGenerator(rescale1.0 / 255) train_generator train_datagen.flow_from_directory( garbage_dataset/train, target_size(224, 224), # 与模型 input_shape 严格一致 batch_size16, class_modecategorical, # 输出四维 one-hot 标签 shuffleTrue ) val_generator val_datagen.flow_from_directory( garbage_dataset/val, target_size(224, 224), batch_size16, class_modecategorical, shuffleFalse # 验证集打乱会破坏每个 epoch 的可比性 ) print(训练集类别索引:, train_generator.class_indices) print(训练集样本数:, train_generator.samples)rescale1.0/255 是把像素值从 0~255 归一化到 0~1直接决定了网络能否收敛。rotation_range10 表示随机旋转正负 10 度width_shift_range0.1 表示水平平移最多 10% 的图宽zoom_range0.1 是随机缩放到 90%~110%这些小扰动让模型见到更多变体缓解样本不足。batch_size16 在显存有限时比较稳妥显存充裕可以调到 32加快训练但会增加内存占用。class_modecategorical 输出四维 one-hot 标签对应最后一层 softmax 的 4 个输出如果改成 sparse 或 binary损失函数和输出层都要跟着调整。验证集不增强、不 shuffle是为了让每个 epoch 的验证指标在同一批数据上做对比。有个建议养成的习惯打印 class_indices 和 samples。前者确认标签顺序后者确认每个子目录真的被读进来了某个文件夹命名不规范出现 0 样本时这里会立刻露出马脚。4.2 训练参数、模型保存与 class_indices 固化网络搭好后编译和训练参数是决定模型能不能学到东西的关键。我常用的配置如下from tensorflow.keras.optimizers import Adam # 四分类用 categorical_crossentropy学习率宁小勿大 model.compile( optimizerAdam(learning_rate1e-4), losscategorical_crossentropy, metrics[accuracy] ) # steps_per_epoch 取整防止一个 epoch 重复取到同一批数据 steps_per_epoch train_generator.samples // train_generator.batch_size validation_steps val_generator.samples // val_generator.batch_size history model.fit( train_generator, steps_per_epochsteps_per_epoch, epochs30, validation_dataval_generator, validation_stepsvalidation_steps ) # 保存完整模型预测时 load_model 一步恢复网络和权重 model.save(my_model.h5)learning_rate 从 1e-4 起步是我在小数据集上的经验值。学习率调大到 1e-2loss 会震荡甚至直接发散典型表现是 accuracy 卡在 0.25 附近——四分类的随机猜测概率就是 0.25模型等于没学。epochs30 对这种小数据集足够看到趋势loss 还在持续下降就往后加不用死守这个值。提示model.save(my_model.h5) 保存的是完整模型predict 时 load_model 一步到位model.save_weights 只存权重使用前必须重新搭一遍网络结构新手不建议走这条路。还有一个容易忽略的点class_indices 的映射关系不会写进 h5 文件。模型记住的是「索引 0 对应 dry」这种训练时的顺序这个顺序由 flow_from_directory 按文件夹字母序生成。所以 train.py 里打印的那行 class_indices最好原样抄进 predict.py 的标签列表两边保持一致模型输出才不会张冠李戴。我第一次跑这个项目时就是没做这一步预测时把可回收垃圾显示成了有害垃圾排查半天才发现是标签顺序对不上这步最玄学也最容易被忽略。5. 垃圾分类项目避坑清单训练与预测阶段的五个黑匣子这类项目的坑高度集中在两处训练阶段的数据读取预测阶段的预处理一致性。下面五条都是实际跑的时候翻过车的按「现象 → 原因 → 解决」写清楚照着排查能省半天时间。5.1 训练阶段的两个高频坑坑 1训练准确率长期卡在 0.25 附近loss 几乎不降。 现象accuracy 在 0.2~0.3 之间波动怎么看都不像在学。 原因四分类随机猜对的概率是 0.25准确率卡在这个值说明模型完全没学到特征。最常见是学习率过大导致 loss 震荡不收敛或者 ImageDataGenerator 里漏写 rescale像素值 0~255 直接喂进网络梯度被放大后无法正常优化。 解决把学习率降到 1e-4 重新训练并确认 train_datagen 里写了 rescale1.0/255。如果仍不降打印几个 batch 的标签检查四类样本是否都在数据流里。坑 2flow_from_directory 报错 Found 0 images或者 class_indices 里缺类。 现象生成器能创建但 samples 为 0打印 class_indices 发现某个类别消失。 原因目录层级不对。常见的是把图片直接放 garbage_dataset/train 下面而不是 train/dry 这种两级结构另一种是文件夹里混入了 .DS_Store、缩略图缓存等非图像文件旧版本 Keras 解析时会被干扰。 解决严格按第 2 章的层级放图片train 下一层必须是类别文件夹再下一层才是图片。跑一遍预检脚本清掉非图像文件再打印 class_indices 核对类别数。5.2 预测阶段的坑坑 3训练准确率 90% 以上预测单张图却全错。 现象训练指标正常拿 val 里某张图测类别差得离谱且每次错的类还不一样。 原因预测预处理和训练预处理不一致。最常见是训练做了 rescale1.0/255预测时直接用原图像素矩阵喂进去或者训练读的是 RGB预测用 cv2.imread 读出来是 BGR通道顺序反过来颜色语义全乱。 解决在 predict.py 里严格复刻训练的预处理链先转 RGB再 resize 到 (224, 224)再除以 255最后 np.expand_dims 加 batch 维度。建议把预处理写成独立函数训练和预测共用同一套逻辑从根上杜绝两边漂移。坑 4图片上中文标签显示成方框或乱码。 现象cv2.putText 画上去的「干垃圾」变成「???」或一排方框。 原因OpenCV 的 putText 底层不支持中文它只认 ASCII 字符集汉字传进去直接被丢弃或替换。这是 OpenCV 做图像处理时绕不开的编码坑。 解决中文绘制改用 PIL。思路是把 OpenCV 的 BGR 图转成 PIL 的 RGB 图用 ImageDraw 配合中文字体文件比如 msyh.ttc 微软雅黑或思源黑体绘制文字画完再转回 BGR 输出。顺手把矩形象底框也一起画文字在复杂背景上才看得清。坑 5load_model(my_model.h5) 报错提示 Unknown layer 或 h5py 相关问题。 现象模型文件在训练环境能加载换台机器或换个环境就报错。 原因TensorFlow/Keras 版本差异。Keras 2.4 和 TensorFlow 2.6 对优化器配置和 h5 解析的处理有区别旧版本保存的权重文件在新版本加载时兼容性出问题。 解决优先保证训练和预测环境的 tensorflow 版本一致还是搞不定就直接用当前环境重训一遍再保存比花时间逆向排查底层兼容快得多。6. 进阶验证把垃圾分类模板拆成你自己的图像分类任务四分类只是这套模板的第一个应用。换成任意图像分类任务比如木材纹理分类、车型识别、PCB 缺陷分类需要改动的只有三处类别文件夹、predict.py 里的标签列表、最后的全连接输出数。把 garbage_dataset/train 下的子文件夹换成你的类别名class_labels 换成对应的中文显示名最后一层 Dense(4) 改成 Dense(N)其中 N 是新的类别数。数据流、训练参数、保存与加载逻辑全部不用动这套模板的边界就在这里网络、数据流、训练流程是通用的你要改的是数据目录、标签映射和输出维度这三处。改完怎么验证模板没被改坏我最常做的验证方式是全量自检把 val 目录里每一张图都传进 predict 函数统计预测标签和文件所在文件夹名是否一致打印整体准确率。from pathlib import Path val_root Path(garbage_dataset/val) correct total 0 for cls_dir in val_root.iterdir(): if not cls_dir.is_dir(): continue for img_path in cls_dir.iterdir(): label, prob predict_image(str(img_path)) total 1 if label cls_dir.name: correct 1 else: print(误判:, img_path, -, label) print(验证准确率: %.2f%% % (correct / total * 100))这段代码直接复用 predict_image等于用真实图片反向检验目录组织、标签映射、预处理一致性三处有没有同步改对。如果验证准确率和训练时的 val accuracy 差距在 5 个百分点以内说明模板衔接正常差距过大优先查预处理的归一化和通道顺序而不是急着调网络结构。从那以后我每接手一个图像分类模板都强制走一遍「目录预检 → 打印 class_indices → 预处理写成共用函数 → 全量验证集自检」这条流水线四步走完再谈调参翻车概率低了很多。这个模板最值得借鉴的正是这种「每一步输入输出都看得见」的结构改成自己任务时能立刻定位该动哪里。希望帮到你。本文还有配套的精品资源点击获取