深度学习-数据集划分

根据具体需要修改

import os
import shutil
from sklearn.model_selection import train_test_split
from tqdm import tqdm

if __name__ == '__main__':
    origin_path = r'D:/Projects/UNet/Dataset'

    result_path = r'D:/Projects/UNet/Dataset/data'

    if os.path.exists(result_path):
        shutil.rmtree(result_path)
    os.mkdir(result_path)
    
    train_path = os.path.join(result_path, 'train')
    test_path = os.path.join(result_path, 'test')

    os.mkdir(train_path)
    os.mkdir(test_path)
    
    for item in ['image', 'mask']:
        os.mkdir(os.path.join(train_path, item))
        os.mkdir(os.path.join(test_path, item))

    paths = os.listdir(os.path.join(origin_path, item))
    print(len(paths))

    train, test = train_test_split(paths, test_size=0.2, random_state=42)
    # print(train, test)
    print(len(train), len(test))

    for path in tqdm(train):
        shutil.copy(os.path.join(origin_path, 'image', path), os.path.join(train_path, 'image', path))
        shutil.copy(os.path.join(origin_path, 'mask', path), os.path.join(train_path, 'mask', path))
    for path in tqdm(test):
        shutil.copy(os.path.join(origin_path, 'image', path), os.path.join(test_path, 'image', path))
        shutil.copy(os.path.join(origin_path, 'mask', path), os.path.join(test_path, 'mask', path))