使用resnet进行简单的垃圾分类（tensorflow2）

三个臭皮姜

已于 2022-07-12 22:58:09 修改

阅读量4.9k

点赞数 18

分类专栏： tensorflow学习笔记 python 文章标签： tensorflow 分类 python

于 2021-06-07 22:55:57 首次发布

本文链接：https://blog.csdn.net/Aiden_yan/article/details/117675851

版权

python 同时被 2 个专栏收录

37 篇文章 3 订阅

订阅专栏

tensorflow学习笔记

13 篇文章 1 订阅

订阅专栏

一.说明

我的环境：

Windows10系统
Pycharm
tensorflow2.3.0

代码目录结构：
在这里插入图片描述

数据集按类别分开成单独的文件夹，每个文件夹下存储对应的数据集图片：
在这里插入图片描述

各文件夹说明：
1.datasets：数据集所在位置
2.model：训练之后生成的模型所在位置
3.pretrain_model：存放着resnet50的预训练模型
4.test：单独进行预测的图片
5.predict.py：调用生成的h5模型进行预测
6.train.py：进行训练

二.训练

1.导入所需的库

import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np
from tensorflow.keras import layers, models
import pathlib
from keras.layers import Input, Activation, BatchNormalization, Flatten
from keras.layers import Dense, Conv2D, MaxPooling2D, ZeroPadding2D, AveragePooling2D
from keras.models import Model

2.设置GPU

gpus = tf.config.list_physical_devices("GPU")

if gpus:
    tf.config.experimental.set_memory_growth(gpus[0], True)  # 设置GPU显存用量按需使用
    tf.config.set_visible_devices([gpus[0]], "GPU")

3.一些准备

# 支持中文
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

# 设置随机种子尽可能使结果可以重现
np.random.seed(1)
# 设置随机种子尽可能使结果可以重现
tf.random.set_seed(1)

4.导入数据集并检查图片数量

# 数据集路径
data_dir = "./datasets"
data_dir = pathlib.Path(data_dir)
image_count = len(list(data_dir.glob('*/*')))
print("图片总数为：", image_count)

5.设置参数

batch_size = 8
img_height = 224
img_width = 224

6.划分训练集和验证集

# 训练集
train_ds = tf.keras.preprocessing.image_dataset_from_directory(
    data_dir,
    validation_split=0.2,
    subset="training",
    seed=123,
    image_size=(img_height, img_width),
    batch_size=batch_size)

# 验证集
val_ds = tf.keras.preprocessing.image_dataset_from_directory(
    data_dir,
    validation_split=0.2,
    subset="validation",
    seed=123,
    image_size=(img_height, img_width),
    batch_size=batch_size)

7.再次检查数据集

# 检查数据
class_names = train_ds.class_names
print(class_names)

for image_batch, labels_batch in train_ds:
    print(image_batch.shape)  # (8, 224, 224, 3)
    print(labels_batch.shape)  # (8,)
    break

8.配置数据集方便计算机训练（非必要）

# 配置数据集
# AUTOTUNE = tf.data.AUTOTUNE
AUTOTUNE = tf.data.experimental.AUTOTUNE
train_ds = train_ds.cache().shuffle(1000).prefetch(buffer_size=AUTOTUNE)
val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)

9.搭建网络模型resnet

# 构建网络模型
def identity_block(input_tensor, kernel_size, filters, stage, block):
    filters1, filters2, filters3 = filters

    name_base = str(stage) + block + '_identity_block_'

    x = Conv2D(filters1, (1, 1), name=name_base + 'conv1')(input_tensor)
    x = BatchNormalization(name=name_base + 'bn1')(x)
    x = Activation('relu', name=name_base + 'relu1')(x)

    x = Conv2D(filters2, kernel_size, padding='same', name=name_base + 'conv2')(x)
    x = BatchNormalization(name=name_base + 'bn2')(x)
    x = Activation('relu', name=name_base + 'relu2')(x)

    x = Conv2D(filters3, (1, 1), name=name_base + 'conv3')(x)
    x = BatchNormalization(name=name_base + 'bn3')(x)

    x = layers.add([x, input_tensor], name=name_base + 'add')
    x = Activation('relu', name=name_base + 'relu4')(x)
    return x


def conv_block(input_tensor, kernel_size, filters, stage, block, strides=(2, 2)):
    filters1, filters2, filters3 = filters

    res_name_base = str(stage) + block + '_conv_block_res_'
    name_base = str(stage) + block + '_conv_block_'

    x = Conv2D(filters1, (1, 1), strides=strides, name=name_base + 'conv1')(input_tensor)
    x = BatchNormalization(name=name_base + 'bn1')(x)
    x = Activation('relu', name=name_base + 'relu1')(x)

    x = Conv2D(filters2, kernel_size, padding='same', name=name_base + 'conv2')(x)
    x = BatchNormalization(name=name_base + 'bn2')(x)
    x = Activation('relu', name=name_base + 'relu2')(x)

    x = Conv2D(filters3, (1, 1), name=name_base + 'conv3')(x)
    x = BatchNormalization(name=name_base + 'bn3')(x)

    shortcut = Conv2D(filters3, (1, 1), strides=strides, name=res_name_base + 'conv')(input_tensor)
    shortcut = BatchNormalization(name=res_name_base + 'bn')(shortcut)

    x = layers.add([x, shortcut], name=name_base + 'add')
    x = Activation('relu', name=name_base + 'relu4')(x)
    return x


def ResNet50(input_shape=[224, 224, 3], classes=1000):
    img_input = Input(shape=input_shape)
    x = ZeroPadding2D((3, 3))(img_input)

    x = Conv2D(64, (7, 7), strides=(2, 2), name='conv1')(x)
    x = BatchNormalization(name='bn_conv1')(x)
    x = Activation('relu')(x)
    x = MaxPooling2D((3, 3), strides=(2, 2))(x)

    x = conv_block(x, 3, [64, 64, 256], stage=2, block='a', strides=(1, 1))
    x = identity_block(x, 3, [64, 64, 256], stage=2, block='b')
    x = identity_block(x, 3, [64, 64, 256], stage=2, block='c')

    x = conv_block(x, 3, [128, 128, 512], stage=3, block='a')
    x = identity_block(x, 3, [128, 128, 512], stage=3, block='b')
    x = identity_block(x, 3, [128, 128, 512], stage=3, block='c')
    x = identity_block(x, 3, [128, 128, 512], stage=3, block='d')

    x = conv_block(x, 3, [256, 256, 1024], stage=4, block='a')
    x = identity_block(x, 3, [256, 256, 1024], stage=4, block='b')
    x = identity_block(x, 3, [256, 256, 1024], stage=4, block='c')
    x = identity_block(x, 3, [256, 256, 1024], stage=4, block='d')
    x = identity_block(x, 3, [256, 256, 1024], stage=4, block='e')
    x = identity_block(x, 3, [256, 256, 1024], stage=4, block='f')

    x = conv_block(x, 3, [512, 512, 2048], stage=5, block='a')
    x = identity_block(x, 3, [512, 512, 2048], stage=5, block='b')
    x = identity_block(x, 3, [512, 512, 2048], stage=5, block='c')

    x = AveragePooling2D((7, 7), name='avg_pool')(x)

    x = Flatten()(x)
    x = Dense(classes, activation='softmax', name='fc1000')(x)
    model = Model(img_input, x, name='resnet50')
    # 加载预训练模型
    model.load_weights("./pretrain_model/resnet50_weights_tf_dim_ordering_tf_kernels.h5")
    return model

10.实例化模型并输出网络结构

model = ResNet50()
model.summary()

11.编译

# 编译
opt = tf.keras.optimizers.Adam(learning_rate=1e-7)

model.compile(optimizer="adam",
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

12.开始训练

epochs = 10  # 训练轮数
history = model.fit(
    train_ds,
    validation_data=val_ds,
    epochs=epochs
)

13.训练结果可视化

# 模型评估可视化
acc = history.history['accuracy']
val_acc = history.history['val_accuracy']

loss = history.history['loss']
val_loss = history.history['val_loss']

epochs_range = range(epochs)

plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)

plt.plot(epochs_range, acc, label='Training Accuracy')
plt.plot(epochs_range, val_acc, label='Validation Accuracy')
plt.legend(loc='lower right')
plt.title('Training and Validation Accuracy')

plt.subplot(1, 2, 2)
plt.plot(epochs_range, loss, label='Training Loss')
plt.plot(epochs_range, val_loss, label='Validation Loss')
plt.legend(loc='upper right')
plt.title('Training and Validation Loss')
plt.show()

14.保存模型

# 保存模型
model.save('model/garbage_model.h5')

训练结果：
在这里插入图片描述

三.调用模型进行预测

1.导入库

from tensorflow.keras.preprocessing import image
from tensorflow.keras.applications.resnet50 import preprocess_input
import numpy as np
from tensorflow.keras.models import load_model

2.加载模型

# 加载模型
model = load_model('./model/garbage_model.h5')

3.设置类别名称

这里按datasets文件夹下默认的顺序来即可，注意不要写错名称。

# 设置类别名称列表
classes_name=['battery','cans', 'cigarette', 'plastic_bottle', 'vegetables']
china_name=["电池","易拉罐","香烟","塑料瓶","蔬菜"]

4.加载要预测的图片并调用模型输出预测结果

# 图片所在路径
img_path = './test/test2.jpg'
# 加载图片
img = image.load_img(img_path, target_size=(224, 224))
# 将图片转成矩阵形式
x = image.img_to_array(img)
# 扩展一维
x = np.expand_dims(x, axis=0)

x = preprocess_input(x)

# 预测
preds = model.predict(x)
max_index = np.argmax(preds, axis=-1)[0]  # 取可能性最大的那个类别

acc=preds[0][int(max_index)]*100

print("英文物体名称是：{}，中文物体名称是：{}，置信度为：{:.5f}%".format(classes_name[max_index],china_name[max_index],acc))

结果（图片来自百度图片）：

在这里插入图片描述
全部代码：
train.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time : 2021/6/7 17:20
# @Author : 若谷
# @File : train.py
# @Software: PyCharm
import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np
from tensorflow.keras import layers, models
import pathlib
from keras.layers import Input, Activation, BatchNormalization, Flatten
from keras.layers import Dense, Conv2D, MaxPooling2D, ZeroPadding2D, AveragePooling2D
from keras.models import Model

# 设置GPU
gpus = tf.config.list_physical_devices("GPU")

if gpus:
    tf.config.experimental.set_memory_growth(gpus[0], True)  # 设置GPU显存用量按需使用
    tf.config.set_visible_devices([gpus[0]], "GPU")

# 导入数据

# 支持中文
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

# 设置随机种子尽可能使结果可以重现
np.random.seed(1)
# 设置随机种子尽可能使结果可以重现
tf.random.set_seed(1)

# 数据集路径
data_dir = "./datasets"
data_dir = pathlib.Path(data_dir)
image_count = len(list(data_dir.glob('*/*')))
print("图片总数为：", image_count)  # 图片总数为：

# 加载数据
batch_size = 8
img_height = 224
img_width = 224

"""
关于image_dataset_from_directory()的详细介绍可以参考文章：https://mtyjkh.blog.csdn.net/article/details/117018789
"""
# 训练集
train_ds = tf.keras.preprocessing.image_dataset_from_directory(
    data_dir,
    validation_split=0.2,
    subset="training",
    seed=123,
    image_size=(img_height, img_width),
    batch_size=batch_size)

# 测试集
val_ds = tf.keras.preprocessing.image_dataset_from_directory(
    data_dir,
    validation_split=0.2,
    subset="validation",
    seed=123,
    image_size=(img_height, img_width),
    batch_size=batch_size)

# 检查数据
class_names = train_ds.class_names
print(class_names)

for image_batch, labels_batch in train_ds:
    print(image_batch.shape)  # (8, 224, 224, 3)
    print(labels_batch.shape)  # (8,)
    break

# 配置数据集
# AUTOTUNE = tf.data.AUTOTUNE
AUTOTUNE = tf.data.experimental.AUTOTUNE
train_ds = train_ds.cache().shuffle(1000).prefetch(buffer_size=AUTOTUNE)
val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)


# 构建网络模型
def identity_block(input_tensor, kernel_size, filters, stage, block):
    filters1, filters2, filters3 = filters

    name_base = str(stage) + block + '_identity_block_'

    x = Conv2D(filters1, (1, 1), name=name_base + 'conv1')(input_tensor)
    x = BatchNormalization(name=name_base + 'bn1')(x)
    x = Activation('relu', name=name_base + 'relu1')(x)

    x = Conv2D(filters2, kernel_size, padding='same', name=name_base + 'conv2')(x)
    x = BatchNormalization(name=name_base + 'bn2')(x)
    x = Activation('relu', name=name_base + 'relu2')(x)

    x = Conv2D(filters3, (1, 1), name=name_base + 'conv3')(x)
    x = BatchNormalization(name=name_base + 'bn3')(x)

    x = layers.add([x, input_tensor], name=name_base + 'add')
    x = Activation('relu', name=name_base + 'relu4')(x)
    return x


def conv_block(input_tensor, kernel_size, filters, stage, block, strides=(2, 2)):
    filters1, filters2, filters3 = filters

    res_name_base = str(stage) + block + '_conv_block_res_'
    name_base = str(stage) + block + '_conv_block_'

    x = Conv2D(filters1, (1, 1), strides=strides, name=name_base + 'conv1')(input_tensor)
    x = BatchNormalization(name=name_base + 'bn1')(x)
    x = Activation('relu', name=name_base + 'relu1')(x)

    x = Conv2D(filters2, kernel_size, padding='same', name=name_base + 'conv2')(x)
    x = BatchNormalization(name=name_base + 'bn2')(x)
    x = Activation('relu', name=name_base + 'relu2')(x)

    x = Conv2D(filters3, (1, 1), name=name_base + 'conv3')(x)
    x = BatchNormalization(name=name_base + 'bn3')(x)

    shortcut = Conv2D(filters3, (1, 1), strides=strides, name=res_name_base + 'conv')(input_tensor)
    shortcut = BatchNormalization(name=res_name_base + 'bn')(shortcut)

    x = layers.add([x, shortcut], name=name_base + 'add')
    x = Activation('relu', name=name_base + 'relu4')(x)
    return x


def ResNet50(input_shape=[224, 224, 3], classes=1000):
    img_input = Input(shape=input_shape)
    x = ZeroPadding2D((3, 3))(img_input)

    x = Conv2D(64, (7, 7), strides=(2, 2), name='conv1')(x)
    x = BatchNormalization(name='bn_conv1')(x)
    x = Activation('relu')(x)
    x = MaxPooling2D((3, 3), strides=(2, 2))(x)

    x = conv_block(x, 3, [64, 64, 256], stage=2, block='a', strides=(1, 1))
    x = identity_block(x, 3, [64, 64, 256], stage=2, block='b')
    x = identity_block(x, 3, [64, 64, 256], stage=2, block='c')

    x = conv_block(x, 3, [128, 128, 512], stage=3, block='a')
    x = identity_block(x, 3, [128, 128, 512], stage=3, block='b')
    x = identity_block(x, 3, [128, 128, 512], stage=3, block='c')
    x = identity_block(x, 3, [128, 128, 512], stage=3, block='d')

    x = conv_block(x, 3, [256, 256, 1024], stage=4, block='a')
    x = identity_block(x, 3, [256, 256, 1024], stage=4, block='b')
    x = identity_block(x, 3, [256, 256, 1024], stage=4, block='c')
    x = identity_block(x, 3, [256, 256, 1024], stage=4, block='d')
    x = identity_block(x, 3, [256, 256, 1024], stage=4, block='e')
    x = identity_block(x, 3, [256, 256, 1024], stage=4, block='f')

    x = conv_block(x, 3, [512, 512, 2048], stage=5, block='a')
    x = identity_block(x, 3, [512, 512, 2048], stage=5, block='b')
    x = identity_block(x, 3, [512, 512, 2048], stage=5, block='c')

    x = AveragePooling2D((7, 7), name='avg_pool')(x)

    x = Flatten()(x)
    x = Dense(classes, activation='softmax', name='fc1000')(x)
    model = Model(img_input, x, name='resnet50')
    # 加载预训练模型
    model.load_weights("./pretrain_model/resnet50_weights_tf_dim_ordering_tf_kernels.h5")
    return model

model = ResNet50()
model.summary()

# 编译
opt = tf.keras.optimizers.Adam(learning_rate=1e-7)

model.compile(optimizer="adam",
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

epochs = 10  # 训练轮数
history = model.fit(
    train_ds,
    validation_data=val_ds,
    epochs=epochs
)

# 模型评估可视化
acc = history.history['accuracy']
val_acc = history.history['val_accuracy']

loss = history.history['loss']
val_loss = history.history['val_loss']

epochs_range = range(epochs)

plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)

plt.plot(epochs_range, acc, label='Training Accuracy')
plt.plot(epochs_range, val_acc, label='Validation Accuracy')
plt.legend(loc='lower right')
plt.title('Training and Validation Accuracy')

plt.subplot(1, 2, 2)
plt.plot(epochs_range, loss, label='Training Loss')
plt.plot(epochs_range, val_loss, label='Validation Loss')
plt.legend(loc='upper right')
plt.title('Training and Validation Loss')
plt.show()

# 保存模型
model.save('model/garbage_model.h5')

predict.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time : 2021/6/7 21:44
# @Author : 若谷
# @File : predict.py
# @Software: PyCharm
from tensorflow.keras.preprocessing import image
from tensorflow.keras.applications.resnet50 import preprocess_input
import numpy as np
from tensorflow.keras.models import load_model

# 加载模型
model = load_model('./model/garbage_model.h5')

# 设置类别名称列表
classes_name=['battery','cans', 'cigarette', 'plastic_bottle', 'vegetables']
china_name=["电池","易拉罐","香烟","塑料瓶","蔬菜"]

# 图片所在路径
img_path = './test/test2.jpg'
# 加载图片
img = image.load_img(img_path, target_size=(224, 224))
# 将图片转成矩阵形式
x = image.img_to_array(img)
# 扩展一维
x = np.expand_dims(x, axis=0)

x = preprocess_input(x)

# 预测
preds = model.predict(x)
max_index = np.argmax(preds, axis=-1)[0]  # 取可能性最大的那个类别

acc=preds[0][int(max_index)]*100

print("英文物体名称是：{}，中文物体名称是：{}，置信度为：{:.5f}%".format(classes_name[max_index],china_name[max_index],acc))