Spaces:
Running
Running
| import torch | |
| import torchaudio | |
| import random | |
| import itertools | |
| import numpy as np | |
| from tools.mix import mix | |
| def normalize_wav(waveform): | |
| waveform = waveform - torch.mean(waveform) | |
| waveform = waveform / (torch.max(torch.abs(waveform)) + 1e-8) | |
| return waveform * 0.5 | |
| def pad_wav(waveform, segment_length): | |
| waveform_length = len(waveform) | |
| if segment_length is None or waveform_length == segment_length: | |
| return waveform | |
| elif waveform_length > segment_length: | |
| return waveform[:segment_length] | |
| else: | |
| pad_wav = torch.zeros(segment_length - waveform_length).to(waveform.device) | |
| waveform = torch.cat([waveform, pad_wav]) | |
| return waveform | |
| def _pad_spec(fbank, target_length=1024): | |
| batch, n_frames, channels = fbank.shape | |
| p = target_length - n_frames | |
| if p > 0: | |
| pad = torch.zeros(batch, p, channels).to(fbank.device) | |
| fbank = torch.cat([fbank, pad], 1) | |
| elif p < 0: | |
| fbank = fbank[:, :target_length, :] | |
| if channels % 2 != 0: | |
| fbank = fbank[:, :, :-1] | |
| return fbank | |
| def read_wav_file(filename, segment_length): | |
| waveform, sr = torchaudio.load(filename) # Faster!!! | |
| try: | |
| waveform = torchaudio.functional.resample(waveform, orig_freq=sr, new_freq=16000)[0] | |
| except: | |
| print ("0 length wav encountered. Setting to random:", filename) | |
| waveform = torch.rand(160000) | |
| try: | |
| waveform = normalize_wav(waveform) | |
| except: | |
| print ("Exception normalizing:", filename) | |
| waveform = torch.ones(160000) | |
| waveform = pad_wav(waveform, segment_length).unsqueeze(0) | |
| waveform = waveform / torch.max(torch.abs(waveform)) | |
| waveform = 0.5 * waveform | |
| return waveform | |
| def get_mel_from_wav(audio, _stft): | |
| audio = torch.nan_to_num(torch.clip(audio, -1, 1)) | |
| audio = torch.autograd.Variable(audio, requires_grad=False) | |
| melspec, log_magnitudes_stft, energy = _stft.mel_spectrogram(audio) | |
| return melspec, log_magnitudes_stft, energy | |
| def wav_to_fbank(paths, target_length=1024, fn_STFT=None): | |
| assert fn_STFT is not None | |
| waveform = torch.cat([read_wav_file(path, target_length * 160) for path in paths], 0) # hop size is 160 | |
| fbank, log_magnitudes_stft, energy = get_mel_from_wav(waveform, fn_STFT) | |
| fbank = fbank.transpose(1, 2) | |
| log_magnitudes_stft = log_magnitudes_stft.transpose(1, 2) | |
| fbank, log_magnitudes_stft = _pad_spec(fbank, target_length), _pad_spec( | |
| log_magnitudes_stft, target_length | |
| ) | |
| return fbank, log_magnitudes_stft, waveform | |
| def uncapitalize(s): | |
| if s: | |
| return s[:1].lower() + s[1:] | |
| else: | |
| return "" | |
| def mix_wavs_and_captions(path1, path2, caption1, caption2, target_length=1024): | |
| sound1 = read_wav_file(path1, target_length * 160)[0].numpy() | |
| sound2 = read_wav_file(path2, target_length * 160)[0].numpy() | |
| mixed_sound = mix(sound1, sound2, 0.5, 16000).reshape(1, -1) | |
| mixed_caption = "{} and {}".format(caption1, uncapitalize(caption2)) | |
| return mixed_sound, mixed_caption | |
| def augment(paths, texts, num_items=4, target_length=1024): | |
| mixed_sounds, mixed_captions = [], [] | |
| combinations = list(itertools.combinations(list(range(len(texts))), 2)) | |
| random.shuffle(combinations) | |
| if len(combinations) < num_items: | |
| selected_combinations = combinations | |
| else: | |
| selected_combinations = combinations[:num_items] | |
| for (i, j) in selected_combinations: | |
| new_sound, new_caption = mix_wavs_and_captions(paths[i], paths[j], texts[i], texts[j], target_length) | |
| mixed_sounds.append(new_sound) | |
| mixed_captions.append(new_caption) | |
| waveform = torch.tensor(np.concatenate(mixed_sounds, 0)) | |
| waveform = waveform / torch.max(torch.abs(waveform)) | |
| waveform = 0.5 * waveform | |
| return waveform, mixed_captions | |
| def augment_wav_to_fbank(paths, texts, num_items=4, target_length=1024, fn_STFT=None): | |
| assert fn_STFT is not None | |
| waveform, captions = augment(paths, texts) | |
| fbank, log_magnitudes_stft, energy = get_mel_from_wav(waveform, fn_STFT) | |
| fbank = fbank.transpose(1, 2) | |
| log_magnitudes_stft = log_magnitudes_stft.transpose(1, 2) | |
| fbank, log_magnitudes_stft = _pad_spec(fbank, target_length), _pad_spec( | |
| log_magnitudes_stft, target_length | |
| ) | |
| return fbank, log_magnitudes_stft, waveform, captions |