从零构建垃圾分类识别系统：基于8万张图片与TensorFlow的实战指南

1. 项目背景与数据集介绍

垃圾分类识别系统听起来高大上，但其实离我们日常生活特别近。去年我帮小区物业做了一套这样的系统，从零开始折腾了两个月，踩了不少坑，也积累了不少实战经验。这次就用8万张图片的数据集为例，带大家走完全流程。

这个数据集是我从多个开源渠道整理来的，包含245类常见垃圾，比如"厨余垃圾_鸡蛋壳"、"可回收物_矿泉水瓶"这样的细分类别。原始数据质量参差不齐，有些图片模糊不清，有些甚至根本不是垃圾图片。我花了整整一周时间清洗数据，最终得到80012张可用图片，全部统一转换成jpg格式，按文件夹分类存放。比如"trash_jpg/厨余垃圾_香蕉皮"这个路径下，就全是香蕉皮的照片。

数据集有四个大类：

• 可回收物（塑料瓶、纸箱等）
• 厨余垃圾（果皮、剩饭等）
• 有害垃圾（电池、药品等）
• 其他垃圾（卫生纸、塑料袋等）

每个大类下又有几十到上百个小类，比如"可回收物"下面还分"可回收物_玻璃瓶"、"可回收物_易拉罐"等。这种层级结构特别适合用TensorFlow的image_dataset_from_directory方法直接读取，省去了手动标注的麻烦。

2. 开发环境搭建

工欲善其事，必先利其器。我推荐用Anaconda创建独立的Python环境，避免包版本冲突。这是我的环境配置清单：

conda create -n trash_classify python=3.8 conda activate trash_classify pip install tensorflow==2.3.0 pillow matplotlib opencv-python pyqt5

这里有个坑要注意：TensorFlow 2.3虽然不算最新版，但经过实测发现它对MobileNet的支持最稳定。我试过用2.6版本，训练时会出现莫名其妙的NaN loss问题。

硬件方面，最好有块NVIDIA显卡。我用的是GTX 1660 Ti 6GB显存，训练MobileNet大约需要4小时。如果没有显卡，可以用Google Colab的免费GPU资源，记得选择运行时类型为GPU就行。

3. 数据预处理实战

拿到8万张图片后，千万别直接扔给模型训练。我总结了一套预处理"组合拳"：

3.1 数据增强策略

在tf.keras里可以很方便地实现数据增强：

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, validation_split=0.2 # 直接划分20%数据做验证集 )

这里有个技巧：rotation_range不要设太大，否则"可回收物_易拉罐"旋转后可能被误认为"其他垃圾_罐头"。我一开始设了40度，准确率直接掉了5%。

3.2 数据集加载

用TF的API加载数据简直不要太方便：

train_ds = train_datagen.flow_from_directory( 'trash_jpg', target_size=(224, 224), batch_size=32, class_mode='categorical', subset='training' ) val_ds = train_datagen.flow_from_directory( 'trash_jpg', target_size=(224, 224), batch_size=32, class_mode='categorical', subset='validation' )

注意target_size要和模型输入尺寸一致。MobileNet默认是224x224，如果你用EfficientNet就要调整。

4. 模型训练与调优

4.1 MobileNet迁移学习

直接上代码：

base_model = tf.keras.applications.MobileNetV2( input_shape=(224, 224, 3), include_top=False, weights='imagenet' ) # 冻结基础模型 base_model.trainable = False # 添加自定义分类头 model = tf.keras.Sequential([ base_model, tf.keras.layers.GlobalAveragePooling2D(), tf.keras.layers.Dense(256, activation='relu'), tf.keras.layers.Dropout(0.5), tf.keras.layers.Dense(245, activation='softmax') ])

这里有几个关键点：

1. 先用include_top=False去掉原模型的分类头
2. 开始训练时要冻结基础模型(trainable=False)
3. 中间加了个Dropout层防止过拟合，这个参数我调了十几次才确定0.5最合适

4.2 训练技巧

我的训练配置是这样的：

model.compile( optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001), loss='categorical_crossentropy', metrics=['accuracy'] ) history = model.fit( train_ds, validation_data=val_ds, epochs=30, callbacks=[ tf.keras.callbacks.EarlyStopping(patience=3), tf.keras.callbacks.ModelCheckpoint('best_model.h5') ] )

重点来了：

• 学习率一定要小，0.0001起步
• 用了EarlyStopping防止过拟合
• ModelCheckpoint会保存验证集上表现最好的模型

在我的机器上，30个epoch大概要4小时。最终验证集准确率能达到82%左右，对于245个类别来说已经很不错了。

5. 模型部署与应用开发

5.1 PyQt5界面开发

训练好的模型要落地应用，我选择了PyQt5做图形界面。核心代码结构：

class MainWindow(QMainWindow): def __init__(self): super().__init__() self.model = tf.keras.models.load_model('best_model.h5') self.class_names = [...] # 245个类别的名称 def classify_image(self): img = Image.open(self.file_path) img = img.resize((224, 224)) img_array = tf.keras.preprocessing.image.img_to_array(img) img_array = tf.expand_dims(img_array, 0) predictions = self.model.predict(img_array) score = tf.nn.softmax(predictions[0]) result = f"分类结果: {self.class_names[np.argmax(score)]}\n置信度: {100*np.max(score):.2f}%" self.result_label.setText(result)

界面布局用Qt Designer拖拽完成，主要包含：

• 图片上传按钮
• 结果显示区域
• 分类按钮
• 退出按钮

5.2 性能优化技巧

在实际部署时发现几个问题：

1. 模型加载慢：改用tf.lite转换后速度提升3倍
2. 内存占用高：添加了图片尺寸检查，超过5MB的图片先压缩
3. 界面卡顿：把预测任务放到单独线程中执行

转换TensorFlow Lite模型的代码：

converter = tf.lite.TFLiteConverter.from_keras_model(model) tflite_model = converter.convert() with open('model.tflite', 'wb') as f: f.write(tflite_model)

6. 常见问题与解决方案

6.1 类别不平衡问题

数据集中"其他垃圾"类别的图片特别多，导致模型偏向预测为这类。我试过三种解决方法：

1. 类权重法：在model.fit中添加class_weight参数
2. 过采样法：用ImageDataGenerator对少数类图片做更多增强
3. 欠采样法：随机删除多数类图片

最终发现方法1效果最好，代码如下：

from sklearn.utils import class_weight import numpy as np class_weights = class_weight.compute_class_weight( 'balanced', classes=np.unique(train_ds.classes), y=train_ds.classes ) class_weights = dict(enumerate(class_weights)) model.fit(..., class_weight=class_weights)

6.2 新类别增量学习

后来小区新增了"可回收物_奶茶杯"这个类别，我摸索出一套增量学习方案：

1. 冻结原模型所有层
2. 只训练新添加的分类头
3. 用小学习率(0.00001)微调几层卷积

这样既不用重新训练整个模型，又能快速适应新类别。

7. 进阶优化方向

如果想进一步提升准确率，可以尝试：

1. 改用EfficientNetV2模型，我的测试显示能提升3-5%准确率
2. 添加注意力机制模块
3. 使用CutMix数据增强
4. 对模糊图片添加超分辨率预处理

不过要提醒的是，模型复杂度增加会显著延长训练时间。在实际项目中，要在准确率和推理速度之间做好权衡。