Update app.py

app.py

@@ -11,32 +11,19 @@ from albumentations import (
     Rotate,
     GridDistortion
 )
+from pathlib import Path
 
-class SegmentationAlbumentationsTransform(ItemTransform):
-    split_idx = 0
-
-    def __init__(self, aug):
-        self.aug = aug
+# --- FUNCIONES Y CLASES NECESARIAS PARA EL PICKLE ---
 
-    def encodes(self, x):
-        img,mask = x
-        aug = self.aug(image=np.array(img), mask=np.array(mask))
-        return PILImage.create(aug["image"]), PILMask.create(aug["mask"])
-
-
-
-# Cargar modelo desde Hugging Face Hub
-model_path = hf_hub_download(repo_id="PablitoGil14/AP-Practica3", filename="model.pkl")
+def get_y_fn(x):
+    return Path(str(x).replace("Images","Labels").replace("color","gt").replace(".jpg",".png"))
 
 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
@@ -44,19 +31,23 @@ class TargetMaskConvertTransform(ItemTransform):
         mask[mask==74]=4
         mask[mask==29]=3
         mask[mask==25]=3
+        return img, PILMask.create(mask)
 
-        # Back to PILMask
-        mask = PILMask.create(mask)
-        return img, mask
+class SegmentationAlbumentationsTransform(ItemTransform):
+    split_idx = 0
+    def __init__(self, aug):
+        self.aug = aug
+    def encodes(self, x):
+        img,mask = x
+        aug = self.aug(image=np.array(img), mask=np.array(mask))
+        return PILImage.create(aug["image"]), PILMask.create(aug["mask"])
 
+# --- CARGAR MODELO ---
 
+model_path = hf_hub_download(repo_id="PablitoGil14/AP-Practica3", filename="model.pkl")
 learn = load_learner(model_path)
 
-
-def get_y_fn(x):
-    return Path(str(x).replace("Images","Labels").replace("color","gt").replace(".jpg",".png"))
-
-
+# --- FUNCIÓN DE PREDICCIÓN ---
 
 def segmentar(img: Image.Image):
     img = img.resize((640, 480))
@@ -68,8 +59,7 @@ def segmentar(img: Image.Image):
     with torch.no_grad():
         preds = learn.model.eval()(x)
         mask = torch.argmax(preds, dim=1).squeeze().cpu().numpy()
-        
-    # Asignar colores según los valores de clase
+
     out_mask = np.zeros_like(mask, dtype=np.uint8)
     out_mask[mask == 1] = 255
     out_mask[mask == 2] = 150
@@ -77,7 +67,8 @@ def segmentar(img: Image.Image):
     out_mask[mask == 4] = 74
     return Image.fromarray(out_mask)
 
-# Interfaz de Gradio
+# --- INTERFAZ GRADIO ---
+
 demo = gr.Interface(
     fn=segmentar,
     inputs=gr.Image(type="pil"),