Spaces:
Running
Running
| import os | |
| import glob | |
| import torch | |
| import torchaudio | |
| from tqdm import tqdm | |
| class StyleDataset(torch.utils.data.Dataset): | |
| def __init__( | |
| self, | |
| audio_dir: str, | |
| subset: str = "train", | |
| sample_rate: int = 24000, | |
| length: int = 131072, | |
| ) -> None: | |
| super().__init__() | |
| self.audio_dir = audio_dir | |
| self.subset = subset | |
| self.sample_rate = sample_rate | |
| self.length = length | |
| self.style_dirs = glob.glob(os.path.join(audio_dir, subset, "*")) | |
| self.style_dirs = [sd for sd in self.style_dirs if os.path.isdir(sd)] | |
| self.num_classes = len(self.style_dirs) | |
| self.class_labels = {"broadcast" : 0, "telephone": 1, "neutral": 2, "bright": 3, "warm": 4} | |
| self.examples = [] | |
| for n, style_dir in enumerate(self.style_dirs): | |
| # get all files in style dir | |
| style_filepaths = glob.glob(os.path.join(style_dir, "*.wav")) | |
| style_name = os.path.basename(style_dir) | |
| for style_filepath in tqdm(style_filepaths, ncols=120): | |
| # load audio file | |
| x, sr = torchaudio.load(style_filepath) | |
| # sum to mono if needed | |
| if x.shape[0] > 1: | |
| x = x.mean(dim=0, keepdim=True) | |
| # resample | |
| if sr != self.sample_rate: | |
| x = torchaudio.transforms.Resample(sr, self.sample_rate)(x) | |
| # crop length after resample | |
| if x.shape[-1] >= self.length: | |
| x = x[...,:self.length] | |
| # store example | |
| example = (x, self.class_labels[style_name]) | |
| self.examples.append(example) | |
| print(f"Loaded {len(self.examples)} examples for {subset} subset.") | |
| def __len__(self): | |
| return len(self.examples) | |
| def __getitem__(self, idx): | |
| example = self.examples[idx] | |
| x = example[0] | |
| y = example[1] | |
| return x, y | |