图像分类 | 基于 Labelme 数据集和 VGG16 预训练模型实现迁移学习

Hi，大家好，我是源于花海。本文主要使用数据标注工具 Labelme 对自行车（bike）和摩托车（motorcycle）这两种训练样本进行标注，使用预训练模型 VGG16 作为卷积基，并在其之上添加了全连接层。基于标注样本的信息和预训练模型的特征提取能力，训练自己构建的图像分类器，从而实现迁移学习。

---

目录

一、导入必要库

二、定义目录变量

三、数据预处理--数据增强 + 标签处理

1. 定义图像数据生成器

2. 标注样本的数据说明

四、导入预训练网络--VGG16

五、模型构建

六、模型训练

七、可视化训练过程

八、模型预测

---

一、导入必要库

导入必要的库（os、json、numpy、matplotlib.pyplot 等，详见如下），设置相关配置（警告和字体），为后续的图像处理和深度学习任务做准备。

• os： 用于与操作系统进行交互，例如文件路径操作等。
• json： 处理JSON格式的数据。
• numpy： 提供对多维数组进行操作的功能。
• matplotlib.pyplot： 用于绘制图表和可视化。
• keras.preprocessing.image： 包含用于图像处理的工具，如ImageDataGenerator。
• keras.applications： 包含一些预训练的深度学习模型，这里导入VGG16。
• keras.layers、keras.models、keras.optimizers： 用于构建深度学习模型的Keras组件。
• PIL.Image： Python Imaging Library，用于图像处理。
• warnings： 用于忽略警告信息。

import os
import json
import numpy as np
import matplotlib.pyplot as plt
from keras.preprocessing.image import ImageDataGenerator, load_img, img_to_array
from keras.applications import VGG16
from keras import layers, models, optimizers
from PIL import Image
import warnings

# 忽略警告信息
warnings.filterwarnings("ignore")
# 设置中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置黑体样式
plt.rcParams['axes.unicode_minus'] = False  # 正常显示负号

二、定义目录变量

首先定义基本目录和获取当前工作目录，设置训练集和验证集的文件夹路径，通过 "os.path.join" 连接各个目录，使用 "os.path.normpath" 规范化路径格式。构建包含训练和验证图像标签的 JSON 文件的目录。

通过定义相关目录和路径来指定训练集、验证集以及图像标签的存储位置，为后续的数据加载和训练做准备。

数据集一共有近 500 张图像，使用 split.py 脚本划分数据集，规定比例为训练集 72%，验证集 18%，测试集 10% 。

# 定义基本目录和获取当前工作目录
base_dir = r'dataset'
current_dir = os.getcwd()

# 设置训练集和验证集的文件夹路径
train_dir = os.path.normpath(os.path.join(current_dir, base_dir, 'images', 'train'))
validation_dir = os.path.normpath(os.path.join(current_dir, base_dir, 'images', 'val'))

# 构建到包含训练和验证图像标签的 JSON 文件的目录
train_labels_dir = os.path.normpath(os.path.join(current_dir, base_dir, 'json'))
validation_labels_dir = os.path.normpath(os.path.join(current_dir, base_dir, 'json'))

三、数据预处理--数据增强 + 标签处理

1. 定义图像数据生成器

定义了一个实时数据增强的生成器函数，该函数通过随机应用多种变换来实现图像数据的动态增强，同时提供了对应的标签信息，大大提高了模型的泛化能力。

• 设置ImageDataGenerator：创建了一个 ImageDataGenerator 对象，用于实施数据增强。设置了多种数据增强的参数，如像素值缩放、随机旋转、水平/垂直平移、剪切、随机缩放和水平翻转。
• 读取并处理标签信息：构建一个字典 labels_dict，用于存储图像文件名和对应的标签类别。通过读取 JSON 文件中的标签信息，将图像文件名与类别建立映射。

图像种类	bike（自行车）	motorcycle（摩托车）
标签量化	0	1

• 生成器主体部分：
  • 获取图像文件列表，然后进入一个无限循环，每次循环生成一个批次的图像数据和对应的标签。
  • 从图像文件列表中随机选择一个批次的图像索引，并加载图像，进行预处理，然后将其添加到批次列表中。
  • 加载对应的标签，并将图像数据和标签作为批次的一部分返回。这是一个无限循环，可用于 Keras 模型的 fit_generator 函数进行模型训练。

# 实时数据增强--提高模型的泛化能力
def data_generator(directory, labels_directory, batch_size, target_size):
    ImageDataGenerator(
        rescale=1. / 255,  # 像素值缩放到 [0, 1] 之间
        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'  # 使用最近邻插值来填充新创建的像素
    )
    
    # 用于存储图像文件名和对应标签类别的映射
    labels_dict = {}
    
    # 读取标签信息并建立映射
    for json_file in os.listdir(labels_directory):
        if json_file.endswith('.json'):  # 确保文件是 JSON 格式的文件
            with open(os.path.join(labels_directory, json_file)) as f:  # 打开 JSON 文件
                json_data = json.load(f)  # 加载 JSON 数据

            # 处理 JSON 数据
            label = 1 if json_data["shapes"][0]["label"] == "motorcycle" else 0  # 根据标签信息确定类别（二分类--0/1）
            img_filename = os.path.basename(json_data["imagePath"])  # 获取图像文件名
            labels_dict[img_filename] = label  # 将图像文件名和对应标签类别存储在字典中

    image_list = [img for img in os.listdir(directory) if img.endswith('.png')]  # 获取图像文件列表

    while True:
        # 从图像文件列表中随机选择一个批次的图像索引
        batch_indices = np.random.choice(len(image_list), batch_size)
        batch_x = []  # 存储当前批次的图像数据
        batch_y = []  # 存储当前批次的类别数据

        for index in batch_indices:
            img_filename = image_list[index]  # 获取图像文件名
            # 加载并预处理图像
            img_path = os.path.join(directory, img_filename)  # 构建图像路径
            img = load_img(img_path, target_size=target_size)  # 加载并调整图像大小
            img_array = img_to_array(img)  # 将图像转换为 NumPy 数组
            img_array = img_array / 255.0  # 将像素值缩放到 [0, 1] 之间
            batch_x.append(img_array)  # 将图像数据添加到批次列表中
            # 加载标签
            batch_y.append(labels_dict[img_filename])  # 将对应的标签添加到批次标签列表中

        # 将批次的图像和标签转换为 NumPy 数组并返回作为生成器的一部分
        yield np.array(batch_x), np.array(batch_y)

2. 标注样本的数据说明

从指定的数据集中读取图像文件的标签信息，筛选并输出指定标签类别的图像文件名和对应的标签类别，并展示其中一部分标注信息。

• 函数参数： def get_labels(dataset_path, json_path, target_label, num_show=5):
  • dataset_path：数据集的路径，包含图像文件。
  • json_path：存储与图像文件对应标签信息的 JSON 文件的路径。
  • target_label：目标标签类别，用于筛选图像。
  • num_show：要显示的图像数量，默认为 5。
• 初始化一个列表，用于存储图像文件名和对应标签类别： image_labels = []
• 遍历数据集中的图像文件：
  • 对数据集中的每个图像文件进行遍历。
  • 构建与图像文件对应的 JSON 文件路径，读取 JSON 文件，并获取标签信息。
  • 如果标签与目标标签一致，将图像文件名和对应标签类别添加到 image_labels 列表中。
• 输出图像文件名和对应标签类别： 设置数据集和 JSON 文件路径，调用 get_labels 函数，只展示每种图像的前五张和后五张的标注信息。

def get_labels(dataset_path, json_path, target_label, num_show=5):
    # 初始化一个列表，用于存储图像文件名和对应标签类别
    image_labels = []

    # 遍历数据集中的图像文件
    for image_file in os.listdir(dataset_path):
        if image_file.endswith('.png'):
            # 构建与图像文件对应的JSON文件路径
            json_file_path = os.path.join(json_path, image_file.replace('.png', '.json'))
            # 读取JSON文件，获取标签信息
            with open(json_file_path) as f:
                json_data = json.load(f)
                label = json_data["shapes"][0]["label"]

            # 将图像文件名和对应标签类别存储在列表中
            if label == target_label:
                image_labels.append((image_file, label))

    # 输出图像文件名和对应标签类别
    print(f"{target_label}的图像文件名和对应标签类别:")
    total_images = len(image_labels)
    for i, (image_file, label) in enumerate(image_labels[:num_show]):
        print(f"文件名: {image_file}, 标签类别: {label}")
    if total_images > num_show * 2:
        print("......")
    for i, (image_file, label) in enumerate(image_labels[-num_show:]):
        print(f"文件名: {image_file}, 标签类别: {label}")

# 数据集和JSON文件路径
dataset_path = './dataset/PNGImages'
json_path = './dataset/json'
# 输出bike和输出motorcycle的图像文件名和标签类别
get_labels(dataset_path, json_path, 'bike')
get_labels(dataset_path, json_path, 'motorcycle')

四、导入预训练网络--VGG16

VGG16 卷积神经网络 —— 13 层卷积层和 5 层池化层 负责进行特征的提取，最后的 3 层全连接层 负责完成分类任务。

VGG16 的卷积核：（每层卷积的滑动步长 stride=1，padding=1）

• conv3-xxx： 卷积层均为 3×3 的卷积核，xxx表示通道数。其步长为 1，用 padding=same 填充；
• input： 输入图片大小为 224×244 的彩色图像，通道为3（RGB image），即 224×224×3；
• maxpool： 最大池化，在 VGG16 中，pooling 采用的是 2×2 的最大池化方法；
• FC-4096： 全连接层中有 4096 个节点，同样地，FC-1000 为该层全连接层有 1000 个节点；
• padding： 对矩阵在外边填充 n 圈，padding=1 即填充 1 圈，5×5 大小的矩阵，填充一圈后变成 7X7 大小；

因 VGG16 网络用于 1000 分类，而该任务只是二分类，故需修改 VGG16 网络的全连接层，同时冻结原网络的特征提取层（卷积层和池化层的权重保持不变），防止权重更新而破坏预训练权重，减少训练时间和计算资源。

使用 Keras 中的 VGG16 模型实例作为卷积基础（conv_base），并使用 ImageNet 上的预训练权重，通过将卷积层设为不可训练，可以在此基础上构建自定义的全连接网络，从而适应特定的图像分类任务，而无需重新训练 VGG16 的卷积层。

conv_base = VGG16(include_top=False,  # 不包含顶层的全连接网络
                  weights='imagenet',  # 使用 ImageNet 数据集上的预训练权重
                  input_shape=(150, 150, 3))  # 输入图像的形状为 150x150 像素，RGB 三通道
conv_base.trainable = False  # 冻结 VGG16 的卷积层，使之不被更新

conv_base.summary()  # 显示VGG16模型的摘要

五、模型构建

构建了一个顺序模型（Sequential Model），使用了预训练的 VGG16 模型作为卷积基础，添加了全连接层来适应特定的图像分类任务。

• 输出层使用 sigmoid 激活函数（适用于二分类）
• 使用 binary_crossentropy 二分类交叉熵作为损失函数。
• 使用 RMSprop 优化器，学习率为 2e-5。

model = models.Sequential()  # 创建顺序模型
model.add(conv_base)  # 将预训练的 VGG16 模型添加到顺序模型中
model.add(layers.Flatten())  # 将卷积层输出的多维数据展平成一维
model.add(layers.Dense(256, activation='relu'))  # 256个神经元的全连接层，ReLU 激活函数
model.add(layers.Dense(1, activation='sigmoid'))  # 1个神经元的输出层，sigmoid 激活函数，进行二分类
# model = VGG16()
model.summary()  # 输出模型的结构摘要

model.compile(loss='binary_crossentropy',  # 使用二分类交叉熵作为损失函数
              optimizer=optimizers.RMSprop(lr=2e-5),  # 使用 RMSprop 优化器，学习率为 2e-5
              metrics=['acc'])  # 监控模型的准确率

六、模型训练

设置数据生成器的批处理和图像大小，设定训练 25 次，观察训练过程及其训练集和验证集的准确率、损失率。

# 模型训练的参数准备
batch_size = 20  # 设置数据生成器的批处理大小
target_size = (150, 150)  # 设置将图像调整大小为 (150, 150) 的目标大小

history = model.fit(
    data_generator(train_dir, train_labels_dir, batch_size, target_size),  # 使用自定义数据生成器产生训练数据
    steps_per_epoch=len(os.listdir(train_dir)) // batch_size,  # 每个 epoch 中迭代的步数
    epochs=25,  # 训练的总 epoch 数
    validation_data=data_generator(validation_dir, validation_labels_dir, batch_size, target_size),
    # 使用自定义数据生成器产生验证数据
    validation_steps=len(os.listdir(validation_dir)) // batch_size
)

七、可视化训练过程

经过 25 轮的训练后，基于下方的 "loss/acc" 的可视化图，可以看出训练集和验证集的准确率稳定地高达 100%，训练集的损失率最低能达到 0.55%，验证集的损失率最低能达到 0.95%，可见该模型的训练效果非常好。

train_acc = history.history['acc']
train_loss = history.history['loss']
val_acc = history.history['val_acc']
val_loss = history.history['val_loss']

epoch = range(1, len(train_acc) + 1)

plt.figure(figsize=(8, 5))
plt.plot(epoch, train_acc, color='green', label='train_acc')  # 训练集准确率
plt.plot(epoch, val_acc, color='blue', label='val_acc')  # 验证集准确率
plt.plot(epoch, train_loss, color='orange', label='train_loss')  # 训练集损失率
plt.plot(epoch, val_loss, color='red', label='val_loss')  # 验证集损失率
plt.title("VGG16 Model")

plt.xlabel('Epochs', fontsize=12)
plt.ylabel('loss/acc', fontsize=12)
plt.legend(fontsize=11)
plt.ylim(0, 1)  # 设置纵坐标范围为0-1
plt.show()

八、模型预测

通过定义图像预测函数，读取测试集文件夹，以 25 个结果为一个批次，输出最终预测的结果。结果显示，测试集的 49 个图片都能被准确地预测出来。

# 定义单张图像的预测函数
def predict_image(model, image_path, target_size):
    img = Image.open(image_path)  # 打开图像文件
    img = img.resize(target_size)  # 调整图像大小
    img_array = img_to_array(img)  # 将图像转换为数组
    img_array = img_array / 255.0  # 对图像进行归一化处理
    img_array = img_array.reshape((1,) + img_array.shape)  # 将图像数组形状调整为符合模型输入要求
    # 使用模型进行图像预测
    prediction = model.predict(img_array)[0, 0]
    return prediction


# 定义显示一组图像及其结果的函数
def display_images_with_results(image_paths, predictions):
    plt.figure(figsize=(15, 10))

    for i, (image_path, prediction) in enumerate(
            zip(image_paths, predictions), 1):
        plt.subplot(5, 5, i)  # 调整子图显示
        img = Image.open(image_path)  # 打开图像文件

        # 根据预测概率确定预测类别
        predicted_class = "motorcycle" if prediction > 0.5 else "bike"

        plt.imshow(img)  # 显示图像
        plt.title(
            f"文件名: {os.path.basename(image_path)}\n预测类别: {predicted_class}")
        plt.axis('off')

    plt.tight_layout()
    plt.show()


# 设置测试集文件夹路径
test_dir = os.path.normpath(os.path.join(current_dir, base_dir, 'images', 'test'))
# 初始化空列表，用于存储每组25个结果的信息
batch_image_paths, batch_predictions, batch_actual_labels = [], [], []

# 循环遍历测试文件夹中的图像文件
for i, image_file in enumerate(os.listdir(test_dir), 1):
    if image_file.endswith('.png'):
        image_path = os.path.join(test_dir, image_file)  # 构建图像文件的完整路径
        # 使用定义的图像预测函数进行预测
        prediction = predict_image(model, image_path, target_size)
        # 将结果添加到当前批次中
        batch_image_paths.append(image_path)
        batch_predictions.append(prediction)
        # 如果达到每组25个结果，调用显示函数并清空当前批次信息
        if i % 25 == 0:
            display_images_with_results(batch_image_paths, batch_predictions)
            batch_image_paths, batch_predictions = [], []

# 如果还有剩余结果不足25个，调用显示函数
if batch_image_paths:
    display_images_with_results(batch_image_paths, batch_predictions)