|
|
| import os
|
| import pandas as pd
|
| import scripts.config as config
|
| from sklearn.model_selection import train_test_split
|
|
|
| nameTrain = config.source + 'train.txt'
|
| nameVal = config.source + 'val.txt'
|
| percentual = 0.7
|
|
|
| def nomeOrigem(originalName):
|
| nome = os.path.join(config.images, originalName)
|
| return nome
|
|
|
| def nomeMascara(originalName):
|
| nome = os.path.join(config.masks, originalName)
|
| nome = nome.replace('.jpg', '.png')
|
| return nome
|
|
|
| print('Starting processing...')
|
| paths = [os.path.join(config.images, nome) for nome in os.listdir(config.images)]
|
| files = [arq for arq in paths if os.path.isfile(arq)]
|
| files = [arq for arq in files if arq.lower().endswith('.jpg')]
|
|
|
| df = pd.DataFrame({'nome_original': files})
|
|
|
| df['newNameA'] = [nomeOrigem(nome) for i, nome in enumerate(files)]
|
| df['newNameB'] = [nomeMascara(nome) for i, nome in enumerate(files)]
|
|
|
| grupo1, grupo2 = train_test_split(df, train_size=percentual, random_state=42)
|
|
|
| grupo1[['newNameA', 'newNameB']].to_csv(nameTrain, sep='\t', index=False, header=False)
|
| grupo2[['newNameA', 'newNameB']].to_csv(nameVal, sep='\t', index=False, header=False)
|
|
|
| print("Groups saved in '" + nameTrain + "' and '" + nameVal + "'")
|
|
|
