[WIP] Add RandAugment

tests/test_rand_aug.py

@@ -21,37 +21,92 @@
 import math
 
 class RandAugment(Operation):
-    def __init__(self, size: int):
+    def __init__(self, 
+                 size: int = 32, 
+                 num_ops: int = 2, 
+                 magnitude: int = 9, 
+                 num_magnitude_bins: int = 31):
         super().__init__()
         self.size = size
+        self.num_ops = num_ops
+        self.magnitude = magnitude
+        num_bins = num_magnitude_bins
+        # index, name (for readability); bins, sign multiplier
+        # those with a -1 can have negative magnitude with probability 0.5
+        self.op_table = [
+            (0, "Identity", np.array(0.0), 1),
+            (1, "ShearX", np.linspace(0.0, 0.3, num_bins), -1),
+            (2, "ShearY", np.linspace(0.0, 0.3, num_bins), -1),
+            (3, "TranslateX", np.linspace(0.0, 150.0 / 331.0 * size, num_bins), -1),
+            (4, "TranslateY", np.linspace(0.0, 150.0 / 331.0 * size, num_bins), -1),
+            (5, "Rotate", np.linspace(0.0, 30.0, num_bins), -1),
+            (6, "Brightness", np.linspace(0.0, 0.9, num_bins), -1),
+            (7, "Color", np.linspace(0.0, 0.9, num_bins), -1),
+            (8, "Contrast", np.linspace(0.0, 0.9, num_bins), -1),
+            (9, "Sharpness", np.linspace(0.0, 0.9, num_bins), -1),
+            (10, "Posterize", 8 - (np.arange(num_bins) / ((num_bins - 1) / 4)).round().astype('uint8'), 1),
+            (11, "Solarize", np.linspace(255.0, 0.0, num_bins), 1),
+            (12, "AutoContrast", np.array(0.0), 1),
+            (13, "Equalize", np.array(0.0), 1),
+        ]
 
     def generate_code(self) -> Callable:
         my_range = Compiler.get_iterator()
+        op_table = self.op_table
+        magnitudes = np.array([(op[2][self.magnitude] if op[2].ndim > 0 else 0) for op in self.op_table])
+        is_signed = np.array([op[3] for op in self.op_table])
+        num_ops = self.num_ops
+#         for i in range(len(magnitudes)):
+#             print(i, op_table[i][1], '%.3f'%magnitudes[i])
         def randaug(im, mem):
             dst, scratch, lut, scratchf = mem
             for i in my_range(im.shape[0]):
-
-                ## TODO actual randaug logic
-
-                ## rotate
-                deg = np.random.random() * 45.0
-                rotate(im[i], dst[i], deg)
-
-                ## brighten
-                blend(im[i], scratch[i][0], 0.5, dst[i])
-
-                ## adjust contrast
-                adjust_contrast(im[i], scratch[i][0], 0.5, dst[i])
-
-                if deg < 10:
-                    ## equalize
-                    equalize(im[i], lut[i], dst[i])
-
-                if 10 < deg < 20:
-                    ## autocontrast -- things are getting slower now.
-                    autocontrast(im[i], scratchf[i][0], dst[i])
-                    # --^ this is a good candidate for moving entirely to OpenCV
-                    # it would involve less casting/scratch memory I think
+                for _ in range(num_ops):
+                    idx = np.random.randint(low=0, high=13+1)
+                    mag = magnitudes[idx]
+                    if np.random.random() < 0.5:
+                        mag = mag * is_signed[idx] 
+
+                    # Not worth fighting numba at the moment.
+                    # TODO
+                    if idx == 1: # ShearX (0.004)
+                        shear(im[i], dst[i], mag, 0)
+
+                    if idx == 2: # ShearY
+                        shear(im[i], dst[i], 0, mag)
+
+                    if idx == 3: # TranslateX
+                        translate(im[i], dst[i], int(mag), 0)
+
+                    if idx == 4: # TranslateY
+                        translate(im[i], dst[i], 0, int(mag))
+
+                    if idx == 5: # Rotate
+                        rotate(im[i], dst[i], mag)
+
+                    if idx == 6: # Brightness
+                        blend(im[i], scratch[i][0], 1.0 + mag, dst[i])
+
+                    if idx == 7: # Color
+                        adjust_saturation(im[i], scratch[i][0], 1.0 + mag, dst[i])
+
+                    if idx == 8: # Contrast
+                        adjust_contrast(im[i], scratch[i][0], 1.0 + mag, dst[i])
+
+                    if idx == 9: # Sharpness
+                        sharpen(im[i], dst[i], 1.0 + mag)
+
+                    if idx == 10: # Posterize
+                        posterize(im[i], int(mag), dst[i])
+
+                    if idx == 11: # Solarize
+                        solarize(im[i], mag, dst[i])
+
+                    if idx == 12: # AutoContrast (TODO: takes 0.04s -> 0.052s) (+0.01s)
+                        autocontrast(im[i], scratchf[i][0], dst[i])
+
+                    if idx == 13: # Equalize (TODO: +0.008s)
+                        equalize(im[i], lut[i], dst[i])
 
             return dst
 
@@ -346,6 +401,7 @@ def test_translate(amt):
                         num_workers=2, order=OrderOption.RANDOM,
                         drop_last=True, pipelines={'image': pipeline})
 
+        import matplotlib.pyplot as plt
         for ims, labs in loader: pass
         start_time = time.time()
         for _ in range(5): #(100):