declare-lab/tango2

1import torch2import torchaudio3import random4import itertools5import numpy as np6from tools.mix import mix7 8 9def normalize_wav(waveform):10    waveform = waveform - torch.mean(waveform)11    waveform = waveform / (torch.max(torch.abs(waveform)) + 1e-8)12    return waveform * 0.513 14 15def pad_wav(waveform, segment_length):16    waveform_length = len(waveform)17    18    if segment_length is None or waveform_length == segment_length:19        return waveform20    elif waveform_length > segment_length:21        return waveform[:segment_length]22    else:23        pad_wav = torch.zeros(segment_length - waveform_length).to(waveform.device)24        waveform = torch.cat([waveform, pad_wav])25        return waveform26    27    28def _pad_spec(fbank, target_length=1024):29    batch, n_frames, channels = fbank.shape30    p = target_length - n_frames31    if p > 0:32        pad = torch.zeros(batch, p, channels).to(fbank.device)33        fbank = torch.cat([fbank, pad], 1)34    elif p < 0:35        fbank = fbank[:, :target_length, :]36 37    if channels % 2 != 0:38        fbank = fbank[:, :, :-1]39 40    return fbank41 42 43def read_wav_file(filename, segment_length):44    waveform, sr = torchaudio.load(filename)  # Faster!!!45    try:46        waveform = torchaudio.functional.resample(waveform, orig_freq=sr, new_freq=16000)[0]47    except:48        print ("0 length wav encountered. Setting to random:", filename)49        waveform = torch.rand(160000)50    51    try:52        waveform = normalize_wav(waveform)53    except:54        print ("Exception normalizing:", filename)55        waveform = torch.ones(160000)56    waveform = pad_wav(waveform, segment_length).unsqueeze(0)57    waveform = waveform / torch.max(torch.abs(waveform))58    waveform = 0.5 * waveform59    return waveform60 61 62def get_mel_from_wav(audio, _stft):63    audio = torch.nan_to_num(torch.clip(audio, -1, 1))64    audio = torch.autograd.Variable(audio, requires_grad=False)65    melspec, log_magnitudes_stft, energy = _stft.mel_spectrogram(audio)66    return melspec, log_magnitudes_stft, energy67 68 69def wav_to_fbank(paths, target_length=1024, fn_STFT=None):70    assert fn_STFT is not None71 72    waveform = torch.cat([read_wav_file(path, target_length * 160) for path in paths], 0)  # hop size is 16073 74    fbank, log_magnitudes_stft, energy = get_mel_from_wav(waveform, fn_STFT)75    fbank = fbank.transpose(1, 2)76    log_magnitudes_stft = log_magnitudes_stft.transpose(1, 2)77 78    fbank, log_magnitudes_stft = _pad_spec(fbank, target_length), _pad_spec(79        log_magnitudes_stft, target_length80    )81 82    return fbank, log_magnitudes_stft, waveform83 84 85def uncapitalize(s):86    if s:87        return s[:1].lower() + s[1:]88    else:89        return ""90 91    92def mix_wavs_and_captions(path1, path2, caption1, caption2, target_length=1024):93    sound1 = read_wav_file(path1, target_length * 160)[0].numpy()94    sound2 = read_wav_file(path2, target_length * 160)[0].numpy()95    mixed_sound = mix(sound1, sound2, 0.5, 16000).reshape(1, -1)96    mixed_caption = "{} and {}".format(caption1, uncapitalize(caption2))97    return mixed_sound, mixed_caption98 99 100def augment(paths, texts, num_items=4, target_length=1024):101    mixed_sounds, mixed_captions = [], []102    combinations = list(itertools.combinations(list(range(len(texts))), 2))103    random.shuffle(combinations)104    if len(combinations) < num_items:105        selected_combinations = combinations106    else:107        selected_combinations = combinations[:num_items]108        109    for (i, j) in selected_combinations:110        new_sound, new_caption = mix_wavs_and_captions(paths[i], paths[j], texts[i], texts[j], target_length)111        mixed_sounds.append(new_sound)112        mixed_captions.append(new_caption)113        114    waveform = torch.tensor(np.concatenate(mixed_sounds, 0))115    waveform = waveform / torch.max(torch.abs(waveform))116    waveform = 0.5 * waveform117    118    return waveform, mixed_captions119 120 121def augment_wav_to_fbank(paths, texts, num_items=4, target_length=1024, fn_STFT=None):122    assert fn_STFT is not None123    124    waveform, captions = augment(paths, texts)125    fbank, log_magnitudes_stft, energy = get_mel_from_wav(waveform, fn_STFT)126    fbank = fbank.transpose(1, 2)127    log_magnitudes_stft = log_magnitudes_stft.transpose(1, 2)128 129    fbank, log_magnitudes_stft = _pad_spec(fbank, target_length), _pad_spec(130        log_magnitudes_stft, target_length131    )132 133    return fbank, log_magnitudes_stft, waveform, captions