本实验为做一个垃圾分类系统，需要识别可回收垃圾、厨余垃圾、有害垃圾和其他垃圾等四个大类。
import tensorflow as tf
import zipfile
import os
import random
from shutil import copy2
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from PIL import Image
import numpy as np
import cv2
import matplotlib.pyplot as plt
%matplotlib inline解压图片数据集
if not os.path.exists('Waste'):
with zipfile.ZipFile('./Waste.zip') as c:
c.extractall()
# 定义data_set_split函数来实现数据拆分
def data_set_split(src_data_folder, target_data_folder, train_scale=0.8, val_scale=0.2, test_scale=0.0):
'''
读取源数据文件夹，生成划分好的文件夹，分为trian、val、test三个文件夹进行
:param src_data_folder: 源文件夹 data_split/src_data
:param target_data_folder: 目标文件夹data_split/target_data
:param train_scale: 训练集比例
:param val_scale: 验证集比例
:param test_scale: 测试集比例
:return:
'''
print("开始数据集划分")
class_names = os.listdir(src_data_folder)
# 在目标目录下创建文件夹
split_names = ['train', 'val', 'test']
for split_name in split_names:
split_path = os.path.join(target_data_folder, split_name)
if os.path.isdir(split_path):
pass
else:
os.mkdir(split_path)
# 然后在split_path的目录下创建类别文件夹
for class_name in class_names:
class_split_path = os.path.join(split_path, class_name)
if os.path.isdir(class_split_path):
pass
else:
os.mkdir(class_split_path)
# 按照比例划分数据集，并进行数据图片的复制
# 首先进行分类遍历
for class_name in class_names:
current_class_data_path = os.path.join(src_data_folder, class_name)
current_all_data = os.listdir(current_class_data_path)
current_data_length = len(current_all_data)
current_data_index_list = list(range(current_data_length))
random.shuffle(current_data_index_list)
train_folder = os.path.join(os.path.join(target_data_folder, 'train'), class_name)
val_folder = os.path.join(os.path.join(target_data_folder, 'val'), class_name)
test_folder = os.path.join(os.path.join(target_data_folder, 'test'), class_name)
train_stop_flag = current_data_length * train_scale
val_stop_flag = current_data_length * (train_scale + val_scale)
current_idx = 0
train_num = 0
val_num = 0
test_num = 0
for i in current_data_index_list:
src_img_path = os.path.join(current_class_data_path, current_all_data[i])
if current_idx <= train_stop_flag:
copy2(src_img_path, train_folder)
# print("{}复制到了{}".format(src_img_path, train_folder))
train_num = train_num + 1
elif (current_idx > train_stop_flag) and (current_idx <= val_stop_flag):
copy2(src_img_path, val_folder)
# print("{}复制到了{}".format(src_img_path, val_folder))
val_num = val_num + 1
else:
copy2(src_img_path, test_folder)
# print("{}复制到了{}".format(src_img_path, test_folder))
test_num = test_num + 1
current_idx = current_idx + 1
print("*********************************{}*************************************".format(class_name))
print(
"{}类按照{}：{}：{}的比例划分完成，一共{}张图片".format(class_name, train_scale, val_scale, test_scale, current_data_length))
print("训练集{}：{}张".format(train_folder, train_num))
print("验证集{}：{}张".format(val_folder, val_num))
print("测试集{}：{}张".format(test_folder, test_num))# 设置原始数据集目录及分割之后的目录，并调用上述函数实现数据拆分
if not os.path.exists('datasets'):
os.mkdir("datasets")
src_data_folder = "Waste" # todo 原始数据集目录
target_data_folder = "datasets" # todo 数据集分割之后存放的目录
data_set_split(src_data_folder, target_data_folder)# 设置原始数据集目录及分割之后的目录，并调用上述函数实现数据拆分
base = './datasets/'
# 训练集数据生成器:数据增强
train_datagen = ImageDataGenerator(
rescale=1./255,
rotation_range=40,
width_shift_range=0.2,
height_shift_range=0.2,
shear_range=0.2,
zoom_range=0.2,
horizontal_flip=True,
fill_mode='nearest'
)
# 验证集数据生成器,不需要进行增强
validation_datagen = ImageDataGenerator(rescale=1./255)
# 从本地路径加载训练集、验证集
train_generator = train_datagen.flow_from_directory(
base+"train/",
target_size=(224,224),
batch_size=16,
shuffle=True
)
val_generator = train_datagen.flow_from_directory(
base+"val/",
target_size=(224,224),
class_mode='categorical'
)
# 查看数据加载器的生成的batch数据
img,label = None,None
print(train_generator.class_indices, len(train_generator.class_indices))
for da in train_generator:
print(type(da), len(da))
print(type(da[0]), type(da[1]))
print(da[0].shape, da[1].shape)
img,label = da[0], da[1]
break
# 显示垃圾图像
plt.imshow(img[0])
# 搭建LeNet卷积神经网络模型
def lenet_model_load(IMG_SHAPE=(224, 224, 3), class_num=245):
# 通过keras构建模型
model = tf.keras.models.Sequential([
tf.keras.layers.experimental.preprocessing.Rescaling(1. / 255, input_shape=IMG_SHAPE), # 归一化，将像素值处理成0到1之间的值
tf.keras.layers.Conv2D(32, (3, 3), activation='relu'), # 卷积层，32个输出通道，3*3的卷积核，激活函数为relu
tf.keras.layers.MaxPooling2D(2, 2), # 池化层，特征图大小减半
# Add another convolution
tf.keras.layers.Conv2D(64, (3, 3), activation='relu'), # 卷积层，32个输出通道，3*3的卷积核，激活函数为relu
tf.keras.layers.MaxPooling2D(2, 2), # 池化层，特征图大小减半
tf.keras.layers.Flatten(), # 将二维的特征图拉直
# The same 128 dense layers, and 10 output layers as in the pre-convolution example:
tf.keras.layers.Dense(128, activation='relu'), # 128个神经元的全连接层
tf.keras.layers.Dense(class_num, activation='softmax') # 输出层，对应数据集具体的类别数目
])
model.summary() # 输出模型信息
# 模型训练
model.compile(optimizer='adam', loss='categorical_crossentropy',
metrics=['accuracy']) # 编译模型，指定模型的优化器是adam，模型使用的损失函数的交叉熵损失函数
return model # 返回模型
# 模型加载
lenet = lenet_model_load()# MobileNet迁移学习
# 模型加载，指定图片处理的大小和是否进行迁移学习
def mobile_model_load(class_num=245, IMG_SHAPE=(224, 224, 3)):
# 微调的过程中不需要进行归一化的处理
base_model = tf.keras.applications.MobileNetV2(input_shape=(224,224,3),
include_top=False,
weights='imagenet') # 加载mobilenetv2模型
base_model.trainable = False # 将主干的特征提取网络参数进行冻结
model = tf.keras.models.Sequential([
tf.keras.layers.experimental.preprocessing.Rescaling(1. / 127.5, offset=-1, input_shape=(224,224,3)),
# 归一化处理，将像素值处理为-1到1之间
base_model,
tf.keras.layers.GlobalAveragePooling2D(), # 全局平均池化
tf.keras.layers.Dense(class_num, activation='softmax') # 设置最后的全连接层，用于分类
])
model.summary() # 输出模型信息
# 模型训练
model.compile(optimizer='adam', loss='categorical_crossentropy',
metrics=['accuracy']) # 用于编译模型，指定模型的优化器是adam优化器，模型的损失函数是交叉熵损失函数
return model # 返回模型
mobile = mobile_model_load(4)# 展示训练过程的曲线
def show_loss_acc(history):
acc = history.history['accuracy']
val_acc = history.history['val_accuracy']
loss = history.history['loss']
val_loss = history.history['val_loss']
# 绘制准确率曲线图
plt.figure(figsize=(8, 8))
plt.subplot(2, 1, 1)
plt.plot(acc, label='Training Accuracy')
plt.plot(val_acc, label='Validation Accuracy')
plt.legend(loc='lower right')
plt.ylabel('Accuracy')
plt.ylim([min(plt.ylim()), 1])
plt.title('Training and Validation Accuracy')
# 绘制模型误差曲线图
plt.subplot(2, 1, 2)
plt.plot(loss, label='Training Loss')
plt.plot(val_loss, label='Validation Loss')
plt.legend(loc='upper right')
plt.ylabel('Cross Entropy')
plt.title('Training and Validation Loss')
plt.xlabel('epoch')
plt.savefig('results_mobilenet_epoch30.png', dpi=100)# 模型训练
epochs=30
model = mobile_model_load(4) # 加载模型
history = model.fit_generator(train_generator, validation_data=val_generator, epochs=epochs) # 开始训练
os.mkdir("models") # 模型保存路径
model.save("models/mobilenet_245_epoch30.h5") # todo 修改了自己的模型保存位置，保存模型
show_loss_acc(history) # 展示曲线图# 定义类别名称
# 垃圾类名
class_names = ['其他垃圾_PE塑料袋', '其他垃圾_U型回形针', '其他垃圾_一次性杯子', '其他垃圾_一次性棉签', '其他垃圾_串串竹签', '其他垃圾_便利贴', '其他垃圾_创可贴',
'其他垃圾_卫生纸', '其他垃圾_厨房手套', '其他垃圾_厨房抹布', '其他垃圾_口罩', '其他垃圾_唱片', '其他垃圾_图钉', '其他垃圾_大龙虾头',
'其他垃圾_奶茶杯', '其他垃圾_干燥剂', '其他垃圾_彩票', '其他垃圾_打泡网', '其他垃圾_打火机', '其他垃圾_搓澡巾', '其他垃圾_果壳', '其他垃圾_毛巾',
'其他垃圾_涂改带', '其他垃圾_湿纸巾', '其他垃圾_烟蒂', '其他垃圾_牙刷', '其他垃圾_电影票', '其他垃圾_电蚊香', '其他垃圾_百洁布', '其他垃圾_眼镜',
'其他垃圾_眼镜布', '其他垃圾_空调滤芯', '其他垃圾_笔', '其他垃圾_胶带', '其他垃圾_胶水废包装', '其他垃圾_苍蝇拍', '其他垃圾_茶壶碎片',
'其他垃圾_草帽', '其他垃圾_菜板', '其他垃圾_车票', '其他垃圾_酒精棉', '其他垃圾_防霉防蛀片', '其他垃圾_除湿袋', '其他垃圾_餐巾纸',
'其他垃圾_餐盒', '其他垃圾_验孕棒', '其他垃圾_鸡毛掸', '厨余垃圾_八宝粥', '厨余垃圾_冰激凌', '厨余垃圾_冰糖葫芦', '厨余垃圾_咖啡',
'厨余垃圾_圣女果', '厨余垃圾_地瓜', '厨余垃圾_坚果', '厨余垃圾_壳', '厨余垃圾_巧克力', '厨余垃圾_果冻', '厨余垃圾_果皮', '厨余垃圾_核桃',
'厨余垃圾_梨', '厨余垃圾_橙子', '厨余垃圾_残渣剩饭', '厨余垃圾_水果', '厨余垃圾_泡菜', '厨余垃圾_火腿', '厨余垃圾_火龙果', '厨余垃圾_烤鸡',
'厨余垃圾_瓜子', '厨余垃圾_甘蔗', '厨余垃圾_番茄', '厨余垃圾_秸秆杯', '厨余垃圾_秸秆碗', '厨余垃圾_粉条', '厨余垃圾_肉类', '厨余垃圾_肠',
'厨余垃圾_苹果', '厨余垃圾_茶叶', '厨余垃圾_草莓', '厨余垃圾_菠萝', '厨余垃圾_菠萝蜜', '厨余垃圾_萝卜', '厨余垃圾_蒜', '厨余垃圾_蔬菜',
'厨余垃圾_薯条', '厨余垃圾_薯片', '厨余垃圾_蘑菇', '厨余垃圾_蛋', '厨余垃圾_蛋挞', '厨余垃圾_蛋糕', '厨余垃圾_豆', '厨余垃圾_豆腐',
'厨余垃圾_辣椒', '厨余垃圾_面包', '厨余垃圾_饼干', '厨余垃圾_鸡翅', '可回收物_不锈钢制品', '可回收物_乒乓球拍', '可回收物_书', '可回收物_体重秤',
'可回收物_保温杯', '可回收物_保鲜膜内芯', '可回收物_信封', '可回收物_充电头', '可回收物_充电宝', '可回收物_充电牙刷', '可回收物_充电线',
'可回收物_凳子', '可回收物_刀', '可回收物_包', '可回收物_单车', '可回收物_卡', '可回收物_台灯', '可回收物_吊牌', '可回收物_吹风机',
'可回收物_呼啦圈', '可回收物_地球仪', '可回收物_地铁票', '可回收物_垫子', '可回收物_塑料制品', '可回收物_太阳能热水器', '可回收物_奶粉桶',
'可回收物_尺子', '可回收物_尼龙绳', '可回收物_布制品', '可回收物_帽子', '可回收物_手机', '可回收物_手电筒', '可回收物_手表', '可回收物_手链',
'可回收物_打包绳', '可回收物_打印机', '可回收物_打气筒', '可回收物_扫地机器人', '可回收物_护肤品空瓶', '可回收物_拉杆箱', '可回收物_拖鞋',
'可回收物_插线板', '可回收物_搓衣板', '可回收物_收音机', '可回收物_放大镜', '可回收物_日历', '可回收物_暖宝宝', '可回收物_望远镜',
'可回收物_木制切菜板', '可回收物_木桶', '可回收物_木棍', '可回收物_木质梳子', '可回收物_木质锅铲', '可回收物_木雕', '可回收物_枕头',
'可回收物_果冻杯', '可回收物_桌子', '可回收物_棋子', '可回收物_模具', '可回收物_毯子', '可回收物_水壶', '可回收物_水杯', '可回收物_沙发',
'可回收物_泡沫板', '可回收物_灭火器', '可回收物_灯罩', '可回收物_烟灰缸', '可回收物_热水瓶', '可回收物_燃气灶', '可回收物_燃气瓶',
'可回收物_玩具', '可回收物_玻璃制品', '可回收物_玻璃器皿', '可回收物_玻璃壶', '可回收物_玻璃球', '可回收物_瑜伽球', '可回收物_电动剃须刀',
'可回收物_电动卷发棒', '可回收物_电子秤', '可回收物_电熨斗', '可回收物_电磁炉', '可回收物_电脑屏幕', '可回收物_电视机', '可回收物_电话',
'可回收物_电路板', '可回收物_电风扇', '可回收物_电饭煲', '可回收物_登机牌', '可回收物_盒子', '可回收物_盖子', '可回收物_盘子', '可回收物_碗',
'可回收物_磁铁', '可回收物_空气净化器', '可回收物_空气加湿器', '可回收物_笼子', '可回收物_箱子', '可回收物_纸制品', '可回收物_纸牌',
'可回收物_罐子', '可回收物_网卡', '可回收物_耳套', '可回收物_耳机', '可回收物_衣架', '可回收物_袋子', '可回收物_袜子', '可回收物_裙子',
'可回收物_裤子', '可回收物_计算器', '可回收物_订书机', '可回收物_话筒', '可回收物_豆浆机', '可回收物_路由器', '可回收物_轮胎', '可回收物_过滤网',
'可回收物_遥控器', '可回收物_量杯', '可回收物_金属制品', '可回收物_钉子', '可回收物_钥匙', '可回收物_铁丝球', '可回收物_铅球',
'可回收物_铝制用品', '可回收物_锅', '可回收物_锅盖', '可回收物_键盘', '可回收物_镊子', '可回收物_闹铃', '可回收物_雨伞', '可回收物_鞋',
'可回收物_音响', '可回收物_餐具', '可回收物_餐垫', '可回收物_饰品', '可回收物_鱼缸', '可回收物_鼠标', '有害垃圾_指甲油', '有害垃圾_杀虫剂',
'有害垃圾_温度计', '有害垃圾_灯', '有害垃圾_电池', '有害垃圾_电池板', '有害垃圾_纽扣电池', '有害垃圾_胶水', '有害垃圾_药品包装', '有害垃圾_药片',
'有害垃圾_药瓶', '有害垃圾_药膏', '有害垃圾_蓄电池', '有害垃圾_血压计']# 加载训练好的模型
model = tf.keras.models.load_model("mobilenet_245_epoch30.h5")
# 模型预测
def predict_img(model, imgpath):
# 加载图像并转换为Numpy数组
img = Image.open(imgpath)
img = img.convert("RGB")
img = np.array(img)
# 调整图像尺寸
img = cv2.resize(img, (224, 224))
# 增加一个维度 (1, 224, 224, 3)
img = img[np.newaxis, :]
print(img.shape)
# 转换类型
img = img.astype("float32")
# 模型推理（预测）
pred = model(img)
# 返回概率最大值所在的下标，即可得到类别
print(np.argmax(pred), class_names[np.argmax(pred)])# 选择一张图像
imgpath = "datasets/val/Harmful_Waste/img_19538.jpg"
def showimg(imgpath):
img = Image.open(imgpath)
img = img.convert("RGB")
img = np.array(img)
plt.imshow(img)
showimg(imgpath)
# 调用函数进行预测
predict_img(model, imgpath)
# 测试另外一张图像
imgpath = "datasets/val/Kitchen_Waste/img_2299.jpg"
showimg(imgpath)
predict_img(model, imgpath)
1万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
