[deep learning] data segmentation

Before training the model , To split the data set , Divide into training sets according to a certain proportion 、 Test set 、 Verification set , In some cases, test sets can also be ignored , The author has divided the test set here （ Training set ： Test set ： Verification set =8：1：1）, Nonsense , Go straight to the code

# h_split.py

import os
from shutil import copy
import random


def mkfile(file):
    if not os.path.exists(file):
        os.makedirs(file)


file = 'h_data/h_photos'
h_class = [cla for cla in os.listdir(file) if ".txt" not in cla]
# cla for cla in os.listdir(file) if ".txt" not in cla
#  lookup file All the files in the directory , That is to say, the classification is based on file The number of folders under the directory determines the category of classification 

#  Create training set 
mkfile('h_data/train')
for cla in helmet_class:
    mkfile('h_data/train/'+cla)

#  Create test set 
mkfile('h_data/test')

#  Create validation set 
mkfile('h_data/val')
for cla in helmet_class:
    mkfile('h_data/val/'+cla)

split_rate = 0.1
for cla in h_class:
    cla_path = file + '/' + cla + '/'
    images = os.listdir(cla_path) # images Stored the names of all pictures 
    num = len(images)
    eval_index = random.sample(images, k=int(num*split_rate))
    # print(eval_index)
    #  from images Extract from the list k individual ,num For the total number of pictures ,split_rate For custom scale 
    #  The randomly extracted data is the verification set , eval_index Name the picture 
    rest_images = list(set(images).difference(set(eval_index)))
    # take images and eval_index The difference between the set , Get not in eval_index Picture in 
    test_index = random.sample(rest_images, k=int(num*split_rate))
    #  Get the test set data 
    for index, image in enumerate(images):
        # enumerate() Function traverses an object in an indexed sequence 
        if image in eval_index:
            image_path = cla_path + image
            new_path = 'h_data/val/' + cla
            copy(image_path, new_path)
        elif image in test_index:
            image_path = cla_path + image
            new_path = 'h_data/test/'
            copy(image_path, new_path)
        else:
            image_path = cla_path + image
            new_path = 'h_data/train/' + cla
            copy(image_path, new_path)
        print("\r[{}] processing [{}/{}]".format(cla, index+1, num), end="")  # processing bar
    print()

print("processing done!")

In particular
1.file： Change to the directory where your data resides

2. The author's directory structure is as follows ：

3. Data split file (h_split.py) Same as h_data Put it in the same directory