Algorithm-Repo/SourceCode/MachineLearningUtils/Classification/PytorchFramework/split-data.py

# coding=utf-8
import os, random, shutil


# 将图片拆分成训练集train(0.8)和验证集val(0.2)

def moveFile(Dir, train_ratio=0.8, val_ratio=0.2):
    if not os.path.exists(os.path.join(Dir, 'train')):
        os.makedirs(os.path.join(Dir, 'train'))

    if not os.path.exists(os.path.join(Dir, 'val')):
        os.makedirs(os.path.join(Dir, 'val'))

    filenames = []
    for root, dirs, files in os.walk(Dir):
        for name in files:
            filenames.append(name)
        break

    filenum = len(filenames)

    num_train = int(filenum * train_ratio)
    sample_train = random.sample(filenames, num_train)

    for name in sample_train:
        shutil.move(os.path.join(Dir, name), os.path.join(Dir, 'train'))

    sample_val = list(set(filenames).difference(set(sample_train)))

    for name in sample_val:
        shutil.move(os.path.join(Dir, name), os.path.join(Dir, 'val'))


def remove_batch_file():
    file_path = r'D:\val'
    while(True):
        try:
            shutil.rmtree(file_path)
        except:
            if os.path.exists(file_path):
                continue
            else:
                break


if __name__ == '__main__':
    Dir = r"F:\data"
    for root, dirs, files in os.walk(Dir):
        for name in dirs:
            folder = os.path.join(root, name)
            print("正在处理:" + folder)
            moveFile(folder)
        print("处理完成")
        break

    # Dir = r"D:\codes\pytorch_heart_classification\data_j"
    # dst_path = "D:\\codes\\pytorch_heart_classification"
    #
    # count = 0
    # class_folder = os.listdir(Dir)
    # for child in class_folder:
    #     name = os.listdir(os.path.join(Dir, child))
    #     for class_name in name:
    #         for img in os.listdir(os.path.join(os.path.join(Dir, child), class_name)):
    #             count += 1
    #             print("第：{}张图片".format(count))
    #
    #         if class_name == "train":
    #             for img in os.listdir(os.path.join(os.path.join(Dir, child), class_name)):
    #                 if not os.path.exists(os.path.join(dst_path, class_name)):
    #                     os.mkdir(os.path.join(dst_path, class_name))
    #
    #                 if not os.path.exists(os.path.join(os.path.join(dst_path, class_name), child)):
    #                     os.mkdir(os.path.join(os.path.join(dst_path, class_name), child))
    #
    #                 shutil.copy(os.path.join(os.path.join(os.path.join(os.path.join(Dir, child)), class_name),img),
    #                             os.path.join(os.path.join(dst_path, class_name), child))
    #
    #         if class_name == "val":
    #             for img in os.listdir(os.path.join(os.path.join(Dir, child), class_name)):
    #                 if not os.path.exists(os.path.join(dst_path, class_name)):
    #                     os.mkdir(os.path.join(dst_path, class_name))
    #
    #                 if not os.path.exists(os.path.join(os.path.join(dst_path, class_name), child)):
    #                     os.mkdir(os.path.join(os.path.join(dst_path, class_name), child))
    #
    #                 shutil.copy(os.path.join(os.path.join(os.path.join(os.path.join(Dir, child)), class_name),img),
    #                             os.path.join(os.path.join(dst_path, class_name), child))
    #