[segmentation] Apply augmentation transforms

pyproject.toml

@@ -429,6 +429,13 @@ lwnet = "mednet.libs.segmentation.config.models.lwnet"
 m2unet = "mednet.libs.segmentation.config.models.m2unet"
 unet = "mednet.libs.segmentation.config.models.unet"
 
+affine = "mednet.libs.common.config.augmentations.affine"
+elastic = "mednet.libs.common.config.augmentations.elastic"
+hflip = "mednet.libs.common.config.augmentations.hflip"
+jitter = "mednet.libs.common.config.augmentations.jitter"
+hflip-affine = "mednet.libs.common.config.augmentations.hflip_affine"
+hflip-jitter-affine = "mednet.libs.common.config.augmentations.hflip_jitter_affine"
+
 # chase-db1 - retinography
 chasedb1 = "mednet.libs.segmentation.config.data.chasedb1.first_annotator"
 chasedb1-2nd = "mednet.libs.segmentation.config.data.chasedb1.second_annotator"

src/mednet/libs/segmentation/models/segmentation_model.py

@@ -90,11 +90,11 @@ class SegmentationModel(Model):
         self.normalizer = make_z_normalizer(dataloader)
 
     def training_step(self, batch, batch_idx):
-        images = batch[0]["image"]
-        ground_truths = batch[0]["target"]
-        masks = batch[0]["mask"]
+        images = self.augmentation_transforms(batch[0]["image"])
+        ground_truths = self.augmentation_transforms(batch[0]["target"])
+        masks = self.augmentation_transforms(batch[0]["mask"])
 
-        outputs = self(self._augmentation_transforms(images))
+        outputs = self(images)
         return self._train_loss(outputs, ground_truths, masks)
 
     def validation_step(self, batch, batch_idx):

src/mednet/libs/segmentation/scripts/experiment.py

@@ -50,6 +50,7 @@ def experiment(
     seed,
     parallel,
     monitoring_interval,
+    augmentations,
     **_,
 ):  # numpydoc ignore=PR01
     r"""Run a complete experiment, from training, to prediction and evaluation.
@@ -85,6 +86,7 @@ def experiment(
         seed=seed,
         parallel=parallel,
         monitoring_interval=monitoring_interval,
+        augmentations=augmentations,
     )
     train_stop_timestamp = datetime.now()