【脚本工具库】划分图像数据集（附源码）

在机器学习和深度学习任务中，数据集的划分是非常重要的一步。我们通常需要将数据集划分为训练集、验证集和测试集，以便进行模型训练和评估。手动操作不仅耗时，而且容易出错。为了解决这个问题，我们可以编写一个Python脚本，使用os、shutil和glob库来自动化这个过程。本文将详细介绍如何编写一个划分图像数据集的脚本。

准备工作

在开始之前，请确保你的系统上已经安装了Python环境。os、shutil和glob是Python标准库，自带无需安装。

脚本源码

以下是完整的Python脚本源码，该脚本可以从指定的源文件夹中按比例随机划分图像数据集，并将其复制到目标文件夹中。

import os
import shutil
import random
from glob import glob

def split_images(source_dir, train_dir, val_dir=None, test_dir=None, train_ratio=0.8, val_ratio=None, test_ratio=None):
    # 确保目标文件夹存在
    os.makedirs(train_dir, exist_ok=True)
    if val_dir:
        os.makedirs(val_dir, exist_ok=True)
    if test_dir:
        os.makedirs(test_dir, exist_ok=True)

    # 获取所有图像文件路径
    image_paths = glob(os.path.join(source_dir, '*'))
    # 打乱图像文件顺序
    random.shuffle(image_paths)

    # 计算分割点
    total_images = len(image_paths)
    train_split_index = int(total_images * train_ratio)

    if val_dir and test_dir:
        val_split_index = train_split_index + int(total_images * val_ratio)
        # 分割图像路径到训练集、验证集和测试集
        train_images = image_paths[:train_split_index]
        val_images = image_paths[train_split_index:val_split_index]
        test_images = image_paths[val_split_index:]
    elif test_dir:
        # 如果只有测试集，将剩余部分分为训练集和测试集
        train_images = image_paths[:train_split_index]
        test_images = image_paths[train_split_index:]
        val_images = []
    else:
        # 只分割为训练集和验证集
        train_images = image_paths[:train_split_index]
        val_images = image_paths[train_split_index:]
        test_images = []

    # 复制训练集图像到目标文件夹
    for image_path in train_images:
        shutil.copy(image_path, train_dir)

    # 复制验证集图像到目标文件夹
    for image_path in val_images:
        shutil.copy(image_path, val_dir)

    # 复制测试集图像到目标文件夹
    for image_path in test_images:
        shutil.copy(image_path, test_dir)

    print(f"总图像数: {total_images}")
    print(f"训练集图像数: {len(train_images)}")
    if val_dir:
        print(f"验证集图像数: {len(val_images)}")
    if test_dir:
        print(f"测试集图像数: {len(test_images)}")

# 使用示例
source_directory = r'D:\A_Data\VOCdevkit\VOC2007\JPEGImages'  # 替换为源图像文件夹路径
training_directory = r'D:\A_Data\VOCdevkit\VOC2007\JPEGImages1'  # 替换为训练集目标文件夹路径
validation_directory = r'D:\A_Data\VOCdevkit\VOC2007\JPEGImages2'  # 替换为验证集目标文件夹路径
test_directory = r'D:\A_Data\VOCdevkit\VOC2007\JPEGImages3'  # 替换为测试集目标文件夹路径

# 示例：划分为训练集和测试集，比例分别为0.8和0.2
split_images(source_directory, training_directory, test_dir=test_directory, train_ratio=0.8, test_ratio=0.2)

# 示例：划分为训练集、验证集和测试集，比例分别为0.7、0.2和0.1
# split_images(source_directory, training_directory, validation_directory, test_directory, train_ratio=0.7, val_ratio=0.2, test_ratio=0.1)

使用说明

1. 修改source_directory为源图像文件夹路径，training_directory、validation_directory和test_directory为目标文件夹路径。
2. 调整train_ratio、val_ratio和test_ratio以设置训练集、验证集和测试集的划分比例。
3. 运行脚本，程序会按设定比例随机选择图像文件，并将其复制到对应的目标文件夹中。

总结

这个脚本可以帮助你轻松地按比例随机划分图像数据集，节省了大量的时间和精力。希望这个教程对你有所帮助。如果你有任何问题或建议，欢迎在评论区留言讨论。

感谢阅读！