尧图建网站 尧图建网站 YAOTU WEB BUILD 免费咨询
ARTICLE DETAIL

资讯详情

深耕网站建设与建站编程的一线实战洞察。

基于TensorFlow的垃圾分类图像识别系统实战:从训练到部署

基于TensorFlow的垃圾分类图像识别系统实战:从训练到部署 简介图像分类是计算机视觉的基础任务其核心在于利用卷积神经网络自动提取图像特征并判断物体类别。借助迁移学习开发者可以在小规模数据集上快速训练出高精度模型大幅降低从零训练的时间与算力成本。该技术已广泛应用于智慧城市、环保监管等场景尤其适合解决垃圾分类中的自动识别难题。围绕基于TensorFlow/Keras与MobileNetV2的轻量级图像分类方案完整展示了从数据采集、增强、迁移学习训练到TFLite部署的完整流程并针对数据不均衡、过拟合、推理部署等常见问题给出工程化解决思路帮助开发者快速构建可落地的垃圾分类识别系统。1. 项目概述这个系统到底要解决什么问题有一次朋友发我一张照片问我这到底算可回收还是干垃圾我盯着照片看了半天也没敢下结论。后来我干脆写了一个图像识别的垃圾分类系统用摄像头一拍模型直接告诉我类别。没想到这个赶鸭子上架的项目后来成了很多人来问代码和做法的对象。这套系统的核心其实不复杂输入一张垃圾图片通过卷积神经网络识别出它属于哪一类再输出对应的分类提示。放在实际场景里就是垃圾桶旁边装个摄像头或者用户拿手机拍一张照片系统立刻告诉他“这是可回收物应该丢蓝色桶”。它能解决的核心问题有两个一是普通居民对分类标准不熟悉经常拿着一样东西犹豫半天二是社区督导员人手有限不可能每个垃圾桶旁边都站一个人。用图像识别来做初筛至少能覆盖大部分常见垃圾把高频问题自动化。这个项目很适合三类人参考计算机视觉方向的学生用来做课设或毕业设计需要完整代码但不想从零造轮子的开发者以及社区、学校、园区想做垃圾分类小试点的运营者。如果你只是刚接触深度学习也不用慌我会把数据准备、模型训练、推理部署和常见坑全部拆开讲照着跑就能出一个可用版本。下面所有代码我都用 Python TensorFlow/Keras 实现这也是目前快速落地图像分类任务最省心的一条路。2. 整体设计思路与方案选型2.1 图像分类和目标检测怎么选很多第一次做这个项目的人都会纠结我是不是该用 YOLO 或者 Faster R-CNN 来做目标检测我的建议是先想清楚你的真实使用场景。图像分类解决的是“这张图里主要物体是什么类别”的问题图像检测解决的是“图里有哪些物体、每个物体在哪、分别是什么类别”的问题。垃圾分类这个场景里如果你是在投放点安装一个固定摄像头每次用户把一个物品放到识别区域再拍照那么画面里通常只有一个主体用分类模型完全够而且训练成本低、推理速度快、数据需求量小。如果你做的是“机器人分拣线”传送带上一堆垃圾混在一起必须知道每个物品的位置和类别那才需要目标检测。我一开始就踩过这个坑上来想用 YOLOv5 做检测结果光是标注数据就花了两周训练出来的框还经常飘后来回到分类模型两天就上线了。所以在动手之前先问自己一个问题我的摄像头一次要看几个对象一个就分类多个再检测。2.2 技术栈与模型选型技术栈选择上我选 TensorFlow/Keras 而不是 PyTorch不是因为 PyTorch 不好而是因为这个项目追求“快速落地”和“部署方便”。Keras 高层 API 写起来很直观TensorFlow 的 SavedModel 和 TFLite 转换也很成熟后面接安卓端、树莓派都方便。模型方面我优先推荐 MobileNetV2。它是轻量级网络的代表原始参数量小在 CPU 上也能跑得动推理一张图通常只要几十毫秒在树莓派上虽然会慢一些但也能接受。如果你对精度要求更高机器配置也够可以换 EfficientNetB0如果部署平台性能非常有限再轻一点的 MobileNetV3 也可以。我这里用 MobileNetV2 做例子因为它在精度和速度之间最均衡。这里有一个核心技巧不要从头训练要用 ImageNet 预训练权重做迁移学习。原因很简单垃圾图片的数据集通常不大从头训练一个深网络很容易过拟合而 ImageNet 预训练模型已经学会了边缘、纹理、形状等通用特征我们只需要在它基础上微调最后几层就能用很少的数据达到不错的效果。这个思路在整个项目里是最关键的决定了你后面能省多少事。模型参数量CPU 推理速度场景建议MobileNetV2约 3.4M快安卓、树莓派、实时摄像头EfficientNetB0约 5.3M中等效果优先的服务端ResNet50约 25.6M较慢数据量较大且不关注延迟2.3 数据集的获取与类别体系做垃圾分类系统最容易被忽视的是数据但最决定上限的也是数据。公开数据集我推荐两个方向TrashNet 这类国外数据集通常只有玻璃、纸张、纸板、塑料、金属、厨余垃圾六类国内也有华为云等平台开放过垃圾分类数据集类别可能到四十多类。用这些数据集起步没问题但要注意它们不一定符合你所在地的四分类标准。我们国内常见分类是可回收物、有害垃圾、厨余垃圾、其他垃圾。公开数据集往往是按材质分类的比如“塑料”“玻璃”使用前需要做一个映射。比如干净的塑料瓶映射到可回收物但是沾了油污的一次性塑料餐盒应该映射到其他垃圾。这个映射工作看起来不起眼实际直接决定上线效果因为同一材质在不同污染程度下归属完全不同。如果公开数据不够可以自己补充拍摄。我的经验是每个类别至少准备 500 张以上图片1000 张更好并且要有意识地覆盖不同背景、不同角度、不同光照、不同干净程度。拍摄时不要只拍“标准样本”脏的、皱的、破的都要有否则模型在真实环境里会认不出来。这一步可以叫“脏数据采集”听着不好听但它是模型真正能用的关键。3. 手把手实现环境、数据与训练代码3.1 环境准备我建议用 Conda 新建一个干净的环境避免和系统 Python 打架。Python 版本选 3.9TensorFlow 用 2.10 左右比较稳定。如果你用的是 Apple Silicon 的 Mac可以装 tensorflow-macos和下面的代码差别不大如果是 Windows 或者 Linux直接按常规安装。conda create -n trash_ai python3.9 -y conda activate trash_ai pip install tensorflow2.10 opencv-python matplotlib numpy1.23.5 scikit-learn这里有三个容易出问题的点。第一numpy 版本不要随便升TensorFlow 2.10 对新版 numpy 的兼容性经常有问题装 1.23.5 是我试过最稳的第二opencv-python 装这一个就够了不需要再装 opencv-contrib-python否则容易出现依赖冲突第三如果你要用 GPU需要提前装好 CUDA 和 cuDNNTensorFlow 2.10 对应的是 CUDA 11.2这个版本敏感度很高装错了会直接报找不到 libcudart.so。3.2 数据预处理与增强我把数据整理成下面的目录结构每一类一个文件夹后面用 Keras 的flow_from_directory自动读取省去自己写 Dataset 的麻烦。dataset/ train/ recyclable/ 001.jpg kitchen/ 002.jpg harmful/ 003.jpg other/ 004.jpg val/ recyclable/ 005.jpg kitchen/ 006.jpg harmful/ 007.jpg other/ 008.jpg数据增强是必须做的。垃圾图片在真实场景里会出现旋转、偏移、缩放如果模型没见过这些变化换个角度就认不出来了。我用的增强参数如下from tensorflow.keras.preprocessing.image import ImageDataGenerator train_datagen ImageDataGenerator( rescale1.0 / 255, rotation_range20, width_shift_range0.2, height_shift_range0.2, shear_range0.2, zoom_range0.2, horizontal_flipTrue, validation_split0.2 ) val_datagen ImageDataGenerator(rescale1.0 / 255) train_generator train_datagen.flow_from_directory( dataset/train, target_size(224, 224), batch_size32, class_modecategorical, subsettraining, shuffleTrue, seed42 ) val_generator train_datagen.flow_from_directory( dataset/train, target_size(224, 224), batch_size32, class_modecategorical, subsetvalidation, shuffleFalse, seed42 )注意验证集一定不要做增强只做归一化。很多人图省事直接把同一个增强器用在验证集上结果验证准确率看起来很高但实际是模型见过这些“变形版本”的原因这属于数据泄漏。rescale1.0/255是把像素值从 0-255 缩放到 0-1这是大多数预训练模型的输入要求。旋转和缩放角度我控制在 20 度以内因为垃圾在摄像头下不会出现夸张的 90 度翻转太强的增强反而会让模型学到错误的东西。3.3 迁移学习训练代码主体模型用 MobileNetV2加载 ImageNet 权重后删掉顶层接一个全局平均池化层、一个 Dropout、一个全连接分类层。Dropout 在这里非常重要它防止全连接层过拟合尤其是当垃圾数据量不够大的时候。import tensorflow as tf from tensorflow.keras.applications import MobileNetV2 from tensorflow.keras.layers import GlobalAveragePooling2D, Dense, Dropout from tensorflow.keras.models import Sequential class_names list(train_generator.class_indices.keys()) print(类别映射, train_generator.class_indices) base_model MobileNetV2( input_shape(224, 224, 3), include_topFalse, weightsimagenet ) base_model.trainable False model Sequential([ base_model, GlobalAveragePooling2D(), Dropout(0.2), Dense(len(class_names), activationsoftmax) ]) model.compile( optimizertf.keras.optimizers.Adam(learning_rate0.001), losscategorical_crossentropy, metrics[accuracy] ) history model.fit( train_generator, validation_dataval_generator, epochs10, callbacks[ tf.keras.callbacks.ModelCheckpoint( trash_best.h5, save_best_onlyTrue, monitorval_accuracy ), tf.keras.callbacks.EarlyStopping( monitorval_loss, patience3, restore_best_weightsTrue ) ] )第一轮我会把所有基础层冻结只训练最后的分类层这样训练速度很快也能避免一开始就破坏预训练特征。跑大概 10 个 epoch看到验证准确率开始平台之后再解冻基础层的后面一部分用小学习率微调。base_model.trainable True for layer in base_model.layers[:100]: layer.trainable False model.compile( optimizertf.keras.optimizers.Adam(learning_rate1e-5), losscategorical_crossentropy, metrics[accuracy] ) model.fit( train_generator, validation_dataval_generator, epochs5 )这里的学习率一定要从 0.001 降到 1e-5因为解冻后所有卷积层的权重都会参与更新学习率太大很容易把预训练学到的特征冲掉导致损失突然暴涨。我见过很多同学在这一步翻车跑出来的结果还不如不解冻就是因为学习率没降。3.4 模型导出与转换训练完成后我把模型同时存成 H5 和 TFLite 两种格式。H5 用于本地 Python 快速验证TFLite 用于安卓端和边缘设备。model.save(trash_classifier.h5) converter tf.lite.TFLiteConverter.from_keras_model(model) tflite_model converter.convert() with open(trash_classifier.tflite, wb) as f: f.write(tflite_model) print(导出完成)TFLite 转换有个小坑如果你的模型中包含一些自定义层或者复杂的算子转换可能报错。遇到这种情况可以先试试把converter.optimizations设置为默认优化或者降级算子的 target spec。我在 MobileNetV2 上转换很顺利所以这也是我推荐它的原因之一。4. 推理部署与效果调优4.1 单张图片识别脚本模型训练好之后先不要急着接摄像头先用单张图片验证一下流程。这里的核心是推理时的预处理必须和训练时完全一致尺寸、缩放、归一化差一个都可能让结果崩掉。import numpy as np from tensorflow.keras.preprocessing import image from tensorflow.keras.models import load_model model load_model(trash_classifier.h5) class_names [recyclable, kitchen, harmful, other] def predict_image(img_path): img image.load_img(img_path, target_size(224, 224)) img_array image.img_to_array(img) img_array np.expand_dims(img_array, axis0) / 255.0 pred model.predict(img_array, verbose0)[0] idx np.argmax(pred) confidence float(pred[idx]) return class_names[idx], confidence print(predict_image(test.jpg))这里建议把置信度也打印出来。如果最高置信度不到 0.6我建议直接输出“无法确定”让用户人工判断。这在真实场景里比硬猜一个结果要可靠得多垃圾分错桶的代价比“识别不了”大得多。4.2 实时摄像头识别单张图片验证通过后再接摄像头。OpenCV 从摄像头取帧缩放后送进模型把识别结果直接画在画面上。这样在投放点或者实验室里就可以实时演示。import cv2 import numpy as np from tensorflow.keras.models import load_model model load_model(trash_classifier.h5) class_names [recyclable, kitchen, harmful, other] cap cv2.VideoCapture(0) if not cap.isOpened(): print(无法打开摄像头) exit() while True: ret, frame cap.read() if not ret: break img cv2.resize(frame, (224, 224)) img img.astype(float32) / 255.0 img np.expand_dims(img, axis0) preds model.predict(img, verbose0)[0] idx np.argmax(preds) label class_names[idx] conf float(preds[idx]) cv2.putText( frame, f{label}: {conf:.2f}, (20, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2 ) cv2.imshow(Trash Classifier, frame) if cv2.waitKey(1) 0xFF ord(q): break cap.release() cv2.destroyAllWindows()这个实时脚本里有一个容易被忽视的点model.predict在每一帧都调用会有点慢如果摄像头画面卡顿可以降低到每 3 帧识别一次或者把推理放到单独的线程。另外建议在识别前先判断画面中心区域是否有物体变化避免摄像头一直对着空背景反复输出同一个结果。4.3 关键参数调整与准确率提升技巧训练垃圾分类模型准确率卡在某个点不上去是常态。我的调优顺序一般是先看训练集准确率再看验证集准确率判断是欠拟合还是过拟合。如果训练集准确率就已经很低说明模型容量不够或者数据问题。可以尝试把 MobileNetV2 换成 EfficientNetB0或者增加训练轮数。如果训练集准确率很高但验证集准确率低说明过拟合优先增加 Dropout、增加数据增强而不是换更大的模型。类别不均衡在这个项目里特别常见。比如“其他垃圾”的图片可能只有 300 张“可回收物”却有 2000 张模型会偏向样本多的类别。解决方式是在fit里加上class_weight。from sklearn.utils.class_weight import compute_class_weight class_counts np.bincount(train_generator.classes) class_weight compute_class_weight( balanced, classesnp.unique(train_generator.classes), ytrain_generator.classes ) class_weight_dict dict(enumerate(class_weight)) model.fit( train_generator, validation_dataval_generator, epochs10, class_weightclass_weight_dict )学习率的选择也很关键。我习惯用 Adam 默认 0.001 起步如果前 3 个 epoch 损失完全没有下降就把学习率除以 10 重试。微调阶段固定用 1e-5这个值几乎不用动太大会毁掉预训练特征太小则收敛慢。4.4 向安卓端和其他平台扩展如果要把系统放到安卓手机上TFLite 是首选。把转换好的trash_classifier.tflite放到安卓工程的 assets 目录然后用 CameraX 获取摄像头预览帧对每一帧做缩放和推理再把分类结果叠加到预览画面上这就是很多人说的“窗口图像识别”的雏形。TFLite 的推理接口比 Keras 原始模型稍麻烦一点需要手动定义输入输出张量但好在 MobileNetV2 的输入输出结构很固定官方文档里的示例可以直接套用。如果你没做过安卓开发也可以用现成的 MediaPipe Tasks 封装它内置了 TFLite 图像分类的 API只需要配置模型路径和标签文件就能跑起来。如果是 Web 端我更推荐 TensorFlow.js它可以直接在浏览器里跑模型不需要后端服务。不过移动端浏览器初次加载模型会比较慢建议用tfjs-converter把模型转成 tfjs 格式并做量化压缩模型体积可以从十几 MB 降到三四 MB。5. 常见问题与排查技巧实录5.1 训练不收敛或准确率波动大这个是最多人问的问题。训练时损失不下降或者验证准确率忽高忽低我从经验里总结了几个排查方向现象可能原因解决办法loss 一直不降学习率太大或太小先试 0.001不降再降为 0.0001准确率在 0% 和 25% 之间跳数据标签错乱随机抽几张训练图检查验证准确率比训练高很多数据泄漏确认同一物品的照片没有同时出现在训练和验证集训练后期验证准确率下降过拟合增加 Dropout增加增强减少 epoch我最常遇到的其实是数据泄漏。很多人从网上下载一个数据集没有做随机划分或者同一个垃圾瓶子的多张照片被分到了不同文件夹模型相当于“背答案”训练的时候看着很高一到真实拍摄就露馅。所以数据划分一定要在类别文件夹层面做随机同时尽量保证同一来源的照片不跨集合。5.2 内存溢出和训练中断训练时报 OOM尤其是不小心开了 GPU 版又没控制显存很常见。最简单的办法是把 batch size 从 32 降到 16甚至降到 8。Batch size 减半显存占用差不多也减半这是性价比最高的操作。另一个办法是不要一次性把所有数据加载进内存。用flow_from_directory的好处就是边读边训练不会占太多内存。如果你自己写了 Dataset一定要检查是不是把所有图片都读成了 numpy 数组放到内存里。我之前帮人排查过一个训练卡死的问题最终发现是 2 万张图片一次性 load 到内存直接占了几十 GB。5.3 类别数据不均衡垃圾分类数据天然不平衡因为“其他垃圾”的样本多且杂“有害垃圾”样本少。除了前面说的class_weight我还会做针对性增强。对样本少的类别可以多生成几份旋转、缩放、亮度变化之后的副本。但这里要注意增强副本不能放到验证集里否则验证结果会虚高。如果某个类别实在太少比如少于 200 张我会考虑是否暂时去掉这一类先跑一个三分类系统。与其让模型硬学一个数据不够的类不如先保证已有类别的准确率后面再逐步补数据。这一步看起来很保守但实际效果比硬凑四分类好得多。5.4 识别结果与预期不符模型在测试集上准确率不错拿到真实环境里却经常错这是部署阶段最常见的挫败感来源。我总结有三类原因第一摄像头拍摄角度和你训练数据的角度差异大第二真实垃圾的脏污、破损程度远超训练集第三两个类别外观非常相似比如沾了水的纸盒和废纸。解决思路不是一味加数据而是先错例分析。把跑错的图片收集起来按类别归档看看模型到底把什么认成了什么。我做过一次错例分析发现很多“其他垃圾”被认成“可回收物”再排查发现是因为训练集中可回收物的背景颜色比较鲜亮模型把背景当成了线索而不是物体本身。后来我在训练集中增加了大量不同背景的图片问题才缓解。5.5 标签映射和分类标准争议上面提到过同一个物品在不同状态下的分类结果可能完全不同。干净的塑料瓶是可回收物但沾满油污的外卖盒属于其他垃圾椰子壳虽然是厨余垃圾但它的硬壳又不太好处理大骨头属于其他垃圾小骨头属于厨余垃圾。图像识别模型很难判断这些语义层面的规则因为模型看到的是整体外观它不知道这个物体“是否被污染”。所以我的建议是系统输出不要只给一个“类别名称”最好附上一句投放提示比如“其他垃圾请丢灰桶”。同时如果识别置信度偏低界面直接显示“建议人工确认”而不是强行给结果。这套逻辑在真实场景里比追求 99% 准确率更重要。6. 项目扩展方向与我的实操体会6.1 从分类走向检测换用 YOLO当你发现分类模型已经不够用比如要同时识别传送带上的多个物品就该换 YOLO 了。数据准备方面需要用 LabelImg 或 CVAT 标注每个物体的边界框和类别格式转成 YOLO 的 labels 文件。训练可以用 YOLOv5 或 YOLOv8官方仓库自带预训练权重迁移学习的思路和分类模型一样。不过我要提醒一点从分类切到检测工作量不是一个量级。分类每类只需要一批图片检测需要逐张标注框1000 张图片一个人标下来至少一两天。如果是做毕业设计或者小试点先用分类模型把流程跑通再考虑检测是我最推荐的节奏。6.2 Web端和小程序端的落地思路如果想让别人不用装 Python 环境也能用可以做一个简单的 Web 页面。最简单的方案是用 Flask 写一个 POST 接口接收图片调用训练好的 Keras 模型返回分类结果。代码量不大但要注意接口加的图片大小限制避免超大图直接打爆内存。更轻量的方案是转成 TensorFlow.js放到静态网站上。模型在浏览器本地运行图片不需要上传服务器隐私性更好也没有部署成本。缺点就是模型首次加载时间长手机端特别明显量化压缩能缓解一部分。如果你要做的只是一个演示 Demo我建议优先考虑 TensorFlow.js 方案。6.3 几个让我少踩坑的习惯做完这个项目我最深的体会是数据质量永远比模型结构重要。我试过把 MobileNetV2 换成更大的 ResNet50准确率只提升了一点后来把训练集里所有模糊、标签错误的图片清掉准确率直接涨了好几个百分点。模型的性能上限很大程度在数据不在网络结构。最后一个很实用的小技巧每次训练前先拿一个 batch 数据跑一次前向和反向确保 loss 能正常下降再正式启动训练。这个小习惯帮我避免过很多次白等几小时最后发现数据流程有 bug 的尴尬。这个项目后续想继续玩还可以加入语音提示、识别记录统计、传感器联动开门方向很多但先把基础的图像识别链路跑通后面每一步都会顺很多。本文还有配套的精品资源点击获取
返回列表