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

1. 项目背景与数据集介绍垃圾分类识别系统听起来高大上但其实离我们日常生活特别近。去年我帮小区物业做了一套这样的系统从零开始折腾了两个月踩了不少坑也积累了不少实战经验。这次就用8万张图片的数据集为例带大家走完全流程。这个数据集是我从多个开源渠道整理来的包含245类常见垃圾比如厨余垃圾_鸡蛋壳、可回收物_矿泉水瓶这样的细分类别。原始数据质量参差不齐有些图片模糊不清有些甚至根本不是垃圾图片。我花了整整一周时间清洗数据最终得到80012张可用图片全部统一转换成jpg格式按文件夹分类存放。比如trash_jpg/厨余垃圾_香蕉皮这个路径下就全是香蕉皮的照片。数据集有四个大类可回收物塑料瓶、纸箱等厨余垃圾果皮、剩饭等有害垃圾电池、药品等其他垃圾卫生纸、塑料袋等每个大类下又有几十到上百个小类比如可回收物下面还分可回收物_玻璃瓶、可回收物_易拉罐等。这种层级结构特别适合用TensorFlow的image_dataset_from_directory方法直接读取省去了手动标注的麻烦。2. 开发环境搭建工欲善其事必先利其器。我推荐用Anaconda创建独立的Python环境避免包版本冲突。这是我的环境配置清单conda create -n trash_classify python3.8 conda activate trash_classify pip install tensorflow2.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( rescale1./255, rotation_range20, width_shift_range0.2, height_shift_range0.2, shear_range0.2, zoom_range0.2, horizontal_flipTrue, validation_split0.2 # 直接划分20%数据做验证集 )这里有个技巧rotation_range不要设太大否则可回收物_易拉罐旋转后可能被误认为其他垃圾_罐头。我一开始设了40度准确率直接掉了5%。3.2 数据集加载用TF的API加载数据简直不要太方便train_ds train_datagen.flow_from_directory( trash_jpg, target_size(224, 224), batch_size32, class_modecategorical, subsettraining ) val_ds train_datagen.flow_from_directory( trash_jpg, target_size(224, 224), batch_size32, class_modecategorical, subsetvalidation )注意target_size要和模型输入尺寸一致。MobileNet默认是224x224如果你用EfficientNet就要调整。4. 模型训练与调优4.1 MobileNet迁移学习直接上代码base_model tf.keras.applications.MobileNetV2( input_shape(224, 224, 3), include_topFalse, weightsimagenet ) # 冻结基础模型 base_model.trainable False # 添加自定义分类头 model tf.keras.Sequential([ base_model, tf.keras.layers.GlobalAveragePooling2D(), tf.keras.layers.Dense(256, activationrelu), tf.keras.layers.Dropout(0.5), tf.keras.layers.Dense(245, activationsoftmax) ])这里有几个关键点先用include_topFalse去掉原模型的分类头开始训练时要冻结基础模型(trainableFalse)中间加了个Dropout层防止过拟合这个参数我调了十几次才确定0.5最合适4.2 训练技巧我的训练配置是这样的model.compile( optimizertf.keras.optimizers.Adam(learning_rate0.0001), losscategorical_crossentropy, metrics[accuracy] ) history model.fit( train_ds, validation_dataval_ds, epochs30, callbacks[ tf.keras.callbacks.EarlyStopping(patience3), 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 性能优化技巧在实际部署时发现几个问题模型加载慢改用tf.lite转换后速度提升3倍内存占用高添加了图片尺寸检查超过5MB的图片先压缩界面卡顿把预测任务放到单独线程中执行转换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 类别不平衡问题数据集中其他垃圾类别的图片特别多导致模型偏向预测为这类。我试过三种解决方法类权重法在model.fit中添加class_weight参数过采样法用ImageDataGenerator对少数类图片做更多增强欠采样法随机删除多数类图片最终发现方法1效果最好代码如下from sklearn.utils import class_weight import numpy as np class_weights class_weight.compute_class_weight( balanced, classesnp.unique(train_ds.classes), ytrain_ds.classes ) class_weights dict(enumerate(class_weights)) model.fit(..., class_weightclass_weights)6.2 新类别增量学习后来小区新增了可回收物_奶茶杯这个类别我摸索出一套增量学习方案冻结原模型所有层只训练新添加的分类头用小学习率(0.00001)微调几层卷积这样既不用重新训练整个模型又能快速适应新类别。7. 进阶优化方向如果想进一步提升准确率可以尝试改用EfficientNetV2模型我的测试显示能提升3-5%准确率添加注意力机制模块使用CutMix数据增强对模糊图片添加超分辨率预处理不过要提醒的是模型复杂度增加会显著延长训练时间。在实际项目中要在准确率和推理速度之间做好权衡。