Update app.py

app.py

@@ -2,6 +2,33 @@ from huggingface_hub import from_pretrained_fastai
 import gradio as gr
 from fastai.vision.all import *
 
+class TargetMaskConvertTransform(ItemTransform):
+    def __init__(self):
+        pass
+    def encodes(self, x):
+        img,mask = x
+
+        #Convert to array
+        mask = np.array(mask)
+
+        # Uvas
+        mask[mask==255]=1
+
+        # Hojas
+        mask[mask==150]=2
+
+        # Poste
+        mask[mask==76]=3
+        mask[mask==74]=3
+
+        # Madera
+        mask[mask==29]=4
+        mask[mask==25]=4
+
+        # Back to PILMask
+        mask = PILMask.create(mask)
+        return img, mask
+
 repo_id = "ancebuc/grapes-segmentation"
 
 learner = from_pretrained_fastai(repo_id)