Spaces:
Building
Building
| from torch.utils.data import Dataset | |
| from pathlib import Path | |
| from typing import Optional | |
| import torch | |
| from torch.utils.data import default_collate | |
| from typing import Tuple | |
| from functools import partial | |
| from gyraudio.audio_separation.properties import ( | |
| AUG_AWGN, AUG_RESCALE, AUG_TRIM, LENGTHS, LENGTH_DIVIDER, TRIM_PROB | |
| ) | |
| class AudioDataset(Dataset): | |
| def __init__( | |
| self, | |
| data_path: Path, | |
| augmentation_config: dict = {}, | |
| snr_filter: Optional[float] = None, | |
| debug: bool = False | |
| ): | |
| self.debug = debug | |
| self.data_path = data_path | |
| self.augmentation_config = augmentation_config | |
| self.snr_filter = snr_filter | |
| self.load_data() | |
| self.length = len(self.file_list) | |
| self.collate_fn = None | |
| if AUG_TRIM in self.augmentation_config: | |
| self.collate_fn = partial(collate_fn_generic, | |
| lengths_lim=self.augmentation_config[AUG_TRIM][LENGTHS], | |
| length_divider=self.augmentation_config[AUG_TRIM][LENGTH_DIVIDER], | |
| trim_prob=self.augmentation_config[AUG_TRIM][TRIM_PROB]) | |
| def filter_data(self, snr): | |
| if self.snr_filter is None: | |
| return True | |
| if snr in self.snr_filter: | |
| return True | |
| else: | |
| return False | |
| def load_data(self): | |
| raise NotImplementedError("load_data method must be implemented") | |
| def augment_data(self, mixed_audio_signal, clean_audio_signal, noise_audio_signal): | |
| if AUG_RESCALE in self.augmentation_config: | |
| current_amplitude = 0.5 + 1.5*torch.rand(1, device=mixed_audio_signal.device) | |
| # logging.debug(current_amplitude) | |
| mixed_audio_signal *= current_amplitude | |
| noise_audio_signal *= current_amplitude | |
| clean_audio_signal *= current_amplitude | |
| if AUG_AWGN in self.augmentation_config: | |
| # noise_std = self.augmentation_config[AUG_AWGN]["noise_std"] | |
| noise_std = 0.01 | |
| current_noise_std = torch.randn(1) * noise_std | |
| # logging.debug(current_noise_std) | |
| extra_awgn = torch.randn(mixed_audio_signal.shape, device=mixed_audio_signal.device) * current_noise_std | |
| mixed_audio_signal = mixed_audio_signal+extra_awgn | |
| # Open question: should we add noise to the noise signal aswell? | |
| return mixed_audio_signal, clean_audio_signal, noise_audio_signal | |
| def __len__(self): | |
| return self.length | |
| def __getitem__(self, idx: int) -> torch.Tensor: | |
| raise NotImplementedError("__getitem__ method must be implemented") | |
| def collate_fn_generic(batch, lengths_lim, length_divider=1024, trim_prob=0.5) -> Tuple[torch.Tensor, torch.Tensor]: | |
| """Collate function to allow trimming (=crop the time dimension) of the signals in a batch. | |
| Args: | |
| batch (list): A list of tuples (triplets), where each tuple contain: | |
| - mixed_audio_signal | |
| - clean_audio_signal | |
| - noise_audio_signal | |
| lengths_lim (list) : A list of containing a minimum length (0) and a maximum length (1) | |
| length_divider (int) : has to be a trimmed length divider | |
| trim_prob (float) : trimming probability | |
| Returns: | |
| - Tensor: A batch of mixed_audio_signal, trimmed to the same length. | |
| - Tensor: A batch of clean_audio_signal | |
| - Tensor: A batch of noise_audio_signal | |
| """ | |
| # Find the length of the shortest signal in the batch | |
| mixed_audio_signal, clean_audio_signal, noise_audio_signal = default_collate(batch) | |
| length = mixed_audio_signal[0].shape[-1] | |
| min_length, max_length = lengths_lim | |
| take_full_signal = torch.rand(1) > trim_prob | |
| if not take_full_signal: | |
| start = torch.randint(0, length-min_length, (1,)) | |
| trim_length = torch.randint(min_length, min(max_length, length-start-1)+1, (1,)) | |
| trim_length = trim_length-trim_length % length_divider | |
| end = start + trim_length | |
| else: | |
| start = 0 | |
| end = length - length % length_divider | |
| mixed_audio_signal = mixed_audio_signal[..., start:end] | |
| clean_audio_signal = clean_audio_signal[..., start:end] | |
| noise_audio_signal = noise_audio_signal[..., start:end] | |
| return mixed_audio_signal, clean_audio_signal, noise_audio_signal | |