Update app.py

app.py

@@ -4,6 +4,31 @@ 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)
+        
+        # Aquí definimos cada clase en la máscara
+        # uva:
+        mask[mask==255]=1
+        # hojas:
+        mask[mask==150]=2
+        # conductores:
+        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 = "paascorb/practica3_Segmentation"
 
 learner = from_pretrained_fastai(repo_id)