Dataset partitioning script for classification tasks

  This script realizes the data set division of classification tasks , With a little modification , It can also be used for target detection to divide training sets 、 Verification set 、 Test set .

import os
from shutil import copy, rmtree
import random


def mk_file(file_path: str):
    if os.path.exists(file_path):
        #  If the folder exists , First delete the original folder and all the files it contains , And re create a new folder 
        rmtree(file_path)
    os.makedirs(file_path)


def main():
    #  Make sure that random repeatable 
    random.seed(0)

    #  Centralize data 10% Divide the data into validation sets 
    split_rate = 0.1

    #  Point to the main directory where classified image data is stored 
    cwd = os.getcwd()
    data_root = os.path.join(cwd, "flower_data")
    origin_flower_path = os.path.join(data_root, "flower_photos")
    assert os.path.exists(origin_flower_path), "path '{}' does not exist.".format(origin_flower_path)

    #  Get the address information of all folders under the home directory 
    flower_class = [cla for cla in os.listdir(origin_flower_path)
                    if os.path.isdir(os.path.join(origin_flower_path, cla))]

    #  Saved training set folder 
    train_root = os.path.join(data_root, "train")
    mk_file(train_root)
    for cla in flower_class:
        #  Create a folder for each category 
        mk_file(os.path.join(train_root, cla))

    #  Create a folder to save the validation set 
    val_root = os.path.join(data_root, "val")
    mk_file(val_root)
    for cla in flower_class:
        #  Create a folder for each category 
        mk_file(os.path.join(val_root, cla))

    #  Classification according to data set category 
    for cla in flower_class:
        cla_path = os.path.join(origin_flower_path, cla)
        images = os.listdir(cla_path)
        num = len(images)

        #  Index of random sampling verification set 
        eval_index = random.sample(images, k=int(num*split_rate))
        for index, image in enumerate(images):
            if image in eval_index:
                #  Copy the files assigned to the validation set to the corresponding directory 
                image_path = os.path.join(cla_path, image)
                new_path = os.path.join(val_root, cla)
                copy(image_path, new_path)  # copy( The original path , New path )
            else:
                #  Copy the files assigned to the training set to the corresponding directory 
                image_path = os.path.join(cla_path, image)
                new_path = os.path.join(train_root, cla)
                copy(image_path, new_path)
            print("\r[{}] processing [{}/{}]".format(cla, index+1, num), end="")  # processing bar
        print()

    print("processing done!")


if __name__ == '__main__':
    main()