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