基于Tensorflow的垃圾分类系统 自训练CNN算法模型与mobileNet迁移学习,可用于图片或视频垃圾分类。 模型经过训练,具有识别能力,可以对视频画面中的物品进行分类,包括4个分类大类,训练集有八万多张图像,训练集验证准确率为91—95%,验证集准确率85—90%,具有较高的准确率。 该模型可以应用于智能垃圾桶、智能回收站等场景中,提高垃圾分类效率。
最近在捣鼓一个基于TensorFlow的垃圾分类系统,感觉挺有意思的,分享一下我的思路和代码。这个系统主要是通过自训练的CNN算法模型和mobileNet迁移学习来实现的,能够对图片或视频中的垃圾进行分类。模型训练了八万多张图像,训练集的准确率在91%到95%之间,验证集的准确率也有85%到90%,效果还不错。
模型架构
首先,我们来看一下模型的架构。我选择了一个自训练的CNN模型,同时也尝试了mobileNet的迁移学习。mobileNet是一个轻量级的模型,适合在移动设备上运行,这对于我们想要部署到智能垃圾桶或回收站的场景来说非常合适。
import tensorflow as tf from tensorflow.keras import layers, models def build_cnn_model(): model = models.Sequential([ layers.Conv2D(32, (3, 3), activation='relu', input_shape=(224, 224, 3)), layers.MaxPooling2D((2, 2)), layers.Conv2D(64, (3, 3), activation='relu'), layers.MaxPooling2D((2, 2)), layers.Conv2D(128, (3, 3), activation='relu'), layers.MaxPooling2D((2, 2)), layers.Flatten(), layers.Dense(128, activation='relu'), layers.Dense(4, activation='softmax') # 4个分类大类 ]) return model # mobileNet迁移学习 def build_mobilenet_model(): base_model = tf.keras.applications.MobileNetV2(input_shape=(224, 224, 3), include_top=False, weights='imagenet') base_model.trainable = False model = models.Sequential([ base_model, layers.GlobalAveragePooling2D(), layers.Dense(128, activation='relu'), layers.Dense(4, activation='softmax') # 4个分类大类 ]) return model数据准备
接下来是数据准备部分。我们使用了八万多张图像进行训练,这些图像涵盖了四个主要的垃圾分类类别。为了增强模型的泛化能力,我还对数据进行了增强处理,比如随机旋转、缩放、翻转等。
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, fill_mode='nearest' ) train_generator = train_datagen.flow_from_directory( 'path_to_train_data', target_size=(224, 224), batch_size=32, class_mode='categorical' )模型训练
模型训练部分,我使用了Adam优化器和交叉熵损失函数。为了监控训练过程,我还添加了准确率和损失值的回调函数。
model = build_mobilenet_model() model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy']) history = model.fit( train_generator, epochs=10, validation_data=validation_generator )模型评估
训练完成后,我对模型进行了评估。训练集的准确率在91%到95%之间,验证集的准确率在85%到90%之间,效果还是相当不错的。
loss, accuracy = model.evaluate(test_generator) print(f'Test accuracy: {accuracy:.2f}')应用场景
这个模型可以应用在智能垃圾桶、智能回收站等场景中。通过摄像头实时捕捉画面,模型可以快速识别垃圾类别,并自动分类。这不仅提高了垃圾分类的效率,还能减少人工分类的错误率。
总结
总的来说,这个基于TensorFlow的垃圾分类系统还是挺实用的。通过自训练的CNN模型和mobileNet迁移学习,我们能够达到较高的分类准确率。未来还可以考虑进一步优化模型,比如增加更多的训练数据,或者尝试其他更高效的模型架构。
好了,今天就先聊到这里,如果你对这个项目感兴趣,欢迎在评论区交流讨论!