Update app.py

app.py

@@ -2,8 +2,24 @@ 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
+        mask = np.array(mask)
+        new_mask = np.zeros_like(mask, dtype=np.uint8)
 
 
+        new_mask[mask==150]=1                   #Clase Leaves
+        new_mask[(mask==29) | (mask==25)]=2     #Clase Wood
+        new_mask[(mask==74) | (mask==76)]=3     #Clase Pole
+        new_mask[mask==255]=4                   #Clase Wood
+
+
+        mask = PILMask.create(new_mask)
+        return img, mask
+
 
 repo_id = "joortif/Practica3"