将多个小型YOLO数据集合并为一个大型数据集

一、将多个小型YOLO数据集合并为一个大型数据集

import os
import shutil
import argparse

def merge_data(source_dir,target_dir,images_dir,labels_dir):
   
    images_target = os.path.join(target_dir, images_dir)
    labels_target = os.path.join(target_dir, labels_dir)
    
    # 创建目标目录
    os.makedirs(images_target, exist_ok=True)
    os.makedirs(labels_target, exist_ok=True)
    
    # 遍历源目录下所有子目录
    for subdir in os.listdir(source_dir):
        # if subdir.startswith("new_"):
        subdir_path = os.path.join(source_dir, subdir)
        print(subdir, "list dir : ",os.listdir(subdir_path))
        # 处理Images文件夹
        images_source = os.path.join(subdir_path, images_dir)
        if os.path.exists(images_source):
            print("number of images : ",len(os.listdir(images_source)))
            for filename in os.listdir(images_source):
                src_file = os.path.join(images_source, filename)
                dst_file = os.path.join(images_target, filename)
                
                # 处理文件名冲突
                counter = 1
                while os.path.exists(dst_file):
                    name, ext = os.path.splitext(filename)
                    dst_file = os.path.join(images_target, f"{name}_{counter}{ext}")
                    counter += 1
                
                shutil.copy2(src_file, dst_file)
        
        # 处理labels文件夹
        labels_source = os.path.join(subdir_path, labels_dir)
        if os.path.exists(labels_source):
            print("number of labels : ",len(os.listdir(labels_source)))
            for filename in os.listdir(labels_source):
                src_file = os.path.join(labels_source, filename)
                dst_file = os.path.join(labels_target, filename)
                
                # 处理文件名冲突
                counter = 1
                while os.path.exists(dst_file):
                    name, ext = os.path.splitext(filename)
                    dst_file = os.path.join(labels_target, f"{name}_{counter}{ext}")
                    counter += 1
                
                shutil.copy2(src_file, dst_file)
    
    print(f"数据合并完成,结果保存在: {target_dir}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='合并多个小型YOLO数据数据集为一个大型YOLO数据集,每个数据集都包含:Images、labels')
    parser.add_argument('--source_dir',required=True,help='小型YOLO数据集所在文件夹')
    parser.add_argument('--target_dir',required=True,help='合并数据集结果保存的文件夹')
    parser.add_argument('--images_dir',default='Images',help='图片保存的子文件夹名称,默认:Images')
    parser.add_argument('--labels_dir',default='labels',help='标注文件保存的子文件夹名称,默认:labels')
    args = parser.parse_args()
    merge_data(args.source_dir,args.target_dir,args.images_dir,args.labels_dir)

二、以上代码合并自己的数据集

在终端执行脚本:

python --source_dir '自己的小型数据集保存路径' --target_dir '合并的大型数据集保存路径' --images_dir '数据集的图片文件夹名称' --labels_dir '数据集的标注文件文件夹名称'

你可能感兴趣的:(YOLO,python,YOLO,python,机器学习)