基于VGG16的图像分类模型训练与评估

本代码使用预训练的VGG16模型进行图像分类，并使用ImageDataGenerator进行数据增强和预处理，训练模型并评估其性能。

1. 导入必要的库和模块

import os
import random
import numpy as np
import pandas as pd
import tensorflow as tf
import matplotlib.pyplot as plt
from tensorflow.keras.models import Sequential, Model
from tensorflow.keras.layers import Dense, Conv2D, Flatten, Dropout, MaxPooling2D, BatchNormalization, GlobalAveragePooling2D
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.preprocessing import image
from tensorflow.keras.applications.vgg16 import VGG16
from sklearn.metrics import accuracy_score

2. 定义训练集和验证集目录、类别列表以及其他参数

train_dir = 'D:/BS/class/archive/seg_train/seg_train' # 训练集目录
val_dir = 'D:/BS/class/archive/seg_test/seg_test' #目录
classes = os.listdir(train_dir)  # 获取类别列表
batch_size = 64  # 批大小
IMG_HEIGHT = 150  # 图像高度
IMG_WIDTH = 150  # 图像宽度
epochs = 10  # 训练轮数

3. 创建图像生成器

# 创建图像生成器
train_image_generator = ImageDataGenerator(
    rescale=1./255,  # 归一化
    horizontal_flip=True  # 水平翻转
)

val_image_generator = ImageDataGenerator(
    rescale=1./255  # 归一化
)

4. 创建训练集和验证集的数据生成器

train_data_gen = train_image_generator.flow_from_directory(
    batch_size=batch_size,  # 批大小
    directory=train_dir,  # 训练集目录
    shuffle=True,  # 是否打乱数据
    target_size=(IMG_HEIGHT, IMG_WIDTH),  # 图像大小
    class_mode='categorical'  # 分类方式
)

val_data_gen = val_image_generator.flow_from_directory(
    batch_size=batch_size,  # 批大小
    directory=val_dir,  # 验证集目录
    shuffle=True,  # 是否打乱数据
    target_size=(IMG_HEIGHT, IMG_WIDTH),  # 图像大小
    class_mode='categorical'  # 分类方式
)

5. 获取训练集和验证集的样本数

total_train = len(train_data_gen)  # 训练集样本数
total_val = len(val_data_gen)  # 验证集样本数

print("总训练数据批次数：", total_train)
print("总验证数据批次数： ", total_val)

6. 加载预训练的VGG16模型作为基础模型

base_model = VGG16(weights='imagenet', include_top=False, input_shape=(IMG_HEIGHT, IMG_WIDTH, 3))

7. 冻结基础模型的参数

for layer in base_model.layers:
    layer.trainable = False

8. 构建模型

model = Sequential()
model.add(base_model)
model.add(Flatten())
model.add(Dense(1024, activation='relu'))
model.add(Dropout(0.3))
model.add(Dense(512, activation='relu'))
model.add(Dropout(0.3))
model.add(Dense(64, activation='relu'))
model.add(Dense(len(classes), activation='softmax'))

9. 编译模型

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

10. 打印模型结构

model.summary()

11. 训练模型

history = model.fit(
    train_data_gen,  # 训练集数据生成器
    steps_per_epoch=total_train,  # 每轮迭代的步数
    epochs=epochs,  # 训练轮数
    validation_data=val_data_gen,  # 验证集数据生成器
    validation_steps=total_val  # 验证集每轮迭代的步数
)

12. 绘制训练曲线

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.ylim([0, 1.0])
plt.title('Training and Validation Loss')
plt.xlabel('epoch')

plt.show()

13. 加载测试集数据生成器

test_image_generator = ImageDataGenerator(rescale=1. / 255)
test_data_gen = test_image_generator.flow_from_directory(
    batch_size=1,
    directory=val_dir,
    shuffle=False,
    target_size=(IMG_HEIGHT, IMG_WIDTH),
    class_mode='categorical'
)

14. 获取测试集样本数

total_test = len(test_data_gen)
print("Total testing data batches: ", total_test)

15. 进行预测并计算准确率

y_pred = model.predict(test_data_gen)
y_pred_class = np.argmax(y_pred, axis=1)
y_true_class = test_data_gen.classes
class_names = list(test_data_gen.class_indices.keys())
accuracy = accuracy_score(y_true_class, y_pred_class)
print('Accuracy: {:.2f}%'.format(accuracy * 100))

16. 随机选择一部分测试样本进行展示

for i in range(5):
    rand_num = random.randint(0, total_test - 1)
    img, label = test_data_gen[rand_num]
    img = img[0]
    label_name = class_names[np.argmax(label)]
    pred_name = class_names[y_pred_class[rand_num]]
    plt.title('True Label: {}, Predicted Label: {}'.format(label_name, pred_name))
    plt.imshow(img)
    plt.show()

17. 计算模型在验证集上的损失和准确率

loss, accuracy = model.evaluate(val_data_gen, steps=total_val)
print(f"Loss: {loss}")
print(f"Accuracy: {accuracy}")

代码解析

导入必要的库和模块：导入需要使用的库和模块，包括os、random、numpy、pandas、tensorflow、matplotlib.pyplot等。
定义训练集和验证集目录、类别列表以及其他参数：指定训练集和验证集的目录路径，获取类别列表，设定批大小、图像的高度和宽度、训练轮数等参数。
创建图像生成器：使用ImageDataGenerator类创建图像生成器对象，用于对图像进行预处理和数据增强操作。这里使用了归一化和水平翻转操作。
创建训练集和验证集的数据生成器：使用flow_from_directory方法创建训练集和验证集的数据生成器，通过指定目录、批大小、图像大小和分类方式等参数，将图像加载和预处理转化为可用于模型训练的数据批次。
获取训练集和验证集的样本数：通过len函数获取训练集和验证集的数据批次数。
加载预训练的VGG16模型作为基础模型：使用VGG16类加载预训练的VGG16模型作为基础模型，通过设置weights参数为'imagenet'、include_top参数为False和input_shape参数为图像的高度、宽度和通道数(3表示彩色图像)来指定模型的参数。
冻结基础模型的参数：将基础模型的所有层设置为不可训练，即冻结参数，使其不参与训练。
构建模型：使用Sequential模型构建模型，将基础模型添加到模型中，并添加全连接层、Dropout层和输出层，构建一个多层感知器模型。
编译模型：使用compile方法编译模型，指定优化器、损失函数和评估指标。
打印模型结构：使用summary方法打印模型的结构和参数数量。
训练模型：使用fit方法训练模型，传入训练集数据生成器、每轮迭代的步数、训练轮数、验证集数据生成器和验证集每轮迭代的步数等参数。
绘制训练曲线：使用matplotlib.pyplot库绘制训练过程中的准确率和损失函数变化曲线。
加载测试集数据生成器：使用ImageDataGenerator类创建测试集的数据生成器对象，用于对测试集进行预处理。
获取测试集样本数：通过len函数获取测试集的数据批次数。
进行预测并计算准确率：使用模型对测试集进行预测，通过np.argmax函数获取预测结果中概率最大的类别索引，计算准确率。
随机选择一部分测试样本进行展示：从测试集中随机选择一些样本，显示其真实标签和模型预测的标签，展示图像。
计算模型在验证集上的损失和准确率：使用evaluate方法计算模型在验证集上的损失和准确率。