augmentation constructor takes str or np.random.Generator

ctc_otfa.py

@@ -5,6 +5,7 @@
 import ctc_utils
 import os 
 import json
+from typing import Union
 from abc import ABC, abstractmethod
 import ctc_otfa
 import tensorflow as tf
@@ -103,18 +104,31 @@ def apply_augmentations(image, augmentations):
 
 
 class augmentation(ABC):
-    def __init__(self, type, distribution, variance):
+    def __init__(self, type:str, distribution: Union[str, np.random.Generator], variance:float):
         self.type = str(type)
-
-        if isinstance(distribution, np.random.Generator):
-            self.distribution = distribution
-        else:
-            raise Exception('Invalid distribution type \"' + str(distribution) + '\"')
         self.variance = float(variance)
 
+        if isinstance(distribution, str):
+            self.distribution = self.str_to_distribution(distribution=distribution) #convert to np.random.Generator
+        else:
+            self.distribution = distribution
+
     def __str__(self) -> str:
         return 'AUGMENTATION: ' +self.type + ', ' + str(self.variance) + ', ' + str(self.distribution)
 
+    def str_to_distribution(self, distribution: str) -> np.random.Generator:
+        distribution.lower()
+        if distribution == "gaussian" or distribution == "normal":
+            return np.random.normal
+        elif distribution == "uniform":
+            return np.random.uniform
+        elif distribution == "random":
+            return np.random.random
+        elif distribution == "binomial":
+            return np.random.binomial
+        else:
+            raise Exception('Invalid distribution type \"' + str(distribution) + '\"')
+
     # This function should take a tensor and return a tensor
     @abstractmethod
     def augment(self):