import os
import random
import math
import shutil
def data_split(old_path, new_path):
if os.path.exists(new_path) == 0:
os.makedirs(new_path)
for dirpath, dirnames, filenames in os.walk(old_path):
for filename in filenames:
if filename.endswith('.jpg'):
filepath = os.path.join(dirpath, filename)
random.shuffle(filenames)
n = len(filenames)
train_ratio = 0.8
val_ratio = 0.1
test_ratio = 0.1
train_count = int(n * train_ratio)
val_count = int(n * val_ratio)
test_count = n - train_count - val_count
train_subdir = os.path.join(new_path, 'train_set')
val_subdir = os.path.join(new_path, 'val_set')
test_subdir = os.path.join(new_path, 'test_set')
if not os.path.exists(train_subdir):
os.makedirs(train_subdir)
if not os.path.exists(val_subdir):
os.makedirs(val_subdir)
if not os.path.exists(test_subdir):
os.makedirs(test_subdir)
train_files = random.sample(filenames, train_count)
val_files = random.sample(filenames, val_count)
test_files = filenames[-test_count:]
for i, train_file in enumerate(train_files):
shutil.copy2(filepath, os.path.join(train_subdir, train_file))
for i, val_file in enumerate(val_files):
shutil.copy2(filepath, os.path.join(val_subdir, val_file))
for i, test_file in enumerate(test_files):
shutil.copy2(filepath, os.path.join(test_subdir, test_file))
if __name__ == '__main__':
data_path = '/home/uto/AIV_5_4筛/'
data_split(data_path, 'data')
随机划分数据集
最新推荐文章于 2024-02-23 16:07:50 发布