Use export.pkl + load_learner en Space

app.py

@@ -1,88 +1,34 @@
 import gradio as gr
-import torch
 import numpy as np
-from fastai.vision.all import PILImage, imagenet_stats, Normalize
+from fastai.vision.all import load_learner, PILImage
 from huggingface_hub import hf_hub_download
-import os
-
-# Define el repo_id y el nombre del archivo del modelo en el Hub
-MODEL_REPO_ID = "alramil/unet-segmentation-model"
-MODEL_FILE_NAME = "unet.pth"
-
-# Descarga el modelo al iniciar la app
-try:
-    model_path = hf_hub_download(repo_id=MODEL_REPO_ID, filename=MODEL_FILE_NAME)
-    print(f"Modelo descargado exitosamente a: {model_path}")
-except Exception as e:
-    print(f"Error al descargar el modelo desde el Hub: {e}")
-    print(f"Asegúrate de que el modelo '{MODEL_REPO_ID}' existe y contiene el archivo '{MODEL_FILE_NAME}'.")
-    exit()
-
-# Carga tu modelo TorchScript desde la ruta descargada
-model = torch.jit.load(model_path)
-model.eval()
-
-# Códigos y valores originales (deberían coincidir con codesAll.txt si lo usas)
-# Cargar los códigos desde el archivo codesAll.txt
-try:
-    with open("codesAll.txt", "r") as f:
-        class_names = [line.strip() for line in f if line.strip()]
-    # Reconstruir class_values basado en los códigos y el palette original (o definir un mapeo si es necesario)
-    # Nota: Este mapeo es crucial y debe coincidir con RemapMaskLabels. Aquí se asume un mapeo fijo por ahora.
-    class_values = [0, 29, 76, 150, 255] # Asegúrate de que esto coincida con RemapMaskLabels
-    if len(class_names) != len(class_values):
-         print("Advertencia: El número de clases en codesAll.txt no coincide con class_values/palette.")
-         # Puedes ajustar class_values o palette aquí si es necesario, o salir/avisar.
-         # Por ahora, seguimos con class_values y palette fijos.
-
-except FileNotFoundError:
-    print("Error: codesAll.txt no encontrado en el Space. No se pudieron cargar los nombres de las clases dinámicamente.")
-    # Usar valores y nombres por defecto si codesAll.txt no se encuentra
-    class_values = [0,29,76,150,255]
-    class_names = ['Background', 'Leaves', 'Wood', 'Pole', 'Grape'] # Valores por defecto
 
+# Descargar y cargar export.pkl
+pkl = hf_hub_download(repo_id="${MODEL_REPO}", filename="export.pkl")
+learn = load_learner(pkl)
+learn.model.eval()
 
+# Paleta de colores para índices 0–4
 palette = np.array([
-    [0,0,0],       # 0: Background (Negro)
-    [0,128,0],     # 29: Leaves (Verde)
-    [139,69,19],   # 76: Wood (Marrón Silla de Montar)
-    [128,0,128], # 150: Pole (Púrpura)
-    [128,128,128], # 255: Grape (Gris)
+    [0,0,0],      # fondo
+    [0,128,0],    # hojas
+    [139,69,19],  # madera
+    [128,0,128],  # poste
+    [128,128,128] # racimo
 ], dtype=np.uint8)
 
-# Transform para normalizar (debe ser el mismo que usaste al entrenar)
-_norm = Normalize.from_stats(*imagenet_stats)
-
 def segment_image(img):
-    pil = img.resize((640,480))
-    arr = np.array(pil).astype(np.float32) / 255.
-    tensor = torch.tensor(arr).permute(2,0,1).unsqueeze(0)
-    tensor = _norm(tensor)
-
-    with torch.no_grad():
-        out = model(tensor)
-
-    seg = out.argmax(dim=1)[0].cpu().numpy().astype(np.uint8)
-
-    # Re-mapear los valores predichos a los valores originales para usar la paleta
-    # Esto es necesario porque RemapMaskLabels mapeó 0,29,76,150,255 a 0,1,2,3,4 para el entrenamiento
-    # Pero la paleta espera los valores originales o necesitas una paleta mapeada a 0,1,2,3,4
-    # Si tu modelo predice 0,1,2,3,4, puedes mapearlos de nuevo a los valores originales
-    # O usar una paleta que mapee 0,1,2,3,4 a los colores.
-    # La forma más sencilla es usar la paleta con los índices predichos (0,1,2,3,4)
-    # Asegúrate de que la paleta está indexada correctamente.
-    # Si el modelo predice 0,1,2,3,4 y la paleta espera índices 0,1,2,3,4, entonces:
-    seg_color = palette[seg]
-
-
-    return seg_color
+    # FastAI internamente hace resize, tensor y normalize
+    pred,_,_ = learn.predict(PILImage.create(img))
+    seg = np.array(pred)  # H×W array de 0–4
+    return palette[seg]
 
 demo = gr.Interface(
     fn=segment_image,
-    inputs=gr.Image(type="pil", label="Sube una imagen"),
-    outputs=gr.Image(type="numpy", label="Máscara Segmentada"),
-    title="Segmentación Semántica de Racimos",
-    description="Sube una imagen de viñedo y obtén la máscara segmentada por clases (Fondo, Hojas, Madera, Poste, Racimo)."
+    inputs=gr.Image(type="pil", label="Sube imagen"),
+    outputs=gr.Image(type="numpy", label="Máscara segmentada"),
+    title="Segmentación de Racimos de Uva",
+    description="Modelo U-Net entrenado con FastAI"
 )
 
 if __name__=="__main__":

requirements.txt

@@ -1,7 +1,4 @@
 fastai
-torch
-torchvision
 gradio
-numpy
-pillow
 huggingface-hub
+numpy