Update app.py

app.py

@@ -8,6 +8,29 @@ import torch
 
 # Cargar modelo desde Hugging Face Hub
 model_path = hf_hub_download(repo_id="PablitoGil14/AP-Practica3", filename="model.pkl")
+
+class TargetMaskConvertTransform(ItemTransform):
+    def __init__(self):
+        pass
+    def encodes(self, x):
+        img,mask = x
+
+        #Convert to array
+        mask = np.array(mask)
+
+        mask[(mask!=255) & (mask!=150) & (mask!=76) & (mask!=74) & (mask!=29) & (mask!=25)]=0
+        mask[mask==255]=1
+        mask[mask==150]=2
+        mask[mask==76]=4
+        mask[mask==74]=4
+        mask[mask==29]=3
+        mask[mask==25]=3
+
+        # Back to PILMask
+        mask = PILMask.create(mask)
+        return img, mask
+
+
 learn = load_learner(model_path)
 
 def segmentar(img: Image.Image):