# split_data.py
# 功能:将两组类型数据分成训练和验证数据集
import os
from shutil import copy
import random
# 指定路径下创建文件夹
def mkfile(file):
if not os.path.exists(file):
os.makedirs(file)
# 获取data文件夹下所有文件夹名(即需要分类的类名)
file_path = r'C:\Users\1\PycharmProjects\test\data_name'
data_class = [cla for cla in os.listdir(file_path)]
# 创建 训练集train 文件夹,并由类名在其目录下创建5个子目录
mkfile('data/train')
for cla in data_class:
mkfile('data/train/' + cla)
# 创建 验证集val 文件夹,并由类名在其目录下创建子目录
mkfile('data/val')
for cla in data_class:
mkfile('data/val/' + cla)
# 划分比例,训练集 : 验证集 = 9 : 1
split_rate = 0.2
# 遍历所有类别的全部图像并按比例分成训练集和验证集
for cla in data_class:
cla_path = file_path + '/' + cla + '/' # 某一类别的子目录
images = os.listdir(cla_path) # iamges 列表存储了该目录下所有图像的名称
num = len(images)
eval_index = random.sample(images, k=int(num * split_rate)) # 从images列表中随机抽取 k 个图像名称
for index, image in enumerate(images):
# eval_index 中保存验证集val的图像名称
if image in eval_index:
image_path = cla_path + image
new_path = 'data/val/' + cla
copy(image_path, new_path) # 将选中的图像复制到新路径
# 其余的图像保存在训练集train中
else:
image_path = cla_path + image
new_path = 'data/train/' + cla
copy(image_path, new_path)
print("\r[{}] processing [{}/{}]".format(cla, index + 1, num), end="") # 当前打印行不动processing bar
print()
print("processing done!")
结果:
原本data目录下两组数据Dogs和Cats的图片被按照比例分别挪到了train和val目录下对应目录中