8/18/2020

Get list from dir and separate train and test (python function)


from sklearn.model_selection import train_test_split
import random
import os
import glob

def train_test_split_from_dir(origin_dir, test_size=0.2):
os.chdir(origin_dir)
#get list
data_list = []
for file in glob.glob("*.jpg"):
data_list.append(file)
random.shuffle(data_list)
train_json_list, test_json_list = train_test_split(data_list, test_size=test_size)

return train_json_list, test_json_list

No comments:

Post a Comment