coco数据集划分脚本

你可以使用以下脚本来划分COCO数据集：

import os
import json
import random
from shutil import copyfile

def split_coco_dataset(coco_dir, train_percentage, val_percentage, output_dir):
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 定义划分百分比
    assert train_percentage + val_percentage < 1, "训练集和验证集的百分比之和不能大于1"
    test_percentage = 1 - train_percentage - val_percentage
    
    # 读取COCO数据集的标注文件
    annotation_path = os.path.join(coco_dir, 'annotations', 'instances_trainval.json')
    with open(annotation_path, 'r') as f:
        coco_data = json.load(f)
    
    # 获取图像列表
    image_list = coco_data['images']
    random.shuffle(image_list)
    
    # 计算划分数量
    num_images = len(image_list)
    num_train = int(num_images * train_percentage)
    num_val = int(num_images * val_percentage)
    
    # 划分图像列表
    train_images = image_list[:num_train]
    val_images = image_list[num_train:num_train+num_val]
    test_images = image_list[num_train+num_val:]
    
    # 将图像复制到相应的划分目录中
    for image in train_images:
        src_file = os.path.join(coco_dir, 'train2017', image['file_name'])
        dst_file = os.path.join(output_dir, 'train', image['file_name'])
        copyfile(src_file, dst_file)
        
    for image in val_images:
        src_file = os.path.join(coco_dir, 'train2017', image['file_name'])
        dst_file = os.path.join(output_dir, 'val', image['file_name'])
        copyfile(src_file, dst_file)
        
    for image in test_images:
        src_file = os.path.join(coco_dir, 'train2017', image['file_name'])
        dst_file = os.path.join(output_dir, 'test', image['file_name'])
        copyfile(src_file, dst_file)

你可以调用 `split_coco_dataset` 函数来划分数据集。`coco_dir` 参数是COCO数据集的根目录，`train_percentage` 和 `val_percentage` 是你想要分配给训练集和验证集的百分比，`output_dir` 是划分后的数据集存储目录。确保在运行脚本之前，已经在指定的输出目录中创建了 `train`、`val` 和 `test` 子目录。

请注意，该脚本将仅复制图像文件，而不会处理标注文件。如果你还想要处理标注文件以匹配划分后的图像，请相应地修改代码。