使用Python代码将数据集划分为训练集、验证集和测试集

此代码能同时将数据集和对应标签划分为训练集、验证集和测试集，比例为6:2:2。

首先数据集在文件夹train，标签在文件夹train_label；数据集和对应标签名字相同。

划分后的数据文件名：train1；划分后的标签文件名：train_label1

import os
import random
from shutil import copy2

# 源文件夹路径
file_path = r"C:\Users\Aisin Gioro dong po\Desktop\train"
file_path_label = r"C:\Users\Aisin Gioro dong po\Desktop\train_label"
# 新文件路径
new_file_path = r"C:\Users\Aisin Gioro dong po\Desktop\train1"
new_file_path_label = r"C:\Users\Aisin Gioro dong po\Desktop\train_label1"
# 划分数据比例6:2:2
split_rate = [0.6, 0.2, 0.2]
class_names = os.listdir(file_path)  # 数据集文件和对应标签的名字是一样的
class_names_label = os.listdir(file_path_label)
# 目标文件夹下创建文件夹
split_names = ['train', 'val', 'test']
split_names_label = ['train', 'val', 'test']
print(class_names)  # ['00000.jpg', '00001.jpg', '00002.jpg'... ]
print(class_names_label)
# 判断是否存在目标文件夹，不存在则创建---->创建train\val\test文件夹
if os.path.isdir(new_file_path):
    pass
else:
    os.makedirs(new_file_path)
if os.path.isdir(new_file_path_label):
    pass
else:
    os.makedirs(new_file_path_label)
for split_name in split_names:
    split_path = new_file_path + "/" + split_name
    split_path_label = new_file_path_label + "/" + split_name
    print(split_path)  # D:/Code/Data/GREENTdata/train, val, test
    print(split_path_label)
    if os.path.isdir(split_path):
        pass
    else:
        os.makedirs(split_path)
    if os.path.isdir(split_path_label):
        pass
    else:
        os.makedirs(split_path_label)

# 按照比例划分数据集，并进行数据图片的复制
for class_name in class_names:
    current_data_path = file_path  # D:/Code/Data/centerlinedata/tem_voc/JPEGImages/
    current_data_path_label = file_path_label
    current_all_data = os.listdir(current_data_path)
    current_all_data_label = os.listdir(current_data_path_label)
    # os.listdir() 方法用于返回指定的文件夹包含的文件或文件夹的名字的列表。这个列表以字母顺序排列。
    current_data_length = len(current_all_data)  # 文件夹下的图片个数
    current_data_length_label = len(current_all_data_label)
    current_data_index_list = list(range(current_data_length))
    current_data_index_list_label = list(range(current_data_length_label))
    # 函数list()将range()的结果直接转换为列表,是一个数字列表。
    random.shuffle(current_data_index_list)  # 将列表顺序打乱
    random.shuffle(current_data_index_list_label)

    train_path = os.path.join(new_file_path, 'train/')  # D:/Code/Data/GREENTdata/train/
    train_path_label = os.path.join(new_file_path_label, 'train/')
    val_path = os.path.join(new_file_path, 'val/')  # D:/Code/Data/GREENTdata/val/
    val_path_label = os.path.join(new_file_path_label, 'val/')
    test_path = os.path.join(new_file_path, 'test/')  # D:/Code/Data/GREENTdata/test/
    test_path_label = os.path.join(new_file_path_label, 'test/')

    train_stop_flag = current_data_length * split_rate[0]
    train_stop_flag_label = current_data_length_label * split_rate[0]
    val_stop_flag = current_data_length * (split_rate[0] + split_rate[1])
    val_stop_flag_label = current_data_length_label * (split_rate[0] + split_rate[1])

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_data_path, current_all_data[i])
    src_img_path_label = os.path.join(current_data_path_label, current_all_data_label[i])
    if current_idx < train_stop_flag:
        copy2(src_img_path, train_path)
        copy2(src_img_path_label, train_path_label)
        train_num += 1
    elif (current_idx > train_stop_flag) and (current_idx <= val_stop_flag):
        copy2(src_img_path, val_path)
        copy2(src_img_path_label, val_path_label)
        val_num += 1
    else:
        copy2(src_img_path, test_path)
        copy2(src_img_path_label, test_path_label)
        test_num += 1
    current_idx += 1
print("Done!", train_num, val_num, test_num)

参考其他博主代码改进得来：Python将数据集划分为训练集、验证集和测试集_python划分数据集验证集和测试集-CSDN博客