TensorFlow+OpenCV实战：垃圾分类图像分类模型训练与预测全流程

简介这份资源面向图像分类入门者与深度学习实践者提供一套基于简单垃圾分类数据集的完整智能分类方案帮助读者理解从数据准备到模型预测的全流程。包内共1046个文件以1041张jpg图片构成训练与测试数据集另含2个Python脚本分别负责训练与预测、1个h5模型文件、1个txt说明及1段mp4演示视频压缩包约824.68MB。已有1332人学习下载适合作为图像分类模板参考。读者可借助train.py完成模型训练通过predict.py对输入图片进行推理并在图片上以中文标注干垃圾、湿垃圾、可回收垃圾、有害垃圾四类结果同时了解数据集制作、TensorFlow与OpenCV结合使用的具体做法快速搭建自己的分类实验环境。1. 从一张垃圾桶照片说起这套垃圾分类模型到底能跑出什么结果你拍一张外卖盒的照片丢给模型它告诉你这是「干垃圾」还是「可回收物」——这件事听起来像是个玩具 demo但真正动手做一遍你会发现它把图像分类的完整链路全串起来了数据采集、标注、增强、搭网络、训练、调参、预测、可视化。这套资源就是干这个的核心文件是my_model.h5和一组测试图img_11548.jpg、img_11351.jpg、img_4491.jpg等配套train.py负责训练、predict.py负责推理输出四个类别干垃圾、湿垃圾、可回收垃圾、有害垃圾。它不追求 SOTA 精度定位是「图像分类模板」——你可以把数据集换成自己的网络结构不动跑通整条流水线。适合刚接触 TensorFlow 和 OpenCV 的从业者拿来练手也适合需要快速搭一个分类 baseline 的人直接改。2. 拆开这个 h5 文件模型结构、数据流与四分类逻辑2.1 为什么是 TensorFlow OpenCV 这套组合垃圾分类本质上是细粒度图像分类的一个简化版。干垃圾和湿垃圾在视觉上的差异往往集中在纹理和边缘——比如纸巾的纤维感 vs 果皮的湿润反光。OpenCV 在这里承担的是预处理角色读图、缩放、颜色空间转换、直方图均衡这些操作比直接用 PIL 更可控尤其是当你需要做形态学处理膨胀、腐蚀来强化边缘特征时OpenCV 的cv2.dilate和cv2.erode是现成的。TensorFlow 这边my_model.h5是一个已经保存权重的 Keras 模型文件。h5 格式的好处是结构和权重打包在一起tf.keras.models.load_model(my_model.h5)一行就能恢复整个网络不需要你重新定义层。常见做法是底层用几个卷积块提特征后面接全局平均池化再连全连接层输出四类 softmax。这种结构参数量小在几千张图的规模上不容易过拟合。选 TensorFlow 而不是 PyTorch 的理由很实际h5 格式在 Keras 生态里加载最省事而且predict.py里用model.predict()拿到概率向量后直接np.argmax就能出类别索引再映射到中文标签。如果你习惯 PyTorch也可以把 h5 转成 ONNX 再加载但那是另一条路这套资源没走。2.2 数据从哪来、怎么进网络垃圾分类数据集通常按文件夹组织每个类别一个子目录dataset/ ├── gan/ │ ├── img_001.jpg │ └── ... ├── shi/ │ ├── img_002.jpg │ └── ... ├── kehuishou/ │ └── ... └── youhai/ └── ...train.py里一般用ImageDataGenerator做流式读取好处是不用一次性把几千张图全塞进内存。核心参数这么设from tensorflow.keras.preprocessing.image import ImageDataGenerator train_datagen ImageDataGenerator( rescale1./255, # 像素归一化到 [0,1] rotation_range20, # 随机旋转 ±20 度 width_shift_range0.1, # 水平平移 10% height_shift_range0.1, # 垂直平移 10% horizontal_flipTrue, # 水平翻转 validation_split0.2 # 划出 20% 做验证 ) train_generator train_datagen.flow_from_directory( dataset/, target_size(224, 224), # 统一缩放到 224x224 batch_size32, class_modecategorical, subsettraining )rescale必须做否则像素值在 0-255 之间梯度更新会不稳定。target_size设成 224×224 是因为大多数预训练骨干网络比如 MobileNetV2的默认输入就是这个尺寸即使你从零搭网络这个尺寸在精度和显存之间也比较平衡。validation_split划出的验证集不参与梯度更新只用来监控过拟合。flow_from_directory会自动根据子目录名生成类别索引顺序是字母序。所以gan、kehuishou、shi、youhai对应的索引可能是 0、1、2、3但具体顺序要在训练后打印train_generator.class_indices确认否则预测时标签会错位。2.3 训练脚本里几个不能省的步骤train.py的骨架大致是这样import tensorflow as tf from tensorflow.keras import layers, models model models.Sequential([ layers.Conv2D(32, (3,3), activationrelu, input_shape(224,224,3)), layers.MaxPooling2D((2,2)), layers.Conv2D(64, (3,3), activationrelu), layers.MaxPooling2D((2,2)), layers.Conv2D(128, (3,3), activationrelu), layers.MaxPooling2D((2,2)), layers.GlobalAveragePooling2D(), layers.Dense(128, activationrelu), layers.Dropout(0.5), layers.Dense(4, activationsoftmax) # 四个类别 ]) model.compile( optimizeradam, losscategorical_crossentropy, metrics[accuracy] ) history model.fit( train_generator, epochs30, validation_dataval_generator ) model.save(my_model.h5)GlobalAveragePooling2D替代Flatten是为了减少参数量降低过拟合风险。Dropout(0.5)在全连接层前随机丢弃一半神经元也是防过拟合的常规操作。categorical_crossentropy对应 one-hot 标签如果你用class_modesparse就得换成sparse_categorical_crossentropy这两个不能混。epochs30是个经验值实际跑的时候要看验证集准确率什么时候不再上升。如果验证 loss 开始往上走而训练 loss 还在降就是过拟合了该早停。可以加EarlyStopping回调callbacks [ tf.keras.callbacks.EarlyStopping(patience5, restore_best_weightsTrue) ]patience5表示验证 loss 连续 5 个 epoch 不改善就停restore_best_weights把权重回滚到最好的那一轮。2.4 预测脚本怎么把中文标签画到图上predict.py要做三件事加载模型、读图推理、把结果写到图片上。中文显示是个坑OpenCV 的cv2.putText不支持中文直接写会变成问号。常见做法是用 PIL 画字import cv2 import numpy as np from PIL import Image, ImageDraw, ImageFont from tensorflow.keras.models import load_model model load_model(my_model.h5) class_names [干垃圾, 湿垃圾, 可回收垃圾, 有害垃圾] img cv2.imread(img_11548.jpg) img_resized cv2.resize(img, (224, 224)) img_array img_resized / 255.0 img_array np.expand_dims(img_array, axis0) preds model.predict(img_array) idx np.argmax(preds) label class_names[idx] confidence preds[0][idx] img_rgb cv2.cvtColor(img, cv2.COLOR_BGR2RGB) pil_img Image.fromarray(img_rgb) draw ImageDraw.Draw(pil_img) font ImageFont.truetype(simhei.ttf, 36) # 需要中文字体文件 draw.text((10, 10), f{label} {confidence:.2f}, fontfont, fill(255, 0, 0)) result cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR) cv2.imwrite(result.jpg, result)simhei.ttf是黑体字体文件Windows 系统在C:/Windows/Fonts/下能找到Linux 上需要自己拷一份或者用fc-list找可用中文字体。ImageFont.truetype的第二个参数是字号图片分辨率高的时候要相应调大否则字太小看不清。np.expand_dims那一步不能省model.predict要求输入是 batch 形式即使你只预测一张图也得凑出一个维度。preds[0]才是这张图的四类概率np.argmax返回最大概率的索引和class_names的顺序对应。3. 从零跑通训练到预测环境、命令与参数调整3.1 环境装什么、版本怎么选TensorFlow 2.x 是必须的因为 h5 加载和 Keras API 都依赖它。Python 版本建议 3.8 到 3.10太新的版本可能和 TensorFlow 的 wheel 不兼容。一条命令装齐pip install tensorflow opencv-python pillow numpy matplotlib如果你有 NVIDIA 显卡可以装tensorflow-gpu但要注意 CUDA 和 cuDNN 版本必须和 TensorFlow 版本匹配。常见坑是 TensorFlow 2.10 之后 GPU 支持合并进了主包不再单独发tensorflow-gpu装错了会报Could not find cudart64_*.dll。不确定的话先用 CPU 版跑通流程再折腾 GPU。验证安装import tensorflow as tf print(tf.__version__) print(tf.config.list_physical_devices(GPU))如果 GPU 列表为空但你有显卡检查 CUDA 路径是否加进了系统环境变量。3.2 训练命令与日志解读假设数据集已经按类别分好文件夹训练直接跑python train.py --data_dir dataset/ --epochs 30 --batch_size 32 --output my_model.h5如果你的train.py没写 argparse那就改脚本里的硬编码路径。训练过程中终端会打印每个 epoch 的 loss 和 accuracyEpoch 1/30 100/100 [] - 45s 450ms/step - loss: 1.2345 - accuracy: 0.4567 - val_loss: 1.0123 - val_accuracy: 0.5234重点看val_accuracy和val_loss。如果val_accuracy在 0.7 左右就上不去了可能是数据量不够或者类别不平衡。垃圾分类数据集里「可回收垃圾」的样本往往最多「有害垃圾」最少这会导致模型偏向多数类。解决办法是给flow_from_directory加class_weight参数或者对少数类做过采样。3.3 预测单张图和批量预测单张预测python predict.py --image img_11548.jpg --model my_model.h5批量预测就是把测试图全跑一遍import os test_images [img_11548.jpg, img_11351.jpg, img_4491.jpg, img_11382.jpg, img_11376.jpg, img_48.jpg, img_4343.jpg, img_4276.jpg, img_11566.jpg] for img_path in test_images: result predict_single(img_path) print(f{img_path}: {result})批量跑的时候注意内存如果图片很多别一次性全读进来用循环逐张处理。另外预测前要确保图片路径正确cv2.imread读不到文件不会报错会返回None后面 resize 就会崩。加个判断if img is None: print(f无法读取 {img_path}) continue3.4 参数怎么调学习率、batch size、输入尺寸学习率是影响最大的超参数。Adam 默认 0.001如果训练 loss 震荡厉害降到 0.0001如果 loss 下降太慢可以试 0.005 但别更高否则容易发散。batch size 受显存限制CPU 训练用 16 或 32 都行GPU 显存 4G 以上可以上 64。输入尺寸 224×224 是标准值如果你要识别更细的纹理比如区分不同材质的塑料可以提到 256×256 或 299×299但训练时间会成倍增加。数据增强的强度也要看情况。rotation_range20对垃圾分类是合理的因为垃圾不会总是正着放。但如果你做的是工业质检类的图像分类旋转增强可能反而不合理因为产品方向是固定的。horizontal_flipTrue对大多数自然图像没问题但如果是文字相关的分类翻转会改变语义得关掉。4. 避坑与排查中文乱码、标签错位、过拟合4.1 预测结果中文显示成方块或问号现象predict.py跑完图片上只有方框或者???看不到「干垃圾」这些字。原因OpenCV 的cv2.putText只支持 ASCII 字符中文字符不在它的字体渲染范围内。即使你传了中文字符串它也会用默认字体渲染成乱码。解决改用 PIL 的ImageDraw.text并且指定一个支持中文的 TTF 字体文件。ImageFont.truetype(simhei.ttf, 36)里的字体路径要写对Linux 上可能是/usr/share/fonts/truetype/...Windows 上是C:/Windows/Fonts/simhei.ttf。如果找不到字体用fc-list :langzh查一下系统里有哪些中文字体。4.2 预测标签和实际类别对不上现象模型预测「可回收垃圾」的图输出却是「有害垃圾」但置信度还挺高。原因class_names列表的顺序和训练时flow_from_directory生成的class_indices不一致。flow_from_directory按文件夹名的字母序分配索引比如gan0, kehuishou1, shi2, youhai3但你在predict.py里写的class_names [干垃圾, 湿垃圾, 可回收垃圾, 有害垃圾]对应的索引是干0, 湿1, 可回收2, 有害3完全错位了。解决训练完打印train_generator.class_indices把那个字典保存下来预测时按同样的映射关系取标签。或者直接在predict.py里用class_indices的反向映射class_indices {gan: 0, kehuishou: 1, shi: 2, youhai: 3} idx_to_class {v: k for k, v in class_indices.items()} label_en idx_to_class[idx] label_cn {gan: 干垃圾, shi: 湿垃圾, kehuishou: 可回收垃圾, youhai: 有害垃圾}[label_en]4.3 训练准确率很高但预测一塌糊涂现象训练集 accuracy 到 0.98验证集也有 0.95但拿新图片预测结果随机跳。原因过拟合加上数据泄漏。如果验证集是从训练集里随机切的而训练集里同一张图的不同增强版本同时出现在训练和验证中验证准确率就是虚高的。另外如果数据集里同一类别的图片高度相似比如都是从同一个视频里截的帧模型学到的是背景而不是物体本身。解决确保验证集和训练集在图片级别完全隔离最好按拍摄批次或来源划分。增强操作只在训练集上做验证集只做rescale。如果数据量太少考虑用预训练权重做迁移学习把 MobileNetV2 的前面层冻结只训练最后的分类层。4.4 h5 模型加载报错现象load_model(my_model.h5)抛出ValueError: Unknown layer或OSError: Unable to open file。原因第一种情况是模型里用了自定义层或自定义损失函数加载时没有传custom_objects。第二种是文件路径不对或者文件损坏。解决如果是自定义层加载时加custom_objects{MyLayer: MyLayer}。如果是文件问题检查文件大小是否正常训练保存时是否完整写入了。另外TensorFlow 2.15 之后推荐用.keras格式替代 h5如果你用的是新版本可以另存为.keras再加载。4.5 显存不够导致训练中断现象训练到一半报ResourceExhaustedError: OOM when allocating tensor。原因batch size 太大或者输入尺寸太大超出了 GPU 显存。解决先把 batch size 减半如果还不行就减输入尺寸。另外可以在训练前设置显存按需增长gpus tf.config.experimental.list_physical_devices(GPU) if gpus: for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True)这样 TensorFlow 不会一次性占满所有显存而是按需分配。5. 把模板变成你自己的分类器换数据、调网络、验效果这套资源最大的价值不是垃圾分类本身而是它提供了一个可复用的图像分类骨架。你只需要替换dataset/下的文件夹和图片改一下class_names和最后Dense层的输出单元数就能迁移到其他四分类甚至 N 分类任务。比如做塑料瓶、玻璃瓶、纸盒、金属罐的材质分类流程完全一样。换数据时要注意几点。第一每个类别的图片数量尽量均衡差距不要超过 3 倍否则用class_weight补偿。第二图片的拍摄条件要一致如果训练集全是白底产品图预测时丢一张复杂背景的图模型大概率翻车。第三类别定义要互斥别出现一张图既像干垃圾又像可回收垃圾的情况标注时就要定好规则。网络结构也可以按需调整。如果类别之间差异很细微比如不同种类的塑料可以把卷积层加深或者把Dense(128)改成Dense(256)。如果数据量只有几百张反而要把网络变浅否则过拟合压不住。一个实用的技巧是先用预训练模型提特征把MobileNetV2(weightsimagenet, include_topFalse)的输出接一个全局池化再连自己的分类层只训练最后几层。这样即使数据少也能拿到不错的精度。验证模型效果不能只看准确率。做一个混淆矩阵看看哪两个类别最容易混from sklearn.metrics import confusion_matrix import seaborn as sns y_pred model.predict(val_generator) y_pred_classes np.argmax(y_pred, axis1) y_true val_generator.classes cm confusion_matrix(y_true, y_pred_classes) sns.heatmap(cm, annotTrue, fmtd, xticklabelsclass_names, yticklabelsclass_names)如果「干垃圾」和「湿垃圾」之间的误判特别多说明这两个类别的视觉特征区分度不够要么加更多有区分力的样本要么在预处理阶段强化纹理特征比如用 OpenCV 做局部二值模式或者边缘检测把结果作为额外通道输入。我自己的习惯是每次换数据集后先跑 5 个 epoch 看验证准确率能不能到 0.6 以上。如果 5 轮下来还在 0.3 左右说明数据或标签有问题别急着调参先回去检查文件夹结构和图片内容。从那以后我每次拿到新数据都强制走一遍「随机抽 20 张肉眼过一遍」的流程省得训练半天才发现标签错了。希望帮到你。本文还有配套的精品资源点击获取