图像分类模型数据集划分教程：如何划分训练集和验证集

图像分类模型数据集划分教程：如何划分训练集和验证集

关于作者

作者：小白熊

作者简介：精通python、matlab、c#语言，擅长机器学习，深度学习，机器视觉，目标检测，图像分类，姿态识别，语义分割，路径规划，智能优化算法，数据分析，各类创新融合等等。

联系邮箱：[email protected]

科研辅导、知识付费答疑、个性化定制以及其他合作需求请联系作者~

1 前言

  在处理图像分类任务时，数据集的划分是一个重要步骤。为了将数据集有效地分为训练集和验证集，我们可以编写一个Python脚本来自动化这一过程。该脚本通过遍历原始数据集中的每个类别，并根据指定的比例随机划分图像到训练集和验证集，从而实现了数据集的自动化划分。同时通过显示进度条和处理完成信息，提供了良好的用户交互体验。

2 代码解析

2.1 导入必要的库

import os  
from shutil import copy, rmtree  
import random  
import warnings

• os：用于处理文件和目录路径。
• shutil.copy：用于复制文件。
• shutil.rmtree：用于递归地删除目录及其内容。
• random：用于生成随机数，这里主要用于随机划分数据集。
• warnings：用于忽略警告信息，以避免脚本运行时出现不必要的警告。

2.2 创建文件夹函数

def mk_file(file_path: str):  
    """  
    创建文件夹，如果文件夹已存在则先删除。  
      
    :param file_path: 要创建的文件夹路径  
    """  
    if os.path.exists(file_path):  
        rmtree(file_path)  
    os.makedirs(file_path)

• mk_file函数用于确保目标文件夹路径存在且为空。如果文件夹已存在，它会先被删除，然后重新创建。

2.3 主函数

def main():
    # 保证随机可复现
    random.seed(0)

    # 将数据集中10%的数据划分到验证集中
    split_rate = 0.1

    data_root = ".\data"  # 保存路径
    origin_path = '.\原数据集'  # 原始路径
    
    assert os.path.exists(origin_path), "path '{}' does not exist.".format(origin_path)

    flower_class = [cla for cla in os.listdir(origin_path)
                    if os.path.isdir(os.path.join(origin_path, cla))]

    # 建立保存训练集的文件夹
    train_root = os.path.join(data_root, "train")
    mk_file(train_root)
    for cla in flower_class:
        # 建立每个类别对应的文件夹
        mk_file(os.path.join(train_root, cla))

    # 建立保存验证集的文件夹
    val_root = os.path.join(data_root, "val")
    mk_file(val_root)
    for cla in flower_class:
        # 建立每个类别对应的文件夹
        mk_file(os.path.join(val_root, cla))

    for cla in flower_class:
        cla_path = os.path.join(origin_path, cla)
        images = os.listdir(cla_path)
        num = len(images)
        # 随机采样验证集的索引
        eval_index = random.sample(images, k=int(num * split_rate))
        for index, image in enumerate(images):
            if image in eval_index:
                # 将分配至验证集中的文件复制到相应目录
                image_path = os.path.join(cla_path, image)
                new_path = os.path.join(val_root, cla)
                copy(image_path, new_path)
            else:
                # 将分配至训练集中的文件复制到相应目录
                image_path = os.path.join(cla_path, image)
                new_path = os.path.join(train_root, cla)
                copy(image_path, new_path)
            print("\r[{}] processing [{}/{}]".format(cla, index + 1, num), end="")  # processing bar
        print()

    print("processing done!")

2.3.1 数据集划分

• 设置随机种子以确保每次运行脚本时，图像分配到训练集和验证集的方式相同。
• 定义数据集划分比例。
• 指定数据集保存路径和原始数据集路径。
• 检查原始数据集路径是否存在，如果不存在则抛出异常。
  
  

2.3.2 创建文件夹结构

• 遍历原始数据集路径下的所有文件夹，获取类别名称。
• 创建训练集和验证集的根文件夹以及每个类别对应的子文件夹。
  
  

2.3.3 遍历每个类别，划分数据集

• 遍历每个类别，获取该类别的图像列表和图像数量。
• 使用random.sample函数随机选择验证集图像的索引。
• 遍历所有图像，根据索引判断图像应该复制到训练集还是验证集文件夹。
• 使用shutil.copy函数复制图像。
• 显示进度条，以便用户了解当前处理进度。

3 完整代码

import os
from shutil import copy, rmtree
import random
import warnings

warnings.filterwarnings("ignore")


def mk_file(file_path: str):
    if os.path.exists(file_path):
        rmtree(file_path)
    os.makedirs(file_path)


def main():
    # 保证随机可复现
    random.seed(0)

    # 将数据集中10%的数据划分到验证集中
    split_rate = 0.1

    data_root = ".\data"  # 保存路径
    origin_path = '.\原数据集'  # 原始路径
    
    assert os.path.exists(origin_path), "path '{}' does not exist.".format(origin_path)

    flower_class = [cla for cla in os.listdir(origin_path)
                    if os.path.isdir(os.path.join(origin_path, cla))]

    # 建立保存训练集的文件夹
    train_root = os.path.join(data_root, "train")
    mk_file(train_root)
    for cla in flower_class:
        # 建立每个类别对应的文件夹
        mk_file(os.path.join(train_root, cla))

    # 建立保存验证集的文件夹
    val_root = os.path.join(data_root, "val")
    mk_file(val_root)
    for cla in flower_class:
        # 建立每个类别对应的文件夹
        mk_file(os.path.join(val_root, cla))

    for cla in flower_class:
        cla_path = os.path.join(origin_path, cla)
        images = os.listdir(cla_path)
        num = len(images)
        # 随机采样验证集的索引
        eval_index = random.sample(images, k=int(num * split_rate))
        for index, image in enumerate(images):
            if image in eval_index:
                # 将分配至验证集中的文件复制到相应目录
                image_path = os.path.join(cla_path, image)
                new_path = os.path.join(val_root, cla)
                copy(image_path, new_path)
            else:
                # 将分配至训练集中的文件复制到相应目录
                image_path = os.path.join(cla_path, image)
                new_path = os.path.join(train_root, cla)
                copy(image_path, new_path)
            print("\r[{}] processing [{}/{}]".format(cla, index + 1, num), end="")  # processing bar
        print()

    print("processing done!")


if __name__ == '__main__':
    main()

4 文件夹结构

原始数据集文件夹结构应如下：

.\原数据集\  
    类别1\  
        图像1.jpg  
        图像2.jpg  
        ...  
    类别2\  
        图像1.jpg  
        图像2.jpg  
        ...  
    ...

脚本运行后，将生成以下文件夹结构用于存储训练集和验证集：

.\data\  
    train\  
        类别1\  
            图像1.jpg  
            ...  
        类别2\  
            图像1.jpg  
            ...  
        ...  
    val\  
        类别1\  
            图像1.jpg  
            ...  
        类别2\  
            图像1.jpg  
            ...  
        ...

5 使用说明

1. 准备数据集：确保你的原始数据集按照上述文件夹结构进行组织，一个类别一个文件夹，文件夹名称就是类别名称。
2. 检查结果：脚本运行后，检查.\data\train\和.\data\val\文件夹，确保图像已正确划分到训练集和验证集中。

  通过这个脚本，你可以轻松地将图像数据集划分为训练集和验证集，为后续的深度学习模型训练做好准备。