简介：这套资料以垃圾分类为切入点，面向入门神经网络与OpenCV图像处理的开发者，适合用来快速搭建一个可用的图像分类演示项目，也可作为课程设计或算法入门实践的参考。压缩包共1046个文件，其中以1041张jpg垃圾分类图片为主，另有2个Python脚本（train.py负责训练、predict.py负责预测）、1个h5模型权重文件、1个txt说明文件和1段mp4演示视频，整体大小824.68MB；mp4可直观展示运行效果，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 输入为例） |
|---|---|---|---|
| 1 | Conv2D | 32 个 3×3 卷积核，ReLU | 126×126×32 |
| 2 | MaxPooling2D | 池化窗口 2×2 | 63×63×32 |
| 3 | Conv2D | 64 个 3×3 卷积核，ReLU | 61×61×64 |
| 4 | MaxPooling2D | 池化窗口 2×2 | 30×30×64 |
| 5 | Conv2D | 128 个 3×3 卷积核，ReLU | 28×28×128 |
| 6 | MaxPooling2D | 池化窗口 2×2 | 14×14×128 |
| 7 | Flatten | 拉平为一维 | 25088 |
| 8 | Dense | 512 个神经元，ReLU | 512 |
| 9 | Dropout | 比率 0.5 | 512 |
| 10 | Dense | 4 个神经元，Softmax | 4 |
卷积核数量从 32 逐步翻倍到 128，是因为浅层提取边缘、颜色等低级特征不需要太多卷积核，深层需要更丰富的语义特征，所以通道数要增加。每次卷积后接池化，把空间尺寸减半，既能降低计算量，也能在一定程度上防止过拟合。全连接层用 512 个神经元算是折中——再多容易过拟合，再少表达能力不够。最后的 Softmax 输出 4 个值，分别对应四类垃圾的概率。
3.3 自己改网络时要注意的尺寸计算
如果你要改输入尺寸或加深网络，需要自己算卷积和池化后的特征图尺寸。公式是：输出尺寸 = (输入尺寸 - 卷积核大小 + 2×填充) / 步长 + 1。池化层同理，不过更简单，直接整除 2 就行。
import tensorflow as tf def build_model(input_shape=(128, 128, 3), num_classes=4): model = tf.keras.Sequential([ tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=input_shape), tf.keras.layers.MaxPooling2D(2, 2), tf.keras.layers.Conv2D(64, (3, 3), activation='relu'), tf.keras.layers.MaxPooling2D(2, 2), tf.keras.layers.Conv2D(128, (3, 3), activation='relu'), tf.keras.layers.MaxPooling2D(2, 2), tf.keras.layers.Flatten(), tf.keras.layers.Dense(512, activation='relu'), tf.keras.layers.Dropout(0.5), tf.keras.layers.Dense(num_classes, activation='softmax') ]) 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( rescale=1./255, rotation_range=20, width_shift_range=0.2, height_shift_range=0.2, shear_range=0.2, zoom_range=0.2, horizontal_flip=True ) val_datagen = ImageDataGenerator(rescale=1./255) train_generator = train_datagen.flow_from_directory( 'dataset/train', target_size=(128, 128), batch_size=32, class_mode='categorical' ) val_generator = val_datagen.flow_from_directory( 'dataset/val', target_size=(128, 128), batch_size=32, class_mode='categorical' )这段代码里，rescale=1./255是把像素值从 0-255 缩放到 0-1 之间，这是神经网络训练的标准做法——大数值输入会让梯度更新不稳定，收敛慢。后面的 rotation_range、shift_range、zoom_range 都是数据增强参数：训练时随机旋转、平移、缩放、翻转图片，相当于免费扩充数据集，让模型见过更多变体，缓解过拟合。注意验证集只做缩放，不能做增强，否则评估结果不真实。target_size必须和网络输入尺寸一致，class_mode='categorical'意味着标签是 one-hot 编码，和网络最后的 Softmax 输出匹配。
4.2 编译参数：优化器、学习率与损失函数
model.compile( optimizer=tf.keras.optimizers.Adam(learning_rate=0.001), loss='categorical_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(monitor='val_loss', patience=10, restore_best_weights=True), ModelCheckpoint('my_model.h5', monitor='val_loss', save_best_only=True) ] history = model.fit( train_generator, steps_per_epoch=train_generator.samples // 32, validation_data=val_generator, validation_steps=val_generator.samples // 32, epochs=50, callbacks=callbacks )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 里有没有rescale=1./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, axis=0) 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, font=font, 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) 后喂给模型，再把置信度低于阈值的帧直接丢弃，避免误报。实际效果取决于你的数据分布——训练集里的图片背景、光照是否接近摄像头画面，如果差得远，实际准确率会明显打折。我自己的经验是这种场景下最好采集一部分摄像头真实画面混入训练集，不然就是「训练一时爽，部署火葬场」的翻车现场。从那以后我每次做图像分类项目，都会强制走一遍「真实部署环境采集几十张图看看效果」的验证流程，希望帮到你。
本文还有配套的精品资源，点击获取
