Connected processing to loading the dataset.

src/caked/base.py

@@ -25,7 +25,7 @@ def load(self, datapath, datatype):
         pass
 
     @abstractmethod
-    def process(self):
+    def process(self, paths, datatype):
         pass
 
     @abstractmethod

src/caked/dataloader.py

@@ -15,6 +15,9 @@
 
 from .base import AbstractDataLoader, AbstractDataset
 
+np.random.seed(42)
+TRANSFORM_OPTIONS = ["rescale", "normalise", "gaussianblur", "shiftmin"]
+
 
 class DiskDataLoader(AbstractDataLoader):
     def __init__(
@@ -24,11 +27,13 @@ def __init__(
         training: bool = True,
         classes: list[str] | None = None,
         pipeline: str = "disk",
+        transformations: str | None = None,
     ) -> None:
         self.dataset_size = dataset_size
         self.save_to_disk = save_to_disk
         self.training = training
         self.pipeline = pipeline
+        self.transformations = transformations
         if classes is None:
             self.classes = []
         else:
@@ -69,10 +74,27 @@ def load(self, datapath, datatype) -> None:
         if self.dataset_size is not None:
             paths = paths[: self.dataset_size]
 
-        self.dataset = DiskDataset(paths=paths, datatype=datatype)
+        if self.transformations is None:
+            self.dataset = DiskDataset(paths=paths, datatype=datatype)
+        else:
+            self.dataset = self.process(paths=paths, datatype=datatype)
 
-    def process(self):
-        return super().process()
+    def process(self, paths: list[str], datatype: str):
+        if self.transformations is None:
+            msg = "No processing to do as no transformations were provided."
+            raise RuntimeError(msg)
+        transforms = self.transformations.split(",")
+        rescale, normalise, gaussianblur, shiftmin = np.in1d(
+            TRANSFORM_OPTIONS, transforms
+        )
+        return DiskDataset(
+            paths=paths,
+            datatype=datatype,
+            rescale=rescale,
+            normalise=normalise,
+            gaussianblur=gaussianblur,
+            shiftmin=shiftmin,
+        )
 
     def get_loader(self, batch_size: int, split_size: float | None = None):
         if self.training:
@@ -130,7 +152,6 @@ def __init__(
         self.rescale = rescale
         self.normalise = normalise
         self.gaussianblur = gaussianblur
-        self.rescale = rescale
         self.transform = input_transform
         self.datatype = datatype
         self.shiftmin = shiftmin

tests/test_disk_io.py

@@ -18,6 +18,8 @@
 DISK_CLASSES_MISSING = ["2b3a", "1b23"]
 DISK_CLASSES_NONE = None
 DATATYPE_MRC = "mrc"
+TRANSFORM_ALL = "rescale,normalise,gaussianblur,shiftmin"
+TRANSFORM_SOME = "rescale,gaussianblur"
 
 
 def test_class_instantiation():
@@ -127,3 +129,33 @@ def test_get_loader_training_fail():
         torch_loader_train, torch_loader_val = test_loader.get_loader(
             split_size=1, batch_size=64
         )
+
+
+def test_processing_data_all_transforms():
+    test_loader = DiskDataLoader(
+        pipeline=DISK_PIPELINE,
+        classes=DISK_CLASSES_FULL,
+        dataset_size=DATASET_SIZE_ALL,
+        training=True,
+        transformations=TRANSFORM_ALL,
+    )
+    test_loader.load(datapath=TEST_DATA_MRC, datatype=DATATYPE_MRC)
+    assert test_loader.dataset.normalise
+    assert test_loader.dataset.shiftmin
+    assert test_loader.dataset.gaussianblur
+    assert test_loader.dataset.rescale
+
+
+def test_processing_data_some_transforms():
+    test_loader = DiskDataLoader(
+        pipeline=DISK_PIPELINE,
+        classes=DISK_CLASSES_FULL,
+        dataset_size=DATASET_SIZE_ALL,
+        training=True,
+        transformations=TRANSFORM_SOME,
+    )
+    test_loader.load(datapath=TEST_DATA_MRC, datatype=DATATYPE_MRC)
+    assert not test_loader.dataset.normalise
+    assert not test_loader.dataset.shiftmin
+    assert test_loader.dataset.gaussianblur
+    assert test_loader.dataset.rescale