实现对训练集以及测试集的9:1拆分

实现对训练集以及测试集的9:1拆分

# import splitfolders


# # train:validation:test=8:1:1
# splitfolders.ratio(input='C:/Users/徽先生/Desktop/pengyuan/pengy/', output='output', seed=1337, ratio=(0.8, 0.1, 0.1))


# 这个是将训练以及测试集划分成了9:1
# _*_ coding: utf-8 _*_

import os
import tqdm #并没有使用到，这个的作用是用到进度条

# 原始的总体打标数据
# img_dir = '/home/liudan/dataset/CarImage'
# 20211012统一标准并修改后的数据
img_dir = 图像所在位置
train_txt = 拆分之后的训练集的位置
val_txt = 拆分之后的测试集的位置

i = 0
train_lst = []
test_lst = []
img_lst = [i for i in os.listdir(img_dir) if not i.endswith('.txt')]#如果不以txt结尾，那么就不放在列表中
print(len(img_lst))
for name in os.listdir(img_dir):
    # if name.endswith('.txt'):
    #     continue
    i += 1
    if i % 10 == 0:
        test_lst.append(os.path.join(img_dir, name))
    else:
        train_lst.append(os.path.join(img_dir, name))
        # train_lst.append(os.path.join(img_dir, name))
print('val: ',len(test_lst))
print('train: ',len(train_lst))
print(i)

with open(train_txt, 'w',  encoding='utf-8') as wf:
    wf.write('\n'.join(train_lst))

with open(val_txt, 'w',  encoding='utf-8') as wf:
    wf.write('\n'.join(test_lst))