added augmentation

Author: AndreaCominelli

config/conf.json

@@ -5,5 +5,11 @@
     "optimizer":"adam",
     "metrics":["mIoU", "oneHotIoU", "mse", "oneHotMeanIoU"],
     "loss":"focal",
-    "monitor_metric":"val_one_hot_mean_io_u"
+    "monitor_metric":"val_one_hot_mean_io_u",
+    "augmentation":{
+        "toGray":1,
+        "horizontalFlip":1,
+        "verticalFlip":1,
+        "clahe":1
+    }
 }

src/utils/augmentation_pipeline.py

@@ -0,0 +1,11 @@
+import albumentations as A
+
+def get_transformation_pipeline(augmentation=[]):
+    augmentation_pipeline = []
+    for t, p in augmentation:
+        transformation = A.OneOf([
+            t
+        ],p = p)
+        augmentation_pipeline.append(transformation)
+    
+    return A.Compose(augmentation_pipeline)

src/utils/conf_reader.py

@@ -1,6 +1,8 @@
 import json
 import tensorflow as tf
-from utils.losses import LOSSES
+import albumentations as A
+from augmentation_pipeline import get_transformation_pipeline
+from losses import LOSSES
 
 OPTIMIZERS = {
     'adam':tf.keras.optimizers.Adam,
@@ -14,6 +16,13 @@
     'mse':tf.keras.metrics.MeanSquaredError()
 }
 
+AUGMENTATION = {
+    "toGray":A.ToGray(),
+    "horizontalFlip":A.HorizontalFlip(),
+    "verticalFlip":A.VerticalFlip(),
+    "clahe":A.CLAHE()
+}
+
 def read_conf(filename):
     with open(filename,'r') as f:
         conf = json.load(f)
@@ -49,4 +58,10 @@ def read_conf(filename):
     else:
         conf['loss'] = LOSSES[conf['loss']]()
 
+    if not 'augmentation' in conf:
+        conf['augmentation'] = A.Compose([])
+    else:
+        conf['augmentation'] = [(AUGMENTATION[x], p) for x, p in conf['augmentation'].items()]
+        conf['augmentation'] = get_transformation_pipeline(conf['augmentation'])
+
     return conf