Update app.py

app.py

@@ -7,10 +7,24 @@ import gradio as gr
 from torchvision import transforms as T
 from typing import Tuple
 
+
+def rand_augment_transform(magnitude=5, num_layers=3):
+    # These are tuned for magnitude=5, which means that effective magnitudes are half of these values.
+    hparams = {
+        'rotate_deg': 30,
+        'shear_x_pct': 0.9,
+        'shear_y_pct': 0.2,
+        'translate_x_pct': 0.10,
+        'translate_y_pct': 0.30
+    }
+    ra_ops = auto_augment.rand_augment_ops(magnitude, hparams, transforms=_RAND_TRANSFORMS)
+    # Supply weights to disable replacement in random selection (i.e. avoid applying the same op twice)
+    choice_weights = [1. / len(ra_ops) for _ in range(len(ra_ops))]
+    return auto_augment.RandAugment(ra_ops, num_layers, choice_weights)
+
 def get_transform(img_size: Tuple[int], augment: bool = False, rotation: int = 0):
     transforms = []
     if augment:
-        from .augment import rand_augment_transform
         transforms.append(rand_augment_transform())
     if rotation:
         transforms.append(lambda img: img.rotate(rotation, expand=True))