图像分类划分原始数据集的代码

代码：只需要提供原始数据集目录和目标数据集目录以及比例哦！！！

# 工具类
import os
import random
import shutil
from shutil import copy2
"""
数据集默认的比例是--训练集:验证集:测试集=8:1:1
"""

def data_set_split(src_data_folder, target_data_folder, train_scale=0.8, val_scale=0.1, test_scale=0.1):
    '''
    读取源数据文件夹，生成划分好的文件夹，分为trian、val、test三个文件夹进行
    :param src_data_folder: 源文件夹
    :param target_data_folder: 目标文件夹
    :param train_scale: 训练集比例
    :param val_scale: 验证集比例
    :param test_scale: 测试集比例
    :return:
    '''
    print("开始数据集划分")
    class_names = os.listdir(src_data_folder)
    # 在目标目录下创建文件夹
    split_names = ['train', 'val', 'test']
    for split_name in split_names:
        split_path = os.path.join(target_data_folder, 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))
        random.shuffle(current_data_index_list)

        train_folder = os.path.join(os.path.join(target_data_folder, 'train'), class_name)
        val_folder = os.path.join(os.path.join(target_data_folder, 'val'), class_name)
        test_folder = os.path.join(os.path.join(target_data_folder, 'test'), class_name)
        train_stop_flag = current_data_length * train_scale
        val_stop_flag = current_data_length * (train_scale + val_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)
                # print("{}复制到了{}".format(src_img_path, train_folder))
                train_num = train_num + 1
            elif (current_idx > train_stop_flag) and (current_idx <= val_stop_flag):
                copy2(src_img_path, val_folder)
                # print("{}复制到了{}".format(src_img_path, val_folder))
                val_num = val_num + 1
            else:
                copy2(src_img_path, test_folder)
                # print("{}复制到了{}".format(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 = r"F:\project\data_split\datasets"
    target_data_folder = r"F:\project\data_split\targets"
    data_set_split(src_data_folder, target_data_folder)

划分的日志信息：
E:\ProgramData\Anaconda3\python.exe F:/project/data_split/data_split.py
开始数据集划分
n02086240****
n02086240类按照0.8：0.1：0.1的比例划分完成，一共1350张图片
训练集F:\project\data_split\targets\train\n02086240：1081张
验证集F:\project\data_split\targets\val\n02086240：135张
测试集F:\project\data_split\targets\test\n02086240：134张
n02087394****
n02087394类按照0.8：0.1：0.1的比例划分完成，一共1350张图片
训练集F:\project\data_split\targets\train\n02087394：1081张
验证集F:\project\data_split\targets\val\n02087394：135张
测试集F:\project\data_split\targets\test\n02087394：134张
n02088364****
n02088364类按照0.8：0.1：0.1的比例划分完成，一共1350张图片
训练集F:\project\data_split\targets\train\n02088364：1081张
验证集F:\project\data_split\targets\val\n02088364：135张
测试集F:\project\data_split\targets\test\n02088364：134张
n02089973****
n02089973类按照0.8：0.1：0.1的比例划分完成，一共804张图片
训练集F:\project\data_split\targets\train\n02089973：644张
验证集F:\project\data_split\targets\val\n02089973：80张
测试集F:\project\data_split\targets\test\n02089973：80张
n02093754****
n02093754类按照0.8：0.1：0.1的比例划分完成，一共1350张图片
训练集F:\project\data_split\targets\train\n02093754：1081张
验证集F:\project\data_split\targets\val\n02093754：135张
测试集F:\project\data_split\targets\test\n02093754：134张
n02096294****
n02096294类按照0.8：0.1：0.1的比例划分完成，一共1350张图片
训练集F:\project\data_split\targets\train\n02096294：1081张
验证集F:\project\data_split\targets\val\n02096294：135张
测试集F:\project\data_split\targets\test\n02096294：134张
n02099601****
n02099601类按照0.8：0.1：0.1的比例划分完成，一共1350张图片
训练集F:\project\data_split\targets\train\n02099601：1081张
验证集F:\project\data_split\targets\val\n02099601：135张
测试集F:\project\data_split\targets\test\n02099601：134张
n02105641****
n02105641类按照0.8：0.1：0.1的比例划分完成，一共1350张图片
训练集F:\project\data_split\targets\train\n02105641：1081张
验证集F:\project\data_split\targets\val\n02105641：135张
测试集F:\project\data_split\targets\test\n02105641：134张
n02111889****
n02111889类按照0.8：0.1：0.1的比例划分完成，一共1350张图片
训练集F:\project\data_split\targets\train\n02111889：1081张
验证集F:\project\data_split\targets\val\n02111889：135张
测试集F:\project\data_split\targets\test\n02111889：134张
n02115641****
n02115641类按照0.8：0.1：0.1的比例划分完成，一共1350张图片
训练集F:\project\data_split\targets\train\n02115641：1081张
验证集F:\project\data_split\targets\val\n02115641：135张
测试集F:\project\data_split\targets\test\n02115641：134张