discretization added to trainer class

Author: safakkbilici

bayesmedaug/augmentations/functional.py

@@ -392,7 +392,7 @@ class ShiftY():
         * x>=0: 100, x<0: -100
     """
     def __init__(self, shift_amount_y):
-        self.shift_amount_y = int(shift_amount_y)
+        self.shift_amount_y = shift_amount_y
         self.to_mask = True
         self.to_img = True
 
@@ -428,7 +428,7 @@ class ShiftX():
         * x>=0: 100, x<0: -100
     """
     def __init__(self, shift_amount_x):
-        self.shift_amount_x = int(shift_amount_x)
+        self.shift_amount_x = shift_amount_x
         self.to_mask = True
         self.to_img = True
 

bayesmedaug/callbacks/trainer.py

@@ -25,7 +25,8 @@
     roundup,
     discrete_angle,
     discrete_shift,
-    discrete_rcrop
+    discrete_rcrop,
+    discrete_shear
 )
 
 class Trainer():
@@ -107,13 +108,17 @@ def train(self, **params):
 
         if "angle" in params.keys():
             params["angle"] = discrete_angle(params["angle"])
-        if "shift_x" in params.keys():
-            params["shift_x"] = discrete_shift(params["shift_x"])
-        if "shift_y" in params.keys():
-            params["shift_y"] = discrete_shift(params["shift_y"])
-
+        if "shift_amount_x" in params.keys():
+            params["shift_amount_x"] = discrete_shift(params["shift_amount_x"])
+        if "shift_amount_y" in params.keys():
+            params["shift_amount_y"] = discrete_shift(params["shift_amount_y"])
         if "crop_height" in params.keys():
             params["crop_height"] = discrete_rcrop(params["crop_height"])
+        if "shear_amount_y" in params.keys():
+            params["shear_amount_y"] = discrete_shear(params["shear_amount_y"])
+        if "shear_amount_x" in params.keys():
+            params["shear_amount_x"] = discrete_shear(params["shear_amount_x"])
+
 
         if type(self.augmentations) == Listed:
             transform = self.augmentations(**params)

bayesmedaug/utils/discretize.py

@@ -18,3 +18,6 @@ def discrete_shift(x):
 
 def discrete_rcrop(x):
     return int(x * 1000)
+
+def discrete_shear(x):
+    return x * 10