将本身的图片数据集分红训练集、验证集和测试集并用 Tensorflow2.0 导入

文章目录

• 项目介绍
• 划分数据集
• 一、在 dataset 目录下建立三个文件夹
• 二、将原来 dataset 文件夹中的子文件夹分别复制到这三个文件夹中
• 三、在 python 中指定全部路径
• 四、将这三个文件夹中全部子文件夹中的图片删除
• 五、将 dataset 目录下子文件中的原始图片按比例分到三个文件夹中
• 5.1 划分 96 张图片给 train 文件夹
• 5.2 划分 32 张图片给 test 文件夹
• 5.3 划分 32 张图片给 validation 文件夹
• 导入图片
• 一、构建图片生成器
• 二、将图片输入图片生成器
• 训练模型
• 一、构建模型
• 二、训练模型
• 三、测试模型

项目介绍

在下载完一个数据集后，不少朋友会为如何划分出训练集、验证集和测试集而烦恼，这篇文章将详细介绍如何划分数据集并将其用于模型训练。

首先，咱们的原始数据集以下图所示：
可见这个数据集中共有 28 类。

双击打开其中一类后如图所示：

因此一类中有 160 张图片。

咱们的目的是将这些图片按照 6：2：2 的比例分别划分到 train 文件夹、test 文件夹以及 validation 文件夹中，也就是说这三个文件夹中的图片数量分别为 96 个、32 个、32 个。

划分数据集

一、在 dataset 目录下建立三个文件夹

如图所示：

二、将原来 dataset 文件夹中的子文件夹分别复制到这三个文件夹中

如图所示：

三、在 python 中指定全部路径

import os
import shutil
PATH = './dataset/'
PATH_train = PATH + 'train/'
PATH_validation = PATH + 'validation/'
PATH_test = PATH + 'test/'

四、将这三个文件夹中全部子文件夹中的图片删除

# 删除子文件夹中全部文件
for path in [PATH_train, PATH_validation, PATH_test]:
    for dirname, _, _ in os.walk(path):
        if dirname != path:
            for each_dirname, _, filename in os.walk(dirname):
                for each_filename in filename:
                    os.remove(str(each_dirname)+'/'+str(each_filename))

五、将 dataset 目录下子文件中的原始图片按比例分到三个文件夹中

5.1 划分 96 张图片给 train 文件夹

for dirname, _, _ in os.walk(PATH):
    if dirname != PATH:
        for each_dirname, _, filename in os.walk(dirname):
            count = 0
            for each_filename in filename:
                count += 1  # 记录这是此类别下遍历的第几张图片
                if count <= 96:
                    img_path = str(each_dirname)+'/'+str(each_filename)
                    target_dir = PATH_train+img_path.split('/')[-2]  # train 目录下的目标类别子文件夹
                    target_path = target_dir+'/'+img_path.split('/')[-1]  # train 目录下的图片目标路径
                    if not os.path.exists(target_dir):
                         os.makedirs(target_dir)
                    shutil.copy(img_path, target_path)  # 复制图片

5.2 划分 32 张图片给 test 文件夹

for dirname, _, _ in os.walk(PATH):
    if dirname != PATH:
        for each_dirname, _, filename in os.walk(dirname):
            count = 0
            for each_filename in filename:
                count += 1  # 记录这是此类别下遍历的第几张图片
                if 96< count <= 128:
                    img_path = str(each_dirname)+'/'+str(each_filename)
                    target_dir = PATH_test+img_path.split('/')[-2]  # test 目录下的目标类别子文件夹
                    target_path = target_dir+'/'+img_path.split('/')[-1]  # test 目录下的图片目标路径
                    if not os.path.exists(target_dir):
                         os.makedirs(target_dir)
                    shutil.copy(img_path, target_path)  # 复制图片

5.3 划分 32 张图片给 validation 文件夹

for dirname, _, _ in os.walk(PATH):
    if dirname != PATH:
        for each_dirname, _, filename in os.walk(dirname):
            count = 0
            for each_filename in filename:
                count += 1  # 记录这是此类别下遍历的第几张图片
                if 128< count:
                    img_path = str(each_dirname)+'/'+str(each_filename)
                    target_dir = PATH_validation+img_path.split('/')[-2]  # validation 目录下的目标类别子文件夹
                    target_path = target_dir+'/'+img_path.split('/')[-1]  # validation 目录下的图片目标路径
                    if not os.path.exists(target_dir):
                         os.makedirs(target_dir)
                    shutil.copy(img_path, target_path)  # 复制图片

至此，train 文件夹、test 文件夹以及 validation 文件夹中已经包含了这些按比例划分好的图片。

导入图片

在这里咱们使用 ImageDataGenerator 类来实现图片的导入任务。

一、构建图片生成器

datagen = ImageDataGenerator(
        rescale=1./255,  # 图片归一化
        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')

二、将图片输入图片生成器

train_generator = datagen.flow_from_directory(PATH_train, target_size=(224, 224), batch_size=16, shuffle=True)
validation_generator = datagen.flow_from_directory(PATH_validation, target_size=(224, 224), batch_size=16, shuffle=True)
test_generator = datagen.flow_from_directory(PATH_test, target_size=(224, 224), batch_size=16)

Found 2688 images belonging to 28 classes.
Found 896 images belonging to 28 classes.
Found 896 images belonging to 28 classes.

训练模型

一、构建模型

这里使用 MobileNet V2 的迁移学习（只取到前 7 个 Block）来构建模型。

mobile = tf.keras.applications.MobileNetV2(include_top=False, weights='imagenet', input_shape=(224, 224, 3))
output_name = 'block_7_add'
layers = mobile.get_layer(output_name).output
simplified_mobile = tf.keras.Model(inputs=mobile.input, outputs=layers)

simplified_mobile.trainable = False

model = tf.keras.Sequential([
  simplified_mobile,
  tf.keras.layers.Dropout(0.5),
  tf.keras.layers.GlobalAveragePooling2D(),
  tf.keras.layers.Dense(28, activation='softmax')
])

二、训练模型

model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.001), loss='categorical_crossentropy', metrics=['accuracy'])

model.fit_generator(
        train_generator,
        epochs=50,
        validation_data=validation_generator)

三、测试模型

loss, accuracy = model.evaluate(test_generator, verbose=0)