Upload app.py

app.py

@@ -0,0 +1,107 @@
+import gradio as gr
+import torch
+import numpy as np
+from PIL import Image
+from torch import nn
+from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor
+import os
+
+# --- 1. CONFIGURACIÓN ---
+MODEL_PATH = "modelo_mejorado.pth"
+
+LABELS = [
+    "fondo", "wheat leaf rust", "wheat powdery mildew",
+    "wheat septoria blotch", "wheat stem rust", "wheat stripe rust"
+]
+
+# Paleta (R, G, B)
+PALETA_COLORES = [
+    [0, 0, 0], [220, 38, 38], [22, 163, 74], 
+    [37, 99, 235], [234, 179, 8], [219, 39, 119]
+]
+
+# En Docker CPU (Free Tier) forzamos CPU para evitar errores de memoria
+device = torch.device("cpu")
+print(f"Usando dispositivo: {device}")
+
+# --- 2. CARGAR MODELO ---
+checkpoint_name = "nvidia/segformer-b4-finetuned-ade-512-512"
+
+try:
+    model_inference = SegformerForSemanticSegmentation.from_pretrained(
+        checkpoint_name,
+        num_labels=len(LABELS),
+        id2label={i: label for i, label in enumerate(LABELS)},
+        label2id={label: i for i, label in enumerate(LABELS)},
+        ignore_mismatched_sizes=True
+    )
+    # Cargar pesos
+    if os.path.exists(MODEL_PATH):
+        state_dict = torch.load(MODEL_PATH, map_location=device)
+        model_inference.load_state_dict(state_dict)
+        print("✅ Modelo cargado correctamente.")
+    else:
+        print("⚠️ NO se encontró el archivo .pth")
+
+    model_inference.to(device)
+    model_inference.eval()
+
+    image_processor = SegformerImageProcessor.from_pretrained(checkpoint_name)
+    image_processor.do_resize = False 
+    image_processor.do_rescale = True 
+
+except Exception as e:
+    print(f"Error fatal cargando modelo: {e}")
+
+# --- 3. PREDICCIÓN ---
+def predecir_enfermedad(image):
+    if image is None: return None, "⚠️ Sube una imagen."
+    
+    original_size = image.size 
+    img_resized = image.resize((512, 512))
+    
+    inputs = image_processor(images=img_resized, return_tensors="pt")
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+    
+    with torch.no_grad():
+        outputs = model_inference(**inputs)
+        logits = outputs.logits
+        logits_upsampled = nn.functional.interpolate(logits, size=(512, 512), mode="bilinear", align_corners=False)
+        pred_mask = logits_upsampled.argmax(dim=1).squeeze().cpu().numpy()
+
+    color_mask = np.zeros((512, 512, 3), dtype=np.uint8)
+    classes_found = []
+    unique_classes = np.unique(pred_mask)
+    
+    for class_id in unique_classes:
+        if class_id == 0: continue
+        classes_found.append(LABELS[class_id])
+        color_mask[pred_mask == class_id] = PALETA_COLORES[class_id]
+
+    mask_pil = Image.fromarray(color_mask).resize(original_size, resample=Image.NEAREST)
+    final_image = Image.blend(image.convert("RGB"), mask_pil.convert("RGB"), alpha=0.45)
+    
+    if len(classes_found) > 0:
+        diagnosis = "Enfermedades:\n" + "\n".join(f"- {c}" for c in classes_found)
+    else:
+        diagnosis = "Planta Sana (Solo fondo)."
+
+    return final_image, diagnosis
+
+# --- 4. INTERFAZ ---
+# Usamos Blocks simples compatibles con versiones estables
+with gr.Blocks(title="Wheat AI") as demo:
+    gr.Markdown("# 🌾 Detector de Enfermedades en Trigo")
+    
+    with gr.Row():
+        img_in = gr.Image(type="pil", label="Imagen")
+        img_out = gr.Image(type="pil", label="Resultado")
+    
+    txt_out = gr.Textbox(label="Diagnóstico")
+    btn = gr.Button("Analizar", variant="primary")
+    
+    btn.click(predecir_enfermedad, img_in, [img_out, txt_out])
+
+# Configuración crítica para Docker en Hugging Face
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)