import os import random import warnings import soundfile as sf import torch from numpy import trim_zeros from speechbrain.pretrained import EncoderClassifier from torch.multiprocessing import Manager from torch.multiprocessing import Process from torch.utils.data import Dataset from tqdm import tqdm from Preprocessing.AudioPreprocessor import AudioPreprocessor from Preprocessing.TextFrontend import ArticulatoryCombinedTextFrontend class AlignerDataset(Dataset): def __init__(self, path_to_transcript_dict, cache_dir, lang, loading_processes=8, # careful with the amount of processes if you use silence removal, only as many processes as you have cores min_len_in_seconds=1, max_len_in_seconds=20, cut_silences=True, rebuild_cache=False, verbose=False, device="cpu", phone_input=False): os.makedirs(cache_dir, exist_ok=True) if not os.path.exists(os.path.join(cache_dir, "aligner_train_cache.pt")) or rebuild_cache: if cut_silences: torch.set_num_threads(1) torch.hub.load(repo_or_dir='snakers4/silero-vad', model='silero_vad', force_reload=False, onnx=False, verbose=False) # download and cache for it to be loaded and used later torch.set_grad_enabled(True) resource_manager = Manager() self.path_to_transcript_dict = resource_manager.dict(path_to_transcript_dict) key_list = list(self.path_to_transcript_dict.keys()) with open(os.path.join(cache_dir, "files_used.txt"), encoding='utf8', mode="w") as files_used_note: files_used_note.write(str(key_list)) random.shuffle(key_list) # build cache print("... building dataset cache ...") self.datapoints = resource_manager.list() # make processes key_splits = list() process_list = list() for i in range(loading_processes): key_splits.append(key_list[i * len(key_list) // loading_processes:(i + 1) * len(key_list) // loading_processes]) for key_split in key_splits: process_list.append( Process(target=self.cache_builder_process, args=(key_split, lang, min_len_in_seconds, max_len_in_seconds, cut_silences, verbose, "cpu", phone_input), daemon=True)) process_list[-1].start() for process in process_list: process.join() self.datapoints = list(self.datapoints) tensored_datapoints = list() # we had to turn all of the tensors to numpy arrays to avoid shared memory # issues. Now that the multi-processing is over, we can convert them back # to tensors to save on conversions in the future. print("Converting into convenient format...") norm_waves = list() filepaths = list() for datapoint in tqdm(self.datapoints): tensored_datapoints.append([torch.Tensor(datapoint[0]), torch.LongTensor(datapoint[1]), torch.Tensor(datapoint[2]), torch.LongTensor(datapoint[3])]) norm_waves.append(torch.Tensor(datapoint[-2])) filepaths.append(datapoint[-1]) self.datapoints = tensored_datapoints # add speaker embeddings self.speaker_embeddings = list() speaker_embedding_func_ecapa = EncoderClassifier.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb", run_opts={"device": str(device)}, savedir="Models/SpeakerEmbedding/speechbrain_speaker_embedding_ecapa") with torch.no_grad(): for wave in tqdm(norm_waves): self.speaker_embeddings.append(speaker_embedding_func_ecapa.encode_batch(wavs=wave.to(device).unsqueeze(0)).squeeze().cpu()) # save to cache torch.save((self.datapoints, norm_waves, self.speaker_embeddings, filepaths), os.path.join(cache_dir, "aligner_train_cache.pt")) else: # just load the datapoints from cache self.datapoints = torch.load(os.path.join(cache_dir, "aligner_train_cache.pt"), map_location='cpu') self.speaker_embeddings = self.datapoints[2] self.datapoints = self.datapoints[0] self.tf = ArticulatoryCombinedTextFrontend(language=lang) print(f"Prepared an Aligner dataset with {len(self.datapoints)} datapoints in {cache_dir}.") def cache_builder_process(self, path_list, lang, min_len, max_len, cut_silences, verbose, device, phone_input): process_internal_dataset_chunk = list() tf = ArticulatoryCombinedTextFrontend(language=lang) _, sr = sf.read(path_list[0]) ap = AudioPreprocessor(input_sr=sr, output_sr=16000, melspec_buckets=80, hop_length=256, n_fft=1024, cut_silence=cut_silences, device=device) for path in tqdm(path_list): if self.path_to_transcript_dict[path].strip() == "": continue wave, sr = sf.read(path) dur_in_seconds = len(wave) / sr if not (min_len <= dur_in_seconds <= max_len): if verbose: print(f"Excluding {path} because of its duration of {round(dur_in_seconds, 2)} seconds.") continue try: with warnings.catch_warnings(): warnings.simplefilter("ignore") # otherwise we get tons of warnings about an RNN not being in contiguous chunks norm_wave = ap.audio_to_wave_tensor(normalize=True, audio=wave) except ValueError: continue dur_in_seconds = len(norm_wave) / 16000 if not (min_len <= dur_in_seconds <= max_len): if verbose: print(f"Excluding {path} because of its duration of {round(dur_in_seconds, 2)} seconds.") continue norm_wave = torch.tensor(trim_zeros(norm_wave.numpy())) # raw audio preprocessing is done transcript = self.path_to_transcript_dict[path] try: cached_text = tf.string_to_tensor(transcript, handle_missing=False, input_phonemes=phone_input, path_to_wavfile=path).squeeze(0).cpu().numpy() except KeyError: tf.string_to_tensor(transcript, handle_missing=True, input_phonemes=phone_input, path_to_wavfile=path).squeeze(0).cpu().numpy() continue # we skip sentences with unknown symbols cached_text_len = torch.LongTensor([len(cached_text)]).numpy() cached_speech = ap.audio_to_mel_spec_tensor(audio=norm_wave, normalize=False, explicit_sampling_rate=16000).transpose(0, 1).cpu().numpy() cached_speech_len = torch.LongTensor([len(cached_speech)]).numpy() process_internal_dataset_chunk.append([cached_text, cached_text_len, cached_speech, cached_speech_len, norm_wave.cpu().detach().numpy(), path]) self.datapoints += process_internal_dataset_chunk def __getitem__(self, index): text_vector = self.datapoints[index][0] tokens = list() for vector in text_vector: if vector[19] == 0: # we don't include word boundaries when performing alignment, since they are not always present in audio. for phone in self.tf.phone_to_vector: if vector.numpy().tolist()[11:] == self.tf.phone_to_vector[phone][11:]: # the first 10 dimensions are for modifiers, so we ignore those when trying to find the phoneme in the ID lookup tokens.append(self.tf.phone_to_id[phone]) # this is terribly inefficient, but it's fine break tokens = torch.LongTensor(tokens) return tokens, \ torch.LongTensor([len(tokens)]), \ self.datapoints[index][2], \ self.datapoints[index][3], \ self.speaker_embeddings[index] def __len__(self): return len(self.datapoints)