Update app.py

app.py

@@ -35,12 +35,32 @@ import torch.nn.functional as F
 #learner = from_pretrained_fastai(repo_id)
 
 from huggingface_hub import from_pretrained_fastai
-
-
-
 import torchvision.transforms as transforms
+from transforms import ItemTransform
+
+from albumentations import (
+    Compose,
+    OneOf,
+    ElasticTransform,
+    GridDistortion,
+    OpticalDistortion,
+    HorizontalFlip,
+    Rotate,
+    Transpose,
+    CLAHE,
+    ShiftScaleRotate
+)
+
+class SegmentationAlbumentationsTransform(ItemTransform):
+    split_idx = 0
+
+    def __init__(self, aug):
+        self.aug = aug
 
-
+    def encodes(self, x):
+        img,mask = x
+        aug = self.aug(image=np.array(img), mask=np.array(mask))
+        return PILImage.create(aug["image"]), PILMask.create(aug["mask"])
 
 class TargetMaskConvertTransform(ItemTransform):
     def __init__(self):