added augmentation pipeline to docs

docs/doc/ml/models.md

@@ -6,7 +6,12 @@ hide:
 
 # Machine learning models
 
-The `qim3d` library aims to ease the creation of ML models for volumetric images
+The `qim3d` library aims to ease the creation of ML models for volumetric images.
+
+::: qim3d.ml.Augmentation
+    options:
+        members:
+            - augment
 
 ::: qim3d.ml.models
     options:

qim3d/ml/_augmentations.py

@@ -5,22 +5,26 @@ class Augmentation:
     Class for defining image augmentation transformations using the MONAI library.
         
     Args:
-        resize (str, optional): Specifies how the images should be reshaped to the appropriate size.
-        trainsform_train (str, optional): Level of transformation for the training set.
-        transform_validation (str, optional): Level of transformation for the validation set.
-        transform_test (str, optional): Level of transformation for the test set.
+        resize (str, optional): Specifies how the images should be reshaped to the appropriate size, either 'crop', 'resize', or 'padding'. Defaults to 'crop'.
+        trainsform_train (str, optional): Level of transformation for the training set, either 'light', 'moderate', 'heavy' or None. Defaults to 'moderate'.
+        transform_validation (str, optional): Level of transformation for the validation set, either 'light', 'moderate', 'heavy' or None. Defaults to None.
+        transform_test (str, optional): Level of transformation for the test set, either 'light', 'moderate', 'heavy' or None. Defaults to None.
 
     Raises:
-        ValueError: If the ´resize´ is neither 'crop', 'resize' nor 'padding'.
+        ValueError: If `resize` is neither 'crop', 'resize' nor 'padding'.
     
     Example:
-        my_augmentation = Augmentation(resize = 'crop', transform_train = 'light')
+        ```python
+        import qim3d
+
+        augmentation =  qim3d.ml.Augmentation(resize = 'crop', transform_train = 'light')
+        ```
     """
     
     def __init__(
             self, 
             resize: str = 'crop', 
-            transform_train: str = 'moderate', 
+            transform_train: str | None = 'moderate', 
             transform_validation: str | None = None,
             transform_test: str | None = None,
             ):
@@ -34,13 +38,17 @@ class Augmentation:
         self.transform_validation = transform_validation
         self.transform_test = transform_test
     
-    def augment(self, img_shape: tuple, level: str | None = None):
+    def augment(
+            self, 
+            img_shape: tuple, 
+            level: str | None = None
+            ): 
         """ 
         Creates an augmentation pipeline based on the specified level.
 
         Args:
-            img_shape (tuple): Dimensions of the image.
-            level (str, optional): Level of augmentation. One of [None, 'light', 'moderate', 'heavy'].
+            img_shape (tuple): Dimensions of the volume as (D, W, H).
+            level (str, optional): Level of augmentation, either 'light', 'moderate', 'heavy' or None. Defaults to None.
 
         Raises:
             ValueError: If `img_shape` is not 3D.

qim3d/ml/_data.py

@@ -78,6 +78,8 @@ class Dataset(torch.utils.data.Dataset):
         target = np.expand_dims(target, axis=0)
 
         if self.transform:
+
+            # Apply augmentation
             transformed = self.transform({"image": image, "label": target})
             image = transformed["image"]
             target = transformed["label"]