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

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

准备工作

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

脚本源码

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

import os
import shutil
import random
from glob import globdef 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)

使用说明

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