Update app.py

app.py

@@ -11,13 +11,9 @@ class TargetMaskConvertTransform(ItemTransform):
     def __init__(self): pass
     def encodes(self, x): return x
 
-# Transformación de Albumentations (SOLUCIÓN AQUÍ)
+# Transformación de Albumentations (SOLUCIÓN FIX)
 class SegmentationAlbumentationsTransform(ItemTransform):
-    split_idx = 0 # <--- ESTA LÍNEA ES LA CLAVE
-    # Al poner split_idx=0, le decimos a FastAI: 
-    # "Esta transformación es SOLO para el conjunto de TRAIN".
-    # Por tanto, al predecir una foto nueva, se saltará esta clase y no dará error.
-    
+    split_idx = 0 # Esta línea evita que se ejecute en inferencia
     def __init__(self, aug): 
         pass
     def encodes(self, x):
@@ -49,14 +45,17 @@ def predict(img):
     # Procesamiento a Numpy
     mask_data = np.array(pred_mask, dtype=np.uint8)
     h, w = mask_data.shape
+    
+    # Creamos la imagen RGB vacía
     color_mask = np.zeros((h, w, 3), dtype=np.uint8)
     
-    # Paleta de colores
-    # 1:Hojas(Verde), 2:Madera(Marrón), 3:Poste(Gris), 4:Uvas(Morado)
-    color_mask[mask_data == 1] = [0, 255, 0]
-    color_mask[mask_data == 2] = [139, 69, 19]
-    color_mask[mask_data == 3] = [192, 192, 192]
-    color_mask[mask_data == 4] = [128, 0, 128]
+    # --- PALETA DE COLORES ACTUALIZADA ---
+    # 0:Fondo(Negro), 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)
     
     return color_mask
 
@@ -66,7 +65,7 @@ iface = gr.Interface(
     inputs=gr.Image(type="pil"),
     outputs=gr.Image(type="numpy"),
     title="Segmentación de Uvas (ResNet50)",
-    description="Modelo capaz de detectar: Uvas (Morado), Hojas (Verde), Madera (Marrón) y Postes (Gris).",
+    description="Modelo capaz de detectar: Uvas (Blanco), Hojas (Verde), Madera (Azul) y Postes (Rojo).",
     examples=[]
 )