determine_file_lists.py

import os
import pickle
import utils.project_functions as pf
from sklearn.model_selection import train_test_split

#  Separação do dataset

# * Nessa seção, o dataset será separado da seguinte forma:
# > * 70% dos falantes serão separados para treinamento, 20% para validação e 10% para testes
# > * Serão utilizados na avaliação do desempenho de testes e validação somente as frases não utilizadas em nenhum outro conjunto
# > * Realizar testes/validação em outro momento também com frases que ocorrem no treinamento (não inéditas)

# Define os caminhos relevantes para este script
dataset_folder_path = os.path.join(f"{os.path.sep}datasets", "ntcd_timit")
clean_folder_path = os.path.join(dataset_folder_path, "Clean")
volunteers_path = os.path.join(clean_folder_path, "volunteers")

# Lista as pastas correspondentes a cada voluntário
spk_list = os.listdir(volunteers_path)

# Obtém as frases ditas por cada voluntário
spk_wav_dict = {}
for speaker in spk_list:
    path = os.path.join(volunteers_path, speaker, "straightcam")
    spk_wav_list = os.listdir(path)
    
    spk_wav_dict[speaker] = spk_wav_list

# Separa a lista de voluntários em treino (70%), validação (20%) e teste (10%)
spk_train_val_list, spk_test_list = train_test_split(spk_list, test_size = 0.1, random_state = 0)
spk_train_list, spk_val_list = train_test_split(spk_train_val_list, test_size = 0.2, random_state = 0)

# Listas de SNRs e ruídos de acordo com as pastas do dataset
SNR_classes_list = ['-5', '0', '5', '10', '15', '20']
Noise_classes_list = ['Babble', 'Cafe', 'Car', 'LR', 'Street', 'White']

# Constrói a lista com todos os arquivos a serem utilizados no treinamento          
train_files, train_wav_set = pf.ds_files_list(spk_train_list, 
                                              spk_wav_dict, 
                                              Noise_classes_list, 
                                              SNR_classes_list, 
                                              dataset_folder_path)

# Constrói a lista com todos os arquivos a serem utilizados na validação        
val_files, val_wav_set = pf.ds_files_list(spk_val_list, 
                                          spk_wav_dict, 
                                          Noise_classes_list, 
                                          SNR_classes_list, 
                                          dataset_folder_path, 
                                          exclude_wav = train_wav_set)

# Constrói a lista com todos os arquivos a serem utilizados nos teste   
test_files, test_wav_set = pf.ds_files_list(spk_test_list, 
                                            spk_wav_dict, 
                                            Noise_classes_list, 
                                            SNR_classes_list, 
                                            dataset_folder_path, 
                                            exclude_wav = train_wav_set | val_wav_set)

# Mostra a distribuição final dos arquivos nos datasets
print('Total de arquvivos para treino: ', len(train_files))
print('Sentenças únicas para treino: ', len(train_wav_set))
print('-'*50)
print('Total de arquvivos para validação: ', len(val_files))
print('Sentenças únicas para validação: ', len(val_wav_set))
print('-'*50)
print('Total de arquvivos para teste: ', len(test_files))
print('Sentenças únicas para teste: ', len(test_wav_set))

total_unique_sentences = len(train_wav_set) + len(val_wav_set) + len(test_wav_set)

print(f'Tamanho relativo do dataset de treinamento: {(len(train_wav_set)/total_unique_sentences):.2%}')
print(f'Tamanho relativo do dataset de validação: {(len(val_wav_set)/total_unique_sentences):.2%}')
print(f'Tamanho relativo do dataset de treinamento: {(len(test_wav_set)/total_unique_sentences):.2%}')

print("\nExemplo de par de arquivos a ser entregue à pipeline de pré-processamento")
print(train_files[0])

with open(os.path.join("data", "file_lists.pkl"), 'wb') as file:
    file_lists = {'train_files': train_files,
                  'val_files': val_files,
                  'test_files': test_files}
    pickle.dump(file_lists, file)