Update app.py

app.py

@@ -3,10 +3,27 @@ import gradio as gr
 import torchvision.transforms as transforms
 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(
@@ -26,6 +43,7 @@ def mask_to_img(mask):
 
 # Definimos una función que se encarga de llevar a cabo las predicciones
 def predict(img):
+  learn = load_learner('export.pkl')
   device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 
   model = learn.cpu()
   model.eval()