Python 实现数据集自动划分（训练集 / 验证集 / 测试集）

代码解析

1. 导入依赖库

首先导入必要的 Python 库，os 用于文件路径操作，random 用于随机打乱文件列表，shutil 用于文件复制：

import os
import random
import shutil

2. 保证结果可复现

设置随机种子，确保每次运行代码的划分结果一致（便于实验对比）：

random.seed(42)  # 固定种子，结果可复现

3. 目录设置与创建

定义数据集根目录及图片、标签存放路径，并自动创建划分后的子目录（train/val/test）：

# 数据集根目录
data_dir = './datasets'
images_dir = os.path.join(data_dir, 'images')  # 原始图片目录
labels_dir = os.path.join(data_dir, 'labels')  # 原始标签目录

# 创建划分后的子目录（若已存在则不报错）
os.makedirs(os.path.join(images_dir, 'train'), exist_ok=True)
os.makedirs(os.path.join(images_dir, 'val'), exist_ok=True)
os.makedirs(os.path.join(images_dir, 'test'), exist_ok=True)
os.makedirs(os.path.join(labels_dir, 'train'), exist_ok=True)
os.makedirs(os.path.join(labels_dir, 'val'), exist_ok=True)
os.makedirs(os.path.join(labels_dir, 'test'), exist_ok=True)

4. 读取与划分文件

获取所有图片文件，按比例划分并随机打乱：

# 获取所有.jpg 格式的图片文件（可根据实际格式修改）
image_files = [f for f in os.listdir(images_dir) if f.endswith('.jpg')]

# 定义划分比例（可根据需求调整）
train_ratio = 0.8  # 训练集占比
val_ratio = 0.1    # 验证集占比
test_ratio = 0.1   # 测试集占比

# 计算各集合的文件数量
total_images = len(image_files)
train_index = int(total_images * train_ratio)      # 训练集结束索引
val_index = int(total_images * (train_ratio + val_ratio))  # 验证集结束索引

# 随机打乱文件列表（保证划分随机性）
random.shuffle(image_files)

# 划分数据集
train_images = image_files[:train_index]           # 训练集
val_images = image_files[train_index:val_index]    # 验证集
test_images = image_files[val_index:]              # 测试集

5. 复制文件到对应目录

将图片和对应的标签文件（假设标签与图片同名，后缀为.txt）复制到划分后的目录：

# 复制训练集文件
for img_file in train_images:
    label_file = img_file.replace('.jpg', '.txt')  # 标签文件名（与图片对应）
    shutil.copy(os.path.join(images_dir, img_file), os.path.join(images_dir, 'train'))
    shutil.copy(os.path.join(labels_dir, label_file), os.path.join(labels_dir, 'train'))

# 复制验证集文件（逻辑同上）
for img_file in val_images:
    label_file = img_file.replace('.jpg', '.txt')
    shutil.copy(os.path.join(images_dir, img_file), os.path.join(images_dir, 'val'))
    shutil.copy(os.path.join(labels_dir, label_file), os.path.join(labels_dir, 'val'))

# 复制测试集文件（逻辑同上）
for img_file in test_images:
    label_file = img_file.replace('.jpg', '.txt')
    shutil.copy(os.path.join(images_dir, img_file), os.path.join(images_dir, 'test'))
    shutil.copy(os.path.join(labels_dir, label_file), os.path.join(labels_dir, 'test'))

6. 输出划分结果

打印各集合的文件数量，确认划分完成：

print(f"数据集划分完成！训练集：{len(train_images)}，验证集：{len(val_images)}，测试集：{len(test_images)}")

使用说明

1. 格式适配：若图片格式为.png 等，需修改 endswith('.jpg') 为对应格式；标签格式不同时同理。
2. 比例调整：修改 train_ratio、val_ratio、test_ratio 可自定义划分比例。

目录结构要求：原始数据集需按如下结构存放（可修改代码中 data_dir 路径适配你的数据）：

datasets/
├─ images/  # 存放所有图片（.jpg 格式）
└─ labels/  # 存放所有标签（.txt 格式，与图片同名）