Update app.py

app.py

@@ -6,6 +6,27 @@ import torch
 # Cargamos el modelo
 learn = load_learner('export.pkl')
 
+# Definimos todo lo necesario para hacer inferencia
+class TargetMaskConvertTransform(ItemTransform):
+    def __init__(self): 
+        pass
+    def encodes(self, x):
+        img,mask = x
+        
+        #Convert to array
+        mask = np.array(mask)
+        
+        # background = 0, leaves = 1, pole = 74 o 76, wood = 25 o 29, grape = 255
+        mask[mask == 255] = 1 # grape
+        mask[mask == 150] = 2 # leaves
+        mask[mask == 76] = 3 ; mask[mask == 74] = 3 # pole
+        mask[mask == 29] = 4 ; mask[mask == 25] = 4 # wood
+        mask[mask >= 5] = 0 # resto: background
+        
+        # Back to PILMask
+        mask = PILMask.create(mask)
+        return img, mask
+
 def transform_image(image, device):
     my_transforms = transforms.Compose([transforms.ToTensor(),
                                         transforms.Normalize(