Update transforms.py

transforms.py

@@ -1,12 +1,13 @@
 from transformers import AutoImageProcessor
 import torchvision.transforms as T
 
-processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k")
-augment = T.Compose([
-    T.RandomResizedCrop(224),
-    T.RandomHorizontalFlip(),
-    T.ColorJitter(0.2,0.2,0.2,0.1),
-    T.RandomRotation(20),
-    T.ToTensor(),
-    T.Normalize(mean=processor.image_mean, std=processor.image_std)
-])
+def build_transforms(backbone_model="google/vit-base-patch16-224-in21k"):
+    processor = AutoImageProcessor.from_pretrained(backbone_model)
+    return T.Compose([
+        T.RandomResizedCrop(224),
+        T.RandomHorizontalFlip(),
+        T.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
+        T.RandomRotation(20),
+        T.ToTensor(),
+        T.Normalize(mean=processor.image_mean, std=processor.image_std)
+    ])