Update app.py

app.py

@@ -4,20 +4,18 @@ import numpy as np
 from huggingface_hub import hf_hub_download
 import albumentations as A 
 
-# --- 1. Definición de Clases ---
+# --- 1. Definición de Clases (Necesarias para el unpickle) ---
 
 # Transformación de Etiquetas
 class TargetMaskConvertTransform(ItemTransform):
     def __init__(self): pass
     def encodes(self, x): return x
 
-# Transformación de Albumentations (SOLUCIÓN FIX)
+# Transformación de Albumentations
+# La definimos solo para que load_learner no falle al leer el archivo.
 class SegmentationAlbumentationsTransform(ItemTransform):
-    split_idx = 0 # Esta línea evita que se ejecute en inferencia
-    def __init__(self, aug): 
-        pass
-    def encodes(self, x):
-        return x
+    def __init__(self, aug): pass
+    def encodes(self, x): return x
 
 # Métrica Jaccard
 def JaccardMulti(inp, targ, axis=1, smooth=1e-6, reduction='mean', square_sum=False):
@@ -34,9 +32,15 @@ FILENAME = "model.pkl"
 print(f"Descargando modelo desde {REPO_ID}...")
 model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
 
-# Cargamos el learner (CPU=True)
+# Cargamos el learner
 learn = load_learner(model_path, cpu=True)
 
+# --- SOLUCIÓN DEL ERROR (CIRUGÍA DE PIPELINE) ---
+# Buscamos la transformación de Albumentations dentro del learner y la ELIMINAMOS de la lista.
+# Así garantizamos que nunca se ejecute durante la inferencia.
+learn.dls.after_item.fs = [t for t in learn.dls.after_item.fs if not isinstance(t, SegmentationAlbumentationsTransform)]
+print("Transformación de Albumentations eliminada del pipeline de inferencia.")
+
 # --- 3. Función de Predicción ---
 def predict(img):
     # Predicción
@@ -49,13 +53,12 @@ def predict(img):
     # Creamos la imagen RGB vacía
     color_mask = np.zeros((h, w, 3), dtype=np.uint8)
     
-    # --- PALETA DE COLORES ACTUALIZADA ---
-    # 0:Fondo(Negro), 1:Hojas(Verde), 2:Madera(Azul), 3:Poste(Rojo), 4:Uvas(Blanco)
-    
+    # --- PALETA DE COLORES ---
+    # 0:Fondo, 1:Hojas(Verde), 2:Madera(Azul), 3:Poste(Rojo), 4:Uvas(Blanco)
     color_mask[mask_data == 1] = [0, 255, 0]     # Hojas -> Verde
-    color_mask[mask_data == 2] = [0, 0, 255]     # Madera -> Azul (Tu cambio)
-    color_mask[mask_data == 3] = [255, 0, 0]     # Poste -> Rojo (Tu cambio)
-    color_mask[mask_data == 4] = [255, 255, 255] # Uvas -> Blanco (Tu cambio)
+    color_mask[mask_data == 2] = [0, 0, 255]     # Madera -> Azul
+    color_mask[mask_data == 3] = [255, 0, 0]     # Poste -> Rojo
+    color_mask[mask_data == 4] = [255, 255, 255] # Uvas -> Blanco
     
     return color_mask