Update app.py

app.py

@@ -1,9 +1,27 @@
 from huggingface_hub import from_pretrained_fastai
 import gradio as gr
 from fastai.vision.all import *
-
 import PIL
 import torchvision.transforms as transforms
+
+class TargetMaskConvertTransform(ItemTransform):
+    def __init__(self): 
+        pass
+    def encodes(self, x):
+        img,mask = x
+        
+        #Convert to array
+        mask = np.array(mask)
+        
+        mask[mask!=255]=0
+        # Change 255 for 1
+        mask[mask==255]=1
+        
+        
+        # Back to PILMask
+        mask = PILMask.create(mask)
+        return img, mask
+
 def transform_image(image):
     my_transforms = transforms.Compose([transforms.ToTensor(),
                                         transforms.Normalize(