UNet Data Augmentation Tool

• Implemented an Augmentation class (under qim3d/utils/augmentations.py) where the user can either choose default augmentations: None, 'light', 'moderate' or 'heavy', or add a customized augmentation.
• Customization can be changed to get **kwargs as input, otherwise should be removed (user can directly put custom transformation in qim3d.utils.Dataset() instead).
• Added new version of augmentations in UNet.ipynb.

qim3d/utils/augmentations.py

@@ -22,11 +22,11 @@ class Augmentation:
         my_augmentation = Augmentation()
         moderate_augment = augmentation.augment(level='moderate')
     """
-    def __init__(self):
+    def __init__(self, resize=256, mean=0.5, std=0.5):
         
-        self.resize = 256
-        self.mean = 0.5
-        self.std = 0.5
+        self.resize = resize
+        self.mean = mean
+        self.std = std
                 
     def augment(self, level=None):
 
@@ -36,7 +36,7 @@ class Augmentation:
             # Check if the custom transformation is an albumentation:
             if not isinstance(level, A.core.composition.Compose):
                 raise ValueError("Custom Transformations need to be an instance of Albumentations Compose class, "
-                                 "or one of the following levels: None, light, moderate, heavy")
+                                 "or one of the following levels: None, 'light', 'moderate', 'heavy'")
             # Custom transformation
             else:
                 return level