图像分类模型构建：基于VGG16的迁移学习

本项目使用VGG16模型进行迁移学习，并结合图像数据增强技术，构建了一个图像分类模型，并展示了模型训练、评估和预测的过程。

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

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  # 训练轮数

2. 创建图像生成器

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

val_image_generator = ImageDataGenerator(
    rescale=1./255  # 归一化
)
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'  # 分类方式
)

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

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

3. 构建模型

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

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

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'))

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

model.summary()

4. 训练模型

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

5. 评估模型

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()

6. 进行预测

img_dir = 'D:/BS/class/archive/seg_train/seg_train'
classes = os.listdir(img_dir)

# 构建图像路径列表
img_paths = []
for class_name in classes:
    class_dir = os.path.join(img_dir, class_name)
    img_names = [img_name for img_name in os.listdir(class_dir) if img_name.endswith('.jpg')]
    class_paths = [os.path.join(class_dir, img_name) for img_name in img_names]
    img_paths.extend(class_paths)

# 随机选择一部分图像路径
selected_img_paths = random.sample(img_paths, 10)

# 加载测试集数据生成器
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'
)

# 获取测试集样本数
total_test = len(test_data_gen)
print('Total testing data batches: ', total_test)

# 进行预测并计算准确率
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))

# 随机选择一些测试样本进行展示
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()

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

7. 显示数据增强后的图片

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

# 获取训练集生成器
data_gen = image_generator.flow_from_directory(
    batch_size=1,
    directory=train_dir,
    shuffle=True,
    target_size=(IMG_HEIGHT, IMG_WIDTH),
    class_mode='categorical'
)

# 获取一个batch的数据
images, labels = next(data_gen)

# 显示图像
plt.imshow(images[0])
plt.show()

总结

本项目使用VGG16模型进行迁移学习，并结合图像数据增强技术，构建了一个图像分类模型。通过训练和评估，模型取得了不错的准确率。

注意事项

请根据您的实际情况修改代码中的路径、参数等信息。
数据增强可以根据您的需求进行调整。
为了提高模型的泛化能力，建议使用更多的数据进行训练。

未来展望

可以尝试使用其他模型进行迁移学习，例如ResNet、Inception等。
可以尝试使用更复杂的数据增强技术，例如旋转、剪切等。
可以尝试使用其他优化器和损失函数，提高模型的训练效率和准确率。