import random import sys from math import pi import librosa import torch import torchaudio from torch.utils.data import Dataset from tqdm import tqdm from data.audio.unsupervised_audio_dataset import UnsupervisedAudioDataset, load_audio from data.util import load_paths_from_cache, find_files_of_type, is_audio_file # Just all ones. from utils.util import opt_get def _integration_fn_fully_enabled(n): return torch.ones((n,)) # Randomly assigns up to 5 blocks of the output tensor the value '1'. Rest is zero def _integration_fn_spiky(n): fn = torch.zeros((n,)) spikes = random.randint(1,5) for _ in range(spikes): sz = random.randint(n//8, n//2) pos = random.randint(0, n) extent = min(n, sz+pos) fn[pos:extent] = 1 return fn # Uses a sinusoidal ramp up and down (of random length) to a peak which is held for a random duration. def _integration_fn_smooth(n): center = random.randint(1, n-2) max_duration=n-center-1 duration = random.randint(max_duration//4, max_duration) end = center+duration ramp_up_sz = random.randint(n//16,n//4) ramp_up = torch.sin(pi*torch.arange(0,ramp_up_sz)/(2*ramp_up_sz)) if ramp_up_sz > center: ramp_up = ramp_up[(ramp_up_sz-center):] ramp_up_sz = center ramp_down_sz = random.randint(n//16,n//4) ramp_down = torch.flip(torch.sin(pi*torch.arange(0,ramp_down_sz)/(2*ramp_down_sz)), dims=[0]) if ramp_down_sz > (n-end): ramp_down = ramp_down[:(n-end)] ramp_down_sz = n-end fn = torch.zeros((n,)) fn[(center-ramp_up_sz):center] = ramp_up fn[center:end] = 1 fn[end:(end+ramp_down_sz)] = ramp_down return fn def load_rir(path, sr, max_sz): rir = load_audio(path, sr).abs() if rir.shape[-1] > max_sz: rir = rir[:, :max_sz] rir = (rir / torch.norm(rir, p=2)).flip([1]) return rir ''' Wraps a unsupervised_audio_dataset and applies noise to the output clips, then provides labels depending on what noise was added. ''' class AudioWithNoiseDataset(Dataset): def __init__(self, opt): self.underlying_dataset = UnsupervisedAudioDataset(opt) self.env_noise_paths = load_paths_from_cache(opt['env_noise_paths'], opt['env_noise_cache']) self.music_paths = load_paths_from_cache(opt['music_paths'], opt['music_cache']) self.openair_paths = find_files_of_type('img', opt['openair_path'], qualifier=is_audio_file)[0] self.min_volume = opt_get(opt, ['min_noise_volume'], .2) self.max_volume = opt_get(opt, ['max_noise_volume'], .5) self.sampling_rate = self.underlying_dataset.sampling_rate self.use_gpu_for_reverb_compute = opt_get(opt, ['use_gpu_for_reverb_compute'], True) self.openair_kernels = None def load_openair_kernels(self): if self.use_gpu_for_reverb_compute and self.openair_kernels is None: # Load the openair reverbs as CUDA tensors. self.openair_kernels = [] for oa in self.openair_paths: self.openair_kernels.append(load_rir(oa, self.underlying_dataset.sampling_rate, self.underlying_dataset.sampling_rate*2).cuda()) def __getitem__(self, item): # Load on the fly to prevent GPU memory sharing across process errors. self.load_openair_kernels() out = self.underlying_dataset[item] clip = out['clip'] augpath = '' augvol = 0 try: # Randomly adjust clip volume, regardless of the selection, between clipvol = (random.random() * (.8-.5) + .5) clip = clip * clipvol label = random.randint(0, 4) # Current excludes GSM corruption. #label = 2 aug = torch.zeros_like(clip) if label > 0 and label < 4: # 0 is basically "leave it alone" augvol = (random.random() * (self.max_volume-self.min_volume) + self.min_volume) if label == 1: # Add environmental noise. augpath = random.choice(self.env_noise_paths) intg_fns = [_integration_fn_fully_enabled] elif label == 2: # Add music augpath = random.choice(self.music_paths) intg_fns = [_integration_fn_fully_enabled] augvol *= .5 # Music is often severely in the background. elif label == 3: # Add another voice. augpath = random.choice(self.underlying_dataset.audiopaths) intg_fns = [_integration_fn_smooth, _integration_fn_fully_enabled] aug = load_audio(augpath, self.underlying_dataset.sampling_rate) if aug.shape[1] > clip.shape[1]: n, cn = aug.shape[1], clip.shape[1] gap = n-cn placement = random.randint(0, gap) aug = aug[:, placement:placement+cn] aug = random.choice(intg_fns)(aug.shape[1]) * aug aug = aug * augvol if aug.shape[1] < clip.shape[1]: gap = clip.shape[1] - aug.shape[1] placement = random.randint(0, gap-1) aug = torch.nn.functional.pad(aug, (placement, gap-placement)) clip = clip + aug elif label == 4: # Perform reverb (to simulate being in a large room with an omni-mic). This is performed by convolving # impulse recordings from openair over the input clip. if self.use_gpu_for_reverb_compute: rir = random.choice(self.openair_kernels) else: augpath = random.choice(self.openair_paths) rir = load_rir(augpath, self.underlying_dataset.sampling_rate, clip.shape[-1]) clip = torch.nn.functional.pad(clip, (rir.shape[1]-1, 0)) if self.use_gpu_for_reverb_compute: clip = clip.cuda() clip = torch.nn.functional.conv1d(clip.unsqueeze(0), rir.unsqueeze(0)).squeeze(0).cpu() elif label == 5: # Apply the GSM codec to simulate cellular phone audio. clip = torchaudio.functional.apply_codec(clip, self.underlying_dataset.sampling_rate, format="gsm") except: print(f"Exception encountered processing {item}, re-trying because this is often just a failed aug.") print(sys.exc_info()) #raise # Uncomment to surface exceptions. return self[item] clip.clip_(-1, 1) out['clip'] = clip out['label'] = label out['aug'] = aug out['augpath'] = augpath out['augvol'] = augvol out['clipvol'] = clipvol return out def __len__(self): return len(self.underlying_dataset) if __name__ == '__main__': params = { 'mode': 'unsupervised_audio_with_noise', 'path': ['\\\\192.168.5.3\\rtx3080_audio\\split\\cleaned\\books0'], 'cache_path': 'E:\\audio\\remote-cache3.pth', 'sampling_rate': 22050, 'pad_to_samples': 80960, 'phase': 'train', 'n_workers': 0, 'batch_size': 4, 'extra_samples': 4, 'env_noise_paths': ['E:\\audio\\UrbanSound\\filtered', 'E:\\audio\\UrbanSound\\MSSND'], 'env_noise_cache': 'E:\\audio\\UrbanSound\\cache.pth', 'music_paths': ['E:\\audio\\music\\FMA\\fma_large', 'E:\\audio\\music\\maestro\\maestro-v3.0.0'], 'music_cache': 'E:\\audio\\music\\cache.pth', 'openair_path': 'D:\\data\\audio\\openair\\resampled' } from data import create_dataset, create_dataloader, util ds = create_dataset(params) dl = create_dataloader(ds, params) i = 0 for b in tqdm(dl): for b_ in range(b['clip'].shape[0]): #torchaudio.save(f'{i}_clip_{b_}_{b["label"][b_].item()}.wav', b['clip'][b_], ds.sampling_rate) #torchaudio.save(f'{i}_clip_{b_}_aug.wav', b['aug'][b_], ds.sampling_rate) print(f'{i} aug path: {b["augpath"][b_]} aug volume: {b["augvol"][b_]} clip volume: {b["clipvol"][b_]}') i += 1