import os
import random
from shutil import copy2
# 这个函数用于分割数据集
def data_set_split(src_data_folder, target_data_folder, slice_data=[0.4, 0.3, 0.3]):
'''
读取源数据文件夹,生成划分好的文件夹,分为train、val、test三个文件夹
:param src_data_folder: 源数据文件夹的路径
:param target_data_folder: 目标数据文件夹的路径
:param slice_data: 数据划分比例 [训练集比例,验证集比例,测试集比例]
:return: None
'''
print("开始数据集划分")
# 获取所有类别文件夹的名称
class_names = os.listdir(src_data_folder)
# 创建train、val、test数据集文件夹
split_names = ['train', 'val', 'test']
for split_name in split_names:
split_path = os.path.join(target_data_folder, split_name)
# 如果数据集文件夹不存在,则创建
if not os.path.isdir(split_path):
os.mkdir(split_path)
# 在每个数据集文件夹下创建类别文件夹
for class_name in class_names:
class_split_path = os.path.join(split_path, class_name)
if not os.path.isdir(class_split_path):
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(target_data_folder, 'train', class_name)
val_folder = os.path.join(target_data_folder, 'val', class_name)
test_folder = os.path.join(target_data_folder, 'test', class_name)
# 计算划分的界限索引
train_stop_flag = int(current_data_length * slice_data[0])
val_stop_flag = int(current_data_length * (slice_data[0] + slice_data[1]))
# 按照比例复制文件到对应的训练、验证、测试集文件夹
train_num, val_num, test_num = 0, 0, 0
for i, idx in enumerate(current_data_index_list):
src_img_path = os.path.join(current_class_data_path, current_all_data[idx])
# 根据索引位置判断是训练、验证还是测试,然后复制到相应文件夹
if i < train_stop_flag:
copy2(src_img_path, train_folder)
train_num += 1
elif i < val_stop_flag:
copy2(src_img_path, val_folder)
val_num += 1
else:
copy2(src_img_path, test_folder)
test_num += 1
# 输出当前类别数据划分的情况
print(f"*********************************{class_name}*************************************")
print("{}类按照{}:{}:{}的比例划分完成,一共{}张图片".format(class_name, *slice_data, 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:\lajishujuji\xinxinxinxinxin\LaJiFenLei"
# 目标数据文件夹的路径,即划分后的数据集将存放的地方
target_data_folder = r"F:\lajishujuji\LaJiFenLai"
# 调用函数进行数据划分,这里划分比例为60%训练集,20%验证集,20%测试集
data_set_split(src_data_folder, target_data_folder, slice_data=[0.6, 0.2, 0.2])
YOLO分类数据集划分
于 2024-03-22 16:09:38 首次发布
文章讲述了如何使用Python脚本对数据集进行划分,将数据分为训练集、验证集和测试集,适用于数据预处理过程。
摘要由CSDN通过智能技术生成