restructuring to separate data processing from torch dataset implementations

mu_map/data/preprocessing.py

-import torch
-
-
-def norm_max(tensor: torch.Tensor):
-    return (tensor - tensor.min()) / (tensor.max() - tensor.min())
-
-
-class MaxNorm:
-    def __call__(self, tensor: torch.Tensor):
-        return norm_max(tensor)
-
-
-def norm_mean(tensor: torch.Tensor):
-    return tensor / tensor.mean()
-
-
-class MeanNorm:
-    def __call__(self, tensor: torch.Tensor):
-        return norm_mean(tensor)
-
-
-def norm_gaussian(tensor: torch.Tensor):
-    return (tensor - tensor.mean()) / tensor.std()
-
-
-class GaussianNorm:
-    def __call__(self, tensor: torch.Tensor):
-        return norm_gaussian(tensor)
-
-
-__all__ = [
-    norm_max.__name__,
-    norm_mean.__name__,
-    norm_gaussian.__name__,
-    MaxNorm.__name__,
-    MeanNorm.__name__,
-    GaussianNorm.__name__,
-]

mu_map/data/mock.py → mu_map/dataset/mock.py

-from mu_map.data.datasets import MuMapDataset
+from mu_map.dataset.default import MuMapDataset
 
 
 class MuMapMockDataset(MuMapDataset):

mu_map/dataset/normalization.py

+from torch import Tensor
+
+from mu_map.dataset.transform import Transform
+
+
+def norm_max(tensor: Tensor) -> Tensor:
+    return (tensor - tensor.min()) / (tensor.max() - tensor.min())
+
+
+class MaxNormTransform(Transform):
+    def __call__(self, inputs: Tensor, outputs_expected: Tensor) -> Tuple[Tensor, Tensor]:
+        return norm_max(inputs), outputs_expected
+
+
+def norm_mean(tensor: Tensor):
+    return tensor / tensor.mean()
+
+
+class MeanNormTransform(Transform):
+    def __call__(self, inputs: Tensor, outputs_expected: Tensor) -> Tuple[Tensor, Tensor]:
+        return norm_mean(inputs), outputs_expected
+
+
+def norm_gaussian(tensor: Tensor):
+    return (tensor - tensor.mean()) / tensor.std()
+
+
+class GaussianNormTransform(Transform):
+    def __call__(self, inputs: Tensor, outputs_expected: Tensor) -> Tuple[Tensor, Tensor]:
+        return norm_gaussian(inputs), outputs_expected
+
+
+__all__ = [
+    norm_max.__name__,
+    norm_mean.__name__,
+    norm_gaussian.__name__,
+    MaxNormTransform.__name__,
+    MeanNormTransform.__name__,
+    GaussianNormTransform.__name__,
+]

mu_map/data/patch_dataset.py → mu_map/dataset/patches.py

@@ -4,7 +4,7 @@ import random
 import numpy as np
 import torch
 
-from mu_map.data.datasets import MuMapDataset
+from mu_map.dataset.default import MuMapDataset
 
 
 class MuMapPatchDataset(MuMapDataset):