yolo数据集格式按照每一个类别的比例划分数据集

写在前面： 写脚本不易，写博客不易，请多点赞关注，谢谢。10多年来，我一直免费给大家毫无保留的分享技术等，不但从来没被打赏过，而且在分享有些模型转化处理的高级脚本中，有些同胞由于自身的问题，没有转化成功就出言伤人，各位同胞们，你们要知道，随着时间推移，各个第三方库的版本一致在更新，比如onnx，torch等，这就导致这些库的接口有所变化，写法也有所变化，包括你自己的环境等等各种原因导致你没能成功，而我是无私免费分享者，却遭到不礼貌的言语伤害，我伤心，但这并没有影响我的分享热情，我会继续分享给大家。
背景： 介于在用yolov系列模型训练模型时，数据集划分对训练结果的影响较大，所以写了一个数据集重新划分的脚本，支持两类检测的数据，标签分别为0和1，针对数据集中的空标签都有所考虑。
脚本功能： 按照不同标签，按比例进行划分，比如按照1：8的比例你设置训练集占80%，即train_ratio设置为0.8，那么你标签为0的数据数量的80%将划分为训练集，标签为1的80%将划分为训练集，针对空标签，你也可以选择不划分，或者将空标签仅放在训练集里面，或者正常按照训练集和验证集来划分，若一个标签中有多个不同类别的标签，也会单独把这一类数据拿出来，按照80%比例进行划分。

1、待划分数据集格式如下：

主目录yolodatas名称可以不一样，但是其下面的子目录文件夹的名字必须保持一致，若yolodatas名称不一样则需要在脚本当中修改，详细见3.
在这里插入图片描述

2、划分完成之后数据集格式如下：

在这里插入图片描述

3、划分脚本参数介绍：

–base_train_images_path       待划分数据的训练集图像路径
–base_val_images_path        待划分数据的验证集图像路径
–save_root_path        新的数据集保存路径，主目录名
–train_ratio        训练集所占的比例
–emptys        空的标签正常分训练集和验证集，否则不用空标签
–emptys_trian_val        空标签放训练集和验证集，否组只放训练集，emptys为True时有效，emptys为False时此参数无效。
说明： 如果原始待划分的数据集格式和1中所示一致，则不需要该脚本中的路径base_train_images_path和base_val_images_path，如果不一致，则需要修改。其他参数参考脚本参数修改即可。

4、转化脚本如下，可拷贝使用：

import os
import argparse
import shutil
import random
def makdirs(opt):
    images_train = os.path.join(opt.save_root_path, 'images/train')
    images_val = os.path.join(opt.save_root_path ,'images/val')

    label_train = os.path.join(opt.save_root_path, 'labels/train')
    label_val = os.path.join(opt.save_root_path,'labels/val')
    os.makedirs(images_train,exist_ok=True)
    os.makedirs(images_val, exist_ok=True)
    os.makedirs(label_train, exist_ok=True)
    os.makedirs(label_val, exist_ok=True)

    return images_train,images_val,label_train,label_val
def readdir(opt):
    #所有图像路径列表
    images_path_list = [os.path.join(opt.base_train_images_path, train_label_name) for train_label_name in
                               os.listdir(opt.base_train_images_path)] \
                              + [os.path.join(opt.base_val_images_path, base_label_val_patch) for base_label_val_patch in
                                 os.listdir(opt.base_val_images_path)]
    #打乱列表
    random.shuffle(images_path_list)
    # 所有标签路径列表
    txt_path_list = [one_images_path_list.replace('images','labels')[:-3]+'txt' for one_images_path_list in images_path_list]

    zero_class_images = []
    zero_class_labels = []


    one_class_images = []
    one_class_labels = []

    one_two_class_images = []
    one_two_class_labels = []

    empty_class_images = []
    empty_class_labels = []


    #所有数据集数量
    file_number = len(images_path_list)

    for idx in range(file_number):
        with open(txt_path_list[idx],'r',encoding='utf-8') as fo:
            lines = fo.readlines()

        if len(lines)>1:
            #标签文件中有多个目标时，将多个目标标签存入列表并统计不同类型标签数量
            class_index_list = []
            #每一行去除\n
            relines = [lls.strip() for lls in lines]

            #遍历每一行，并统计每个标签的数量
            for line in relines:
                #每一行第一个标签名放入列表
                # print(line,line[0])
                class_index_list.append(int(line[0]))
            # print(class_index_list)

            # print("class_index_list:", class_index_list)

            #标签类
            class_label_list = list(set(class_index_list))
            # print("class_label_list:",class_label_list)

            if len(class_label_list)==1:
                if class_label_list[0]==0:
                    zero_class_images.append(images_path_list[idx])
                    zero_class_labels.append(txt_path_list[idx])
                elif class_label_list[0]==1:
                    one_class_images.append(images_path_list[idx])
                    one_class_labels.append(txt_path_list[idx])
                else:
                    print("注意还有其他标签，请修改此处代码，确保所有标签都添加都了标签列表里面！！！！")

            else:
                print("***********************************")
                one_two_class_images.append(images_path_list[idx])
                one_two_class_labels.append(txt_path_list[idx])


        elif len(lines)==1:
            #签文件中有一个目标时
            relines = [lls.strip() for lls in lines]
            class_ids = int(relines[0][0])
            if class_ids == 0:
                zero_class_images.append(images_path_list[idx])
                zero_class_labels.append(txt_path_list[idx])
            elif class_ids == 1:
                one_class_images.append(images_path_list[idx])
                one_class_labels.append(txt_path_list[idx])
            else:
                print("注意还有其他标签，请修改此处代码，确保所有标签都添加都了标签列表里面！！！！")
        elif len(lines)==0:
            #标签为空文件时
            empty_class_images.append(images_path_list[idx])
            empty_class_labels.append(txt_path_list[idx])
        else:
            print("error")
    return zero_class_images,zero_class_labels,one_class_images,one_class_labels,one_two_class_images,one_two_class_labels,empty_class_images,empty_class_labels

def copytofile(opt):
    images_train, images_val, label_train, label_val= makdirs(opt)
    zero_class_images, zero_class_labels, one_class_images, one_class_labels, one_two_class_images, one_two_class_labels, empty_class_images, empty_class_labels =readdir(opt)
    if len(zero_class_images)!=len(zero_class_labels) or len(one_class_images)!=len(one_class_labels):
        print("数据集列表不匹配1！！！！")
        return 0
    if len(one_two_class_images)!=len(one_two_class_labels) or len(empty_class_images)!=len(empty_class_labels):
        print("数据集列表不匹配2！！！！")
        return 0
    number_train_zero = int(opt.train_ratio*len(zero_class_images))
    number_train_one = int(opt.train_ratio*len(one_class_images))
    number_train_one_two = int(opt.train_ratio*len(one_two_class_images))
    number_train_empty = int(opt.train_ratio*len(empty_class_images))
    len_zero,len_one,len_one_two,len_empty = len(zero_class_images),len(one_class_images),len(one_two_class_images),len(empty_class_images)
    for z in range(len_zero):
        #标签为0的数据
        if z<number_train_zero:
            shutil.copy2(zero_class_images[z], os.path.join(images_train,zero_class_images[z].split('\\')[-1]))
            shutil.copy2(zero_class_labels[z], os.path.join(label_train, zero_class_labels[z].split('\\')[-1]))
        else:
            shutil.copy2(zero_class_images[z], os.path.join(images_val, zero_class_images[z].split('\\')[-1]))
            shutil.copy2(zero_class_labels[z], os.path.join(label_val, zero_class_labels[z].split('\\')[-1]))

    for o in range(len_one):
        #标签为1的数据
        print(one_class_images[o],os.path.join(images_train,one_class_images[o].split('\\')[-1]))
        print(one_class_labels[o],os.path.join(label_train, one_class_labels[o].split('\\')[-1]))
        if o<number_train_one:
            shutil.copy2(one_class_images[o], os.path.join(images_train,one_class_images[o].split('\\')[-1]))
            shutil.copy2(one_class_labels[o], os.path.join(label_train, one_class_labels[o].split('\\')[-1]))
        else:
            shutil.copy2(one_class_images[o], os.path.join(images_val, one_class_images[o].split('\\')[-1]))
            shutil.copy2(one_class_labels[o], os.path.join(label_val, one_class_labels[o].split('\\')[-1]))


    for ot in range(len_one_two):
        #标签为0和1的数据
        if ot<number_train_one_two:
            shutil.copy2(one_two_class_images[ot], os.path.join(images_train,one_two_class_images[ot].split('\\')[-1]))
            shutil.copy2(one_two_class_labels[ot], os.path.join(label_train, one_two_class_labels[ot].split('\\')[-1]))
        else:
            shutil.copy2(one_two_class_images[ot], os.path.join(images_val, one_two_class_images[ot].split('\\')[-1]))
            shutil.copy2(one_two_class_labels[ot], os.path.join(label_val, one_two_class_labels[ot].split('\\')[-1]))

    if opt.emptys:
        if opt.emptys_trian_val:
            for en in range(len_empty):
                #标签为空的数据
                if en<number_train_empty:
                    shutil.copy2(empty_class_images[en], os.path.join(images_train,empty_class_images[en].split('\\')[-1]))
                    shutil.copy2(empty_class_labels[en], os.path.join(label_train, empty_class_labels[en].split('\\')[-1]))
                else:
                    shutil.copy2(empty_class_images[en], os.path.join(images_val, empty_class_images[en].split('\\')[-1]))
                    shutil.copy2(empty_class_labels[en], os.path.join(label_val, empty_class_labels[en].split('\\')[-1]))
        else:
            for en in range(len_empty):
                #标签为空的数据
                shutil.copy2(empty_class_images[en], os.path.join(images_train,empty_class_images[en].split('\\')[-1]))
                shutil.copy2(empty_class_labels[en], os.path.join(label_train, empty_class_labels[en].split('\\')[-1]))


    print('#------------------------------------------------------总数据集数量为-----------------------------------------------------------------------#')
    print("划分前总数据集：",len_zero+len_one+len_one_two+len_empty)
    print('#-----------------------------------------------------所有数据集排布情况---------------------------------------------------------------------#')
    print("zero:%s\none:%s\nzero_two:%s\nempty:%s" % (len_zero,len_one,len_one_two,len_empty))
    print('#---------------------------------------------------数据集划分后训练集分布--------------------------------------------------------------------#')
    print("zero:%s\none:%s\nzero_two:%s\nempty:%s\n训练集总数：%s"%(number_train_zero,number_train_one,number_train_one_two,number_train_empty,number_train_zero+number_train_one+number_train_one_two+number_train_empty))
    print('#---------------------------------------------------数据集划分后验证集分布--------------------------------------------------------------------#')
    print("zero:%s\none:%s\nzero_two:%s\nempty:%s\n验证集总数：%s" % (len_zero-number_train_zero, len_one-number_train_one, len_one_two-number_train_one_two, len_empty-number_train_empty,len_zero+len_one+len_one_two+len_empty-(number_train_zero+number_train_one+number_train_one_two+number_train_empty)))
    print('#---------------------------------------------------数据集划分后数据集总数--------------------------------------------------------------------#')
    print("划分后总数据集：",number_train_zero+number_train_one+number_train_one_two+number_train_empty+len_zero+len_one+len_one_two+len_empty-(number_train_zero+number_train_one+number_train_one_two+number_train_empty))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--base_train_images_path', type=str, default='yolodatas/images/train',help='原始图像的训练集路径')
    parser.add_argument('--base_val_images_path', type=str, default='yolodatas/images/val', help='原始图像的验证集路径')
    parser.add_argument('--save_root_path', nargs='+', type=str, default='./mynewyolodata', help='新的数据集保存路径')
    parser.add_argument('--train_ratio', type=float, default=0.8, help='训练集所占的比例')
    parser.add_argument('--emptys', type=bool, default=True, help='空的标签正常分训练集和验证集，否则不用空标签')
    parser.add_argument('--emptys_trian_val', type=bool, default=True, help='空标签放训练集和验证集，否组只放训练集')
    opt = parser.parse_args()

    copytofile(opt)