导入模块
import tensorflow as tf
from tensorflow.keras.models import Sequential
from sklearn.model_selection import train_test_split,cross_val_score
from tensorflow.keras.layers import Dense,Dropout
from tensorflow.keras.layers import Flatten #扁平化处理工具
#导入卷积和池化
from tensorflow.keras.layers import Convolution2D,MaxPooling2D
from tensorflow.keras.optimizers import Adam,SGD
import os
import numpy as np
将垃圾图片读取成数组添加到列表中，将垃圾分类标签添加到列表中
feature = []
target = []
class_names_to_ids = {'cardboard': 0, 'glass': 1, 'metal': 2, 'paper':3, 'plastic':4, 'trash':5}
data_dir = 'dataset/'
for class_name in class_names_to_ids.keys():
images_list = os.listdir(data_dir + class_name)
for image_name in images_list:
all_path = data_dir + class_name + '/' + image_name
img_arr = plt.imread(all_path)
feature.append(img_arr)
target.append(class_names_to_ids[class_name])
将列表转换为数组
feature = np.array(feature)
target = np.array(target)
划分数据集
x_train,x_test,y_train,y_test = train_test_split(feature,target,test_size=0.1,random_state=2020)
将训练集和测试集的特征数据做归一化处理，将训练集和测试集的标签转为独热编码
#特征数据做归一化处理
x_train = x_train.reshape(-1,28,28,1)/255.0
x_test = x_test.reshape(-1,28,28,1)/255.0
# 把训练集和测试集的标签转为独热编码
y_train = tf.keras.utils.to_categorical(y_train,num_classes=10)
y_test = tf.keras.utils.to_categorical(y_test,num_classes=10)
定义顺序模型
# 定义顺序模型
model = Sequential()
第一个卷积层
# 第一个卷积层
# input_shape 输入数据
# filters 滤波器个数 32，生成 32 张特征图
# kernel_size 卷积窗口大小 5*5
# strides 步长 1
# padding padding 方式 same/valid
# activation 激活函数
model.add(
Convolution2D(
input_shape = (384,512,3),
filters = 32,
kernel_size = 5,
strides = 1,
padding = 'same',
activation = 'relu' )
)
第一个池化层
# 第一个池化层
# pool_size 池化窗口大小 2*2
# strides 步长 2
# padding padding 方式 same/valid
model.add(
MaxPooling2D(
pool_size = 2,
strides = 2,
padding = 'same',)
)
第二个卷积层
# 第二个卷积层
# filters 滤波器个数 64，生成 64 张特征图
# kernel_size 卷积窗口大小 5*5
# strides 步长 1
# padding padding 方式 same/valid
# activation 激活函数
model.add(
Convolution2D(64,5,strides=1,padding='same',activation='relu')
)
第二个池化层
# 第二个池化层
# pool_size 池化窗口大小 2*2
# strides 步长 2
# padding padding 方式 same/valid
model.add(MaxPooling2D(2,2,'same'))
把第二个池化层的输出进行数据扁平化
# 把第二个池化层的输出进行数据扁平化
model.add(Flatten())
第一个全连接层
# 第一个全连接层
model.add(Dense(1024,activation = 'relu'))
# Dropout关闭部分神经元，防止过拟合
model.add(Dropout(0.5))
第二个全连接层，输出为6类垃圾的概率值
# 第二个全连接层
model.add(Dense(6,activation='softmax'))
定义优化器，模型训练及保存
# 定义优化器
adam = Adam(lr=1e-4)#该优化器是其他优化器的集大成者，较为通用
# 定义优化器，loss function，训练过程中计算准确率
model.compile(
optimizer=adam,
loss='categorical_crossentropy',
metrics=['accuracy'])
# 训练模型
# 一个epoch指代所有的数据送入网络中完成一次前向计算及反向传播的过程，
# 所谓Batch就是每次送入网络中训练的一部分数据，而Batch Size就是每个batch中训练样本的数量
model.fit(x_train,y_train,batch_size=64,epochs=10,validation_data=(x_test, y_test))
# 保存模型
model.save('mnist.h5')
image.png
