New Model, Many Changes

datasets/urban_8k.py

@@ -0,0 +1,95 @@
+import pickle
+from collections import defaultdict
+from pathlib import Path
+
+import librosa as librosa
+from torch.utils.data import Dataset
+import torch
+
+import variables as V
+from ml_lib.modules.util import F_x
+
+
+class BinaryMasksDataset(Dataset):
+    _to_label = defaultdict(lambda: -1)
+    _to_label.update(dict(clear=V.CLEAR, mask=V.MASK))
+
+    @property
+    def sample_shape(self):
+        return self[0][0].shape
+
+    def __init__(self, data_root, setting, mel_transforms, transforms=None, stretch_dataset=False,
+                 use_preprocessed=True):
+        self.use_preprocessed = use_preprocessed
+        self.stretch = stretch_dataset
+        assert isinstance(setting, str), f'Setting has to be a string, but was: {type(setting)}.'
+        assert setting in V.DATA_OPTIONS, f'Setting must match one of: {V.DATA_OPTIONS}.'
+        super(BinaryMasksDataset, self).__init__()
+
+        self.data_root = Path(data_root)
+        self.setting = setting
+        self._wav_folder = self.data_root / 'wav'
+        self._mel_folder = self.data_root / 'mel'
+        self.container_ext = '.pik'
+        self._mel_transform = mel_transforms
+
+        self._labels = self._build_labels()
+        self._wav_files = list(sorted(self._labels.keys()))
+        self._transforms = transforms or F_x(in_shape=None)
+
+    def _build_labels(self):
+        labeldict = dict()
+        with open(Path(self.data_root) / 'lab' / 'labels.csv', mode='r') as f:
+            # Exclude the header
+            _ = next(f)
+            for row in f:
+                if self.setting not in row:
+                    continue
+                filename, label = row.strip().split(',')
+                labeldict[filename] = self._to_label[label.lower()] if not self.setting == 'test' else filename
+        if self.stretch and self.setting == V.DATA_OPTIONS.train:
+            additional_dict = ({f'X{key}': val for key, val in labeldict.items()})
+            additional_dict.update({f'XX{key}': val for key, val in labeldict.items()})
+            additional_dict.update({f'XXX{key}': val for key, val in labeldict.items()})
+            labeldict.update(additional_dict)
+
+        # Delete File if one exists.
+        if not self.use_preprocessed:
+            for key in labeldict.keys():
+                try:
+                    (self._mel_folder / (key.replace('.wav', '') + self.container_ext)).unlink()
+                except FileNotFoundError:
+                    pass
+
+        return labeldict
+
+    def __len__(self):
+        return len(self._labels)
+
+    def _compute_or_retrieve(self, filename):
+
+        if not (self._mel_folder / (filename + self.container_ext)).exists():
+            raw_sample, sr = librosa.core.load(self._wav_folder / (filename.replace('X', '') + '.wav'))
+            mel_sample = self._mel_transform(raw_sample)
+            self._mel_folder.mkdir(exist_ok=True, parents=True)
+            with (self._mel_folder / (filename + self.container_ext)).open(mode='wb') as f:
+                pickle.dump(mel_sample, f, protocol=pickle.HIGHEST_PROTOCOL)
+
+        with (self._mel_folder / (filename + self.container_ext)).open(mode='rb') as f:
+            mel_sample = pickle.load(f, fix_imports=True)
+        return mel_sample
+
+    def __getitem__(self, item):
+
+        key: str = list(self._labels.keys())[item]
+        filename = key.replace('.wav', '')
+        mel_sample = self._compute_or_retrieve(filename)
+        label = self._labels[key]
+
+        transformed_samples = self._transforms(mel_sample)
+
+        if self.setting != V.DATA_OPTIONS.test:
+            # In test, filenames instead of labels are returned. This is a little hacky though.
+            label = torch.as_tensor(label, dtype=torch.float)
+
+        return transformed_samples, label