如何在Python中调整从Roboflow下载的数据集划分比例？

核心思路 #

先将所有数据（图片+对应标注）合并为一个集合，按目标比例重新随机划分，最后将划分好的数据分别写入训练/验证/测试文件夹，按需更新标注关联信息。

具体步骤（以YOLO格式为例，其他格式可类比） #

1. 读取并合并原数据集文件 #

假设下载后的数据集结构为：

dataset/
├── train/
│   ├── images/
│   └── labels/
├── valid/
│   ├── images/
│   └── labels/
└── test/
    ├── images/
    └── labels/

执行代码合并所有图片路径：

import os
import shutil
import random

# 定义路径
original_dataset_path = "path/to/your/original/dataset"
new_dataset_path = "path/to/your/new/dataset"

# 创建新数据集文件夹结构
for split in ["train", "valid", "test"]:
    os.makedirs(os.path.join(new_dataset_path, split, "images"), exist_ok=True)
    os.makedirs(os.path.join(new_dataset_path, split, "labels"), exist_ok=True)

# 收集所有有效图片路径
all_image_paths = []
for split in ["train", "valid", "test"]:
    img_dir = os.path.join(original_dataset_path, split, "images")
    for filename in os.listdir(img_dir):
        if filename.lower().endswith((".jpg", ".jpeg", ".png")):
            all_image_paths.append(os.path.join(img_dir, filename))

# 随机打乱数据，避免类别分布不均
random.shuffle(all_image_paths)

2. 按目标比例划分数据集 #

比如设置训练集70%、验证集20%、测试集10%：

total_count = len(all_image_paths)
train_count = int(total_count * 0.7)
valid_count = int(total_count * 0.2)
test_count = total_count - train_count - valid_count

# 切割数据集
train_images = all_image_paths[:train_count]
valid_images = all_image_paths[train_count:train_count+valid_count]
test_images = all_image_paths[train_count+valid_count:]

3. 复制文件到对应子集文件夹 #

同步复制对应的标签文件（YOLO格式下标签与图片同名，后缀为.txt）：

def copy_data_to_split(image_list, split_name):
    for img_path in image_list:
        img_filename = os.path.basename(img_path)
        # 拼接对应标签文件路径
        label_filename = os.path.splitext(img_filename)[0] + ".txt"
        label_path = os.path.join(os.path.dirname(os.path.dirname(img_path)), "labels", label_filename)
        
        # 复制图片
        shutil.copy(img_path, os.path.join(new_dataset_path, split_name, "images", img_filename))
        # 复制标签（若存在）
        if os.path.exists(label_path):
            shutil.copy(label_path, os.path.join(new_dataset_path, split_name, "labels", label_filename))

# 执行复制操作
copy_data_to_split(train_images, "train")
copy_data_to_split(valid_images, "valid")
copy_data_to_split(test_images, "test")

额外注意事项 #

• 若使用VOC/COCO格式，需同步修改XML/JSON标注文件中的图片路径，或重新生成子集索引文件（如VOC的train.txt、val.txt）。
• 划分完成后，建议统计各子集的类别分布，确保各子集的类别比例与整体一致，避免模型训练偏差。
• 若数据集包含特殊格式的标注文件，需针对性调整标签文件的读取和复制逻辑。

内容的提问来源于stack exchange，提问作者抹茶抹茶