File size: 1,236 Bytes
7b615ae |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
#encoding=utf8
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 + "'")
|