helloproject-ai/split_train_val.py

28 lines
1.0 KiB
Python

from os import makedirs, listdir
from os.path import join
from settings import datadir
from shutil import rmtree, copyfile
from random import random
valid_rate = 0.1
makedirs(join(datadir(), 'dataset'), exist_ok=True)
rmtree(join(datadir(), 'dataset', 'train'), ignore_errors=True)
rmtree(join(datadir(), 'dataset', 'val'), ignore_errors=True)
makedirs(join(datadir(), 'dataset', 'train'), exist_ok=True)
makedirs(join(datadir(), 'dataset', 'val'), exist_ok=True)
for name in listdir(join(datadir(), 'sample_set')):
print(name)
makedirs(join(datadir(), 'dataset', 'train', name))
makedirs(join(datadir(), 'dataset', 'val', name))
for file in listdir(join(datadir(), 'sample_set', name)):
if random() > valid_rate:
copyfile(src=join(datadir(), 'sample_set', name, file),
dst=join(datadir(), 'dataset', 'train', name, file))
else:
copyfile(src=join(datadir(), 'sample_set', name, file),
dst=join(datadir(), 'dataset', 'val', name, file))
# print(name, file)