Python实现数据集自动划分的示例代码(训练集/验证集/测试集)

import os
import random
import shutil
# 分割训练集和验证集
# 设置随机种子以保证结果可复现
random.seed(42)

# 数据集根目录
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(labels_dir, 'train'), exist_ok=True)
os.makedirs(os.path.join(labels_dir, 'val'), exist_ok=True)
os.makedirs(os.path.join(images_dir, 'test'), exist_ok=True)
os.makedirs(os.path.join(labels_dir, 'test'), exist_ok=True)

# 获取所有的图像文件名
image_files = [f for f in os.listdir(os.path.join(data_dir, 'images')) 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:]


for img_file in train_images:
    label_file = img_file.replace('.jpg', '.txt')
    shutil.copy(os.path.join(data_dir, 'images', img_file), os.path.join(images_dir, 'train'))
    shutil.copy(os.path.join(data_dir, 'labels', 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(data_dir, 'images', img_file), os.path.join(images_dir, 'val'))
    shutil.copy(os.path.join(data_dir, 'labels', 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(data_dir, 'images', img_file), os.path.join(images_dir, 'test'))
    shutil.copy(os.path.join(data_dir, 'labels', label_file), os.path.join(labels_dir, 'test'))

print(f"数据集划分完成！训练集: {len(train_images)}，验证集: {len(val_images)}，测试集: {len(test_images)}")
import sys
sys.exit(0)

import os
import random
import shutil

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

# 数据集根目录
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)

# 获取所有.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:]          # 测试集

# 复制训练集文件
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'))

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

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