Audio IO Ready

datasets/binar_masks.py

@@ -1,3 +1,4 @@
+import pickle
 from collections import defaultdict
 from pathlib import Path
 
@@ -5,23 +6,28 @@ import librosa as librosa
 from torch.utils.data import Dataset
 
 import variables as V
+from ml_lib.modules.utils import F_x
 
 
-class BinaryMasks(Dataset):
-    _to_label = defaultdict(-1)
+class BinaryMasksDataset(Dataset):
+    _to_label = defaultdict(lambda: -1)
     _to_label['clear'] = V.CLEAR
     _to_label['mask'] = V.MASK
+    settings = ['test', 'devel', 'train']
 
-    def __init__(self, data_root, setting):
-        assert isinstance(setting, str)
-        assert setting in ['test', 'devel', 'train']
-        super(BinaryMasks, self).__init__()
+    def __init__(self, data_root, setting, transforms=None):
+        assert isinstance(setting, str), f'Setting has to be a string, but was: {self.settings}.'
+        assert setting in self.settings, f'Setting must match one of: {self.settings}.'
+        assert callable(transforms) or None, f'Transforms has to be callable, but was: {transforms}'
+        super(BinaryMasksDataset, self).__init__()
 
         self.data_root = Path(data_root)
         self.setting = setting
+        self._transforms = transforms or F_x()
         self._labels = self._build_labels()
         self._wav_folder = self.data_root / 'wav'
-        self._files = list(sorted(self._labels.keys()))
+        self._wav_files = list(sorted(self._labels.keys()))
+        self._mel_folder = self.data_root / 'raw_mel'
 
     def _build_labels(self):
         with open(Path(self.data_root) / 'lab' / 'labels.csv', mode='r') as f:
@@ -39,15 +45,16 @@ class BinaryMasks(Dataset):
         return len(self._labels)
 
     def __getitem__(self, item):
-        key = self._files[item]
-        sample = librosa.core.load(self._wav_folder / self._files[key])
+        key = self._wav_files[item]
+        filename = key[:-4] + '.pik'
+
+        if not (self._mel_folder / filename).exists():
+            raw_sample, sr = librosa.core.load(self._wav_folder / self._wav_files[item])
+            transformed_sample = self._transforms(raw_sample)
+            self._mel_folder.mkdir(exist_ok=True, parents=True)
+            with (self._mel_folder / filename).open(mode='wb') as f:
+                pickle.dump(transformed_sample, f, protocol=pickle.HIGHEST_PROTOCOL)
+        with (self._mel_folder / filename).open(mode='rb') as f:
+            sample = pickle.load(f, fix_imports=True)
         label = self._labels[key]
         return sample, label
-
-
-
-
-
-
-
-