按比例划分数据集train,val,test

#工具类
import os
import random
import shutil
from shutil import copy2
# shutil是python 中的高级文件操作模块，与os模块形成互补的关系，
# os主要提供了文件或文件夹的新建、删除、查看等方法，还提供了对文件以及目录的路径操作。
# shutil模块提供了移动、复制、 压缩、解压等操作，恰好与os互补，共同一起使用，基本能完成所有文件的操作
def data_set_split(src_data_folder,target_data_floder,train_scale=0.9,val_scale=0.1,test_scale=0):
    print('数据划分开始')
    class_names = os.listdir(src_data_folder)
    # 获取指定文件夹下的所有文件,用os.listdir读取到的文件路径是乱序的，调用list.sort()函数可以对文件按顺序排序。
    # 在目标目录下创建文件夹
    split_names = ['train','val','test']
    for split_name in split_names:
        split_path = os.path.join(target_data_floder,split_name)
        if os.path.isdir(split_path):
            pass
        else:
            os.mkdir(split_path)
        # 然后在split_path的目录下创建类别文件夹 
        for class_name in class_names:
            class_split_path = os.path.join(split_path,class_name)
            if os.path.isdir(class_split_path):
                pass
            else:
                os.mkdir(class_split_path)
    # 按照比例划分数据集，并进行数据图片的复制
    # 首先进行分类遍历
    for class_name in class_names:
        current_class_data_path = os.path.join(src_data_folder,class_name)
        current_all_data = os.listdir(current_class_data_path)    
        current_data_length = len(current_all_data)
        current_data_index_list = list(range(current_data_length))
        # a=list(range(7))  [0, 1, 2, 3, 4, 5, 6]
        random.shuffle(current_data_index_list)
        # random.shuffle()用于将一个列表中的元素打乱顺序，值得注意的是使用这个方法不会生成新的列表，只是将原列表的次序打乱
        # random.shuffle(a) print(a)  [3, 4, 2, 0, 5, 1, 6]

        train_folder =os.path.join(os.path.join(target_data_floder,'train'),class_name)
        val_folder =os.path.join(os.path.join(target_data_floder,'val'),class_name)
        test_folder =os.path.join(os.path.join(target_data_floder,'test'),class_name)
        train_stop_flag = current_data_length*train_scale
        val_stop_flag = current_data_length*(val_scale+train_scale)
        current_idx = 0
        train_num = 0
        val_num = 0
        test_num = 0
        for i in current_data_index_list:
            src_img_path = os.path.join(current_class_data_path,current_all_data[i])
            if current_idx <= train_stop_flag:
                copy2(src_img_path,train_folder)
                # copy2(fsrc，path)fsrc：源文件path：目标地址
                train_num = train_num+1
            elif (current_idx > train_stop_flag) and (current_idx <= val_stop_flag):
                copy2(src_img_path, val_folder)
                val_num = val_num + 1
            else:
                copy2(src_img_path, test_folder)
                test_num = test_num + 1
            current_idx = current_idx+1
        print("*********************************{}*************************************".format(class_name))
        print(
            "{}类按照{}：{}：{}的比例划分完成，一共{}张图片".format(class_name, train_scale, val_scale, test_scale, current_data_length))
        print("训练集{}：{}张".format(train_folder, train_num))
        print("验证集{}：{}张".format(val_folder, val_num))
        print("测试集{}：{}张".format(test_folder, test_num))


if __name__ == '__main__':
    src_data_folder = "/home/mini_car_class"   #原始数据路径
    target_data_folder = "/home/split_mini_car_class"  #划分后的数据路径
    data_set_split(src_data_folder, target_data_folder)