| 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) |
| 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) |
|
|
| 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 |
