简介这套资料以垃圾分类为切入点面向入门神经网络与OpenCV图像处理的开发者适合用来快速搭建一个可用的图像分类演示项目也可作为课程设计或算法入门实践的参考。压缩包共1046个文件其中以1041张jpg垃圾分类图片为主另有2个Python脚本train.py负责训练、predict.py负责预测、1个h5模型权重文件、1个txt说明文件和1段mp4演示视频整体大小824.68MBmp4可直观展示运行效果txt则补充基础使用提示。目前已有1331人浏览学习。通过学习可以获得完整数据集、数据预处理与制作思路、训练和预测代码模型输出可对应干垃圾、湿垃圾、可回收垃圾、有害垃圾四类并将类别以中文形式直接显示在图片上这套从数据到预测的流程可当作图像分类任务的基础模板便于理解图片读取、模型保存与加载等环节也能迁移到其他分类项目中。1. 垃圾分类图像分类一套能直接跑通的极小模型模板做图像分类的人手头最缺的往往不是模型结构而是一套「从图片到预测结果」的完整闭环。这份资源就是一个很好的起点它用干垃圾、湿垃圾、可回收垃圾、有害垃圾四类图片训练出一个 h5 格式的卷积神经网络模型并配好了训练脚本和预测脚本。你不需要自己从零写数据加载、标签映射和模型保存逻辑跑通之后稍加改动就能迁移到其他图像分类任务上。整套东西的核心价值在于「小」模型规模小、数据集规模小、代码量小。正因为小每一步都看得清——图片怎么喂进去、标签怎么对应、训练参数怎么调、预测结果怎么显示。它适合两类人一类是想交图像处理或机器学习课程作业的学生另一类是想快速上手 TensorFlow 图像分类流程的入门开发者。下文我会从数据组织、网络结构、训练配置、预测脚本到常见坑完整拆一遍这套模板。2. 数据集与预处理图片改名、标签映射与训练集划分2.1 先搞懂数据目录的结构这套资源的图片文件名是 img_11548.jpg、img_11351.jpg 这类编号模型文件是 my_model.h5脚本是 train.py 和 predict.py。如果你直接打开压缩包看到的应该是一个没有按类别分子目录的扁平目录——也就是说类别信息不在路径里而要通过别的方式告诉训练脚本。常见做法是在数据目录下建四个子文件夹分别命名为 dry、wet、recyclable、harmful对应干垃圾、湿垃圾、可回收垃圾、有害垃圾。然后把散落的图片按实际类别移动进去。这一步不做train.py 里读不到标签训练出来的模型也没有意义。mkdir -p dataset/train/dry dataset/train/wet dataset/train/recyclable dataset/train/harmful mkdir -p dataset/val/dry dataset/val/wet dataset/val/recyclable dataset/val/harmful # 示例把一张干垃圾图片放入训练集 cp img_11548.jpg dataset/train/dry/ # 示例把一张湿垃圾图片放入验证集 cp img_11351.jpg dataset/val/wet/这里我把数据集拆成了 train 和 val 两个大目录各自下面再按类别分子目录。这么做的好处是 TensorFlow 的 image_dataset_from_directory 或 Keras 的 ImageDataGenerator.flow_from_directory 可以直接按目录名生成标签不需要自己维护 CSV 或数组。目录名必须用英文字母不能直接写「干垃圾」这类中文名否则后续标签映射会非常痛苦。2.2 标签映射目录名到中文类别的桥模型输出的是四个类别的概率分布但预测脚本要在图片上显示「干垃圾」这样的中文这里就需要一个映射表。实际项目里我一般会在一个 config.py 或直接在 predict.py 顶部写死这个映射# label_map.py class_names [dry, wet, recyclable, harmful] class_names_cn [干垃圾, 湿垃圾, 可回收垃圾, 有害垃圾] # 模型输出的是索引比如 argmax 结果是 0就对应 class_names[0] dry # 再通过 zip 把英文索引映射到中文 label_dict dict(zip(class_names, class_names_cn))这段代码的逻辑很简单class_names列表的顺序必须和训练时flow_from_directory自动生成的类别顺序一致否则预测结果会错位。zip把英文名和中文名一一对应之后查询时直接用label_dict[英文名]就能拿到中文。一个常见的坑是训练时目录名排过序比如 flow_from_directory 会按字母序给类别编号dry 是 0、harmful 是 1、recyclable 是 2、wet 是 3如果你手工写映射时顺序不对预测就会张冠李戴。2.3 训练集 / 验证集划分的比例选择这套模板的数据集本身不大如果图片总数只有几百张训练集和验证集按 8:2 划分是比较稳妥的。验证集比例太低评估结果波动会很大比例太高训练数据不够模型学不到东西。我的习惯是先按 8:2 分跑完一轮看验证集准确率是否明显低于训练集。如果是再考虑把验证集压到 15% 或者加上数据增强来缓解过拟合。划分时注意一点类别要按比例分。比如干垃圾有 200 张、湿垃圾只有 80 张不能简单地把所有图片随机打乱后按 8:2 切否则可能出现湿垃圾在验证集里只有十来张的极端情况。稳妥做法是用 sklearn 的 train_test_split 按类别分层划分或者干脆每个类别目录单独抽样。数据量小的时候分层抽样比全局随机抽样靠谱得多。3. 网络结构DIY 小型卷积网络与参数设计3.1 为什么选小型卷积网络而不是预训练模型这套模板主要面向学习和课程设计没有选择 ResNet、VGG16 这类预训练大模型而是自己搭了一个轻量卷积网络输入尺寸也不需要太大控制在 128×128 或 224×224 以内。原因很现实预训练模型动辄上百 MB训练时需要加载在 ImageNet 上学到的权重微调时参数量大、显存占用高、训练时间长对一张普通笔记本显卡或纯 CPU 环境并不友好。而小型卷积网络在几百张图片的小数据集上只要训练得当准确率也能到 80% 以上且单轮训练时间以秒计。实际项目里我一般是先跑小型网络验证数据没问题再考虑换成预训练模型刷精度。反过来直接上大模型数据量不够时很容易过拟合而且排查问题的成本会高很多——你不知道是数据有问题还是模型太大学不动。3.2 网络的每一层在做什么这条 DIY 网络本质上是简化的 VGG 风格卷积层提取特征池化层降维全连接层做分类。具体结构大致如下层序号层类型参数配置输出尺寸以 128×128×3 输入为例1Conv2D32 个 3×3 卷积核ReLU126×126×322MaxPooling2D池化窗口 2×263×63×323Conv2D64 个 3×3 卷积核ReLU61×61×644MaxPooling2D池化窗口 2×230×30×645Conv2D128 个 3×3 卷积核ReLU28×28×1286MaxPooling2D池化窗口 2×214×14×1287Flatten拉平为一维250888Dense512 个神经元ReLU5129Dropout比率 0.551210Dense4 个神经元Softmax4卷积核数量从 32 逐步翻倍到 128是因为浅层提取边缘、颜色等低级特征不需要太多卷积核深层需要更丰富的语义特征所以通道数要增加。每次卷积后接池化把空间尺寸减半既能降低计算量也能在一定程度上防止过拟合。全连接层用 512 个神经元算是折中——再多容易过拟合再少表达能力不够。最后的 Softmax 输出 4 个值分别对应四类垃圾的概率。3.3 自己改网络时要注意的尺寸计算如果你要改输入尺寸或加深网络需要自己算卷积和池化后的特征图尺寸。公式是输出尺寸 (输入尺寸 - 卷积核大小 2×填充) / 步长 1。池化层同理不过更简单直接整除 2 就行。import tensorflow as tf def build_model(input_shape(128, 128, 3), num_classes4): model tf.keras.Sequential([ tf.keras.layers.Conv2D(32, (3, 3), activationrelu, input_shapeinput_shape), tf.keras.layers.MaxPooling2D(2, 2), tf.keras.layers.Conv2D(64, (3, 3), activationrelu), tf.keras.layers.MaxPooling2D(2, 2), tf.keras.layers.Conv2D(128, (3, 3), activationrelu), tf.keras.layers.MaxPooling2D(2, 2), tf.keras.layers.Flatten(), tf.keras.layers.Dense(512, activationrelu), tf.keras.layers.Dropout(0.5), tf.keras.layers.Dense(num_classes, activationsoftmax) ]) return model model build_model() model.summary()这里的关键参数是input_shape和num_classes。input_shape必须和训练时图片缩放的目标尺寸一致否则模型加载权重时会报维度不匹配。num_classes必须等于你的类别数四类垃圾就是 4。Dropout(0.5)只对全连接层起作用——训练时随机丢弃一半神经元让网络不过度依赖某几个节点的输出推理时会自动关闭。如果你想加深网络复制一组「Conv2D MaxPooling2D」加在中间即可但要保证每层参数与上一层输出尺寸匹配否则 summary 阶段就会报错。4. 训练与验证train.py 逐行拆解与训练日志分析4.1 数据加载与增强配置train.py 的核心逻辑分为三部分数据加载、模型编译、训练与保存。先看数据加载这一段它用的是 Keras 的 ImageDataGenerator这比 tf.data 上手简单特别适合本场景。from tensorflow.keras.preprocessing.image import ImageDataGenerator train_datagen ImageDataGenerator( rescale1./255, rotation_range20, width_shift_range0.2, height_shift_range0.2, shear_range0.2, zoom_range0.2, horizontal_flipTrue ) val_datagen ImageDataGenerator(rescale1./255) train_generator train_datagen.flow_from_directory( dataset/train, target_size(128, 128), batch_size32, class_modecategorical ) val_generator val_datagen.flow_from_directory( dataset/val, target_size(128, 128), batch_size32, class_modecategorical )这段代码里rescale1./255是把像素值从 0-255 缩放到 0-1 之间这是神经网络训练的标准做法——大数值输入会让梯度更新不稳定收敛慢。后面的 rotation_range、shift_range、zoom_range 都是数据增强参数训练时随机旋转、平移、缩放、翻转图片相当于免费扩充数据集让模型见过更多变体缓解过拟合。注意验证集只做缩放不能做增强否则评估结果不真实。target_size必须和网络输入尺寸一致class_modecategorical意味着标签是 one-hot 编码和网络最后的 Softmax 输出匹配。4.2 编译参数优化器、学习率与损失函数model.compile( optimizertf.keras.optimizers.Adam(learning_rate0.001), losscategorical_crossentropy, metrics[accuracy] )Adam 优化器在这种小规模任务上是零调试成本的默认选择它自适应调节每个参数的学习率不太需要手动调整。学习率 0.001 是 Adam 的默认值对这类小网络通常够用。如果训练时发现损失值振荡很大可以把学习率降到 0.0005如果损失下降极慢可以试着增大到 0.002但风险是训练不稳定。categorical_crossentropy是多分类的标准损失函数它计算预测概率分布和真实 one-hot 标签之间的交叉熵——预测越接近真实分布损失越小。如果你的标签不是 one-hot 而是整数那就需要用sparse_categorical_crossentropy这两者不能混用。4.3 回调函数早停与模型保存from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint callbacks [ EarlyStopping(monitorval_loss, patience10, restore_best_weightsTrue), ModelCheckpoint(my_model.h5, monitorval_loss, save_best_onlyTrue) ] history model.fit( train_generator, steps_per_epochtrain_generator.samples // 32, validation_dataval_generator, validation_stepsval_generator.samples // 32, epochs50, callbackscallbacks )EarlyStopping 是防止过拟合的最后一道防线它持续监控验证集损失如果连续 10 轮没有下降就提前结束训练并把权重恢复到验证集表现最好的那一轮。ModelCheckpoint 只保存验证集损失最低的模型这样即使后面几轮过拟合了h5 文件里存着的还是最优版本。steps_per_epoch的计算口径值得说它是训练样本数除以 batch_size。比如训练集有 400 张图batch_size 是 32那一个 epoch 需要跑 12.5 步代码里用整除得到 12最后那 8 张图在这个 epoch 里就不参与训练了。问题不大但如果你追求完整遍历可以用 math.ceil 向上取整。真正要注意的是验证集大小不能小于 batch_size否则 validation_steps 会算出 0验证直接报错不跑。4.4 训练日志哪些数字说明模型在健康收敛训练时终端会逐轮打印 loss、accuracy、val_loss、val_accuracy 四列。你要盯的是 val_loss 和 val_accuracy。健康的训练曲线是训练损失和验证损失都逐步下降两者差距不大。如果训练损失降得很低但验证损失反而上涨这就是典型的过拟合信号说明模型在背训练集而不是学规律。如果两者都降不下去问题多半出在数据或模型结构上——比如数据量太少、类别不平衡、学习率过大导致梯度振荡。判断训练是否结束不要只看准确率。一个小数据集上训练准确率很容易冲到 98% 以上这时候你要看验证准确率有没有跟上。本场景下四分类能达到 85% 以上的验证准确率就算不错毕竟样本量和模型容量都有限。如果验证准确率只有 60%-70%优先检查数据有没有标错、增强是否过度、训练集和验证集是否来自同一分布。5. 避坑清单训练不收敛、显存不足、预测全错的排查路径5.1 训练损失不降反升学习率过大加数据没归一化现象train.py 跑起来后loss 在 20 到 50 轮内完全没有下降趋势甚至微微上涨accuracy 在 25% 左右徘徊。原因最常见的是两个——数据没有做归一化像素值在 0-255 的范围直接喂给网络梯度更新时数值波动太大或者学习率设得过高Adam 在损失面上越过最低点来回振荡。解决先确认 ImageDataGenerator 里有没有rescale1./255。如果有把learning_rate从 0.001 改到 0.0005 重新跑。如果还是不行检查类别顺序——flow_from_directory会自动把目录名按字母序排序dry、harmful、recyclable、wet 分别对应 0、1、2、3手动写的标签映射如果没按这个顺序损失就算下降预测也是错的。5.2 报错 Failed to get convolution algorithm显存不够或 CUDA 配置有问题现象训练刚开始就抛错提示卷积算法初始化失败有时候还会附带Resource exhausted: OOM字样。原因GPU 显存不够或者 TensorFlow 在尝试申请显存时被系统拒绝。常见于显存只有 2-4GB 的入门级显卡128×128 输入加 32 batch 的配置显存占用会在 1.5GB 到 3GB 之间浮动。解决把batch_size从 32 降到 16 或 8模型占用显存会直线下降。再不行就把target_size从 (128, 128) 降到 (96, 96)——注意网络结构里输入维度也要同步改。如果是纯 CPU 环境训练会慢但不会报这类错耐心等就行因为模型很小。5.3 预测结果全是同一类类别不平衡加 Softmax 输出偏差现象predict.py 跑任何一张图输出都是「干垃圾」准确率看起来像是「生效了」但实际完全没有分类能力。原因四类样本数量严重不均。比如干垃圾有 500 张图其他三类各 50 张模型发现全部预测为干垃圾就能拿到 80% 的准确率它就这么「偷懒」了。Softmax 输出的概率分布偏向样本多的类别argmax 几乎永远是 0。解决检查数据分布四类样本数最好不要超过 2:1。如果差距大给每个类别设置不同的类别权重或者对样本少的类别做额外的数据增强。Keras 里可以在model.fit里传class_weight参数计算方式是总样本数 / (类别数 × 该类别样本数)。还有一个临时验证手段打印预测概率分布看是不是每个类别的概率都趋近于均值。5.4 h5 文件加载报错输入尺寸和训练时不匹配现象predict.py 加载 my_model.h5 没有问题但跑model.predict时报维度错误或者结果明显不对。原因h5 文件里保存的模型输入形状和预测脚本里target_size不一致。比如训练用的是 (128, 128)但 predict.py 里cv2.resize或load_img的target_size写成了 (224, 224)。解决用model.summary()打印模型结构看第一层的 input_shape 是多少再去改 predict.py 里的target_size两处必须严格一致。这个坑之所以常见是因为模型文件本身不报错——TensorFlow 只在真正传入数据时才发现形状对不上。6. 预测脚本与进阶中文标注显示与批量图片处理技巧predict.py 的职责是加载训练好的模型对单张图片输出类别概率并把中文类别名绘制到图片上。核心代码如下import cv2 import numpy as np from tensorflow.keras.models import load_model from PIL import Image, ImageDraw, ImageFont model load_model(my_model.h5) def predict_image(img_path): img cv2.imread(img_path) img_rgb cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img_resized cv2.resize(img_rgb, (128, 128)) img_array np.expand_dims(img_resized.astype(float32) / 255.0, axis0) pred model.predict(img_array)[0] class_idx np.argmax(pred) confidence pred[class_idx] return class_idx, confidence def draw_chinese_text(img, text, position(10, 10)): img_pil Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) font ImageFont.truetype(simhei.ttf, 24) draw ImageDraw.Draw(img_pil) draw.text(position, text, fontfont, fill(255, 0, 0)) return cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR) class_idx, confidence predict_image(img_11548.jpg) print(类别索引:, class_idx, 置信度:, round(confidence, 4)) img cv2.imread(img_11548.jpg) img draw_chinese_text(img, 干垃圾) cv2.imwrite(output.jpg, img)这里有两个容易翻车的点。第一OpenCV 的imread默认读入 BGR 格式直接给神经网络训练时用的是 RGB所以要cvtColor转一次否则模型看到的颜色通道是反的预测结果会受影响。第二cv2.putText不支持中文所以用 PIL 把图片转成 PIL 图像绘制文字再转回 OpenCV 格式。simhei.ttf是黑体字体文件Windows 系统在 C:\Windows\Fonts 目录下Linux 系统需要单独安装中文字体否则ImageFont.truetype会抛文件找不到的异常。如果想对批量图片做预测循环调用predict_image即可但注意逐张load_model是低效的——模型加载一次放内存里循环里只调predict能省掉大量重复开销。我把这个习惯带到了所有项目里模型加载和数据处理分开写模型只 load 一次数据流走循环。进阶一点这套模板完全可以改成摄像头实时分类把cv2.imread换成cv2.VideoCapture(0)读取视频帧每帧缩放到 (128, 128) 后喂给模型再把置信度低于阈值的帧直接丢弃避免误报。实际效果取决于你的数据分布——训练集里的图片背景、光照是否接近摄像头画面如果差得远实际准确率会明显打折。我自己的经验是这种场景下最好采集一部分摄像头真实画面混入训练集不然就是「训练一时爽部署火葬场」的翻车现场。从那以后我每次做图像分类项目都会强制走一遍「真实部署环境采集几十张图看看效果」的验证流程希望帮到你。本文还有配套的精品资源点击获取