Create app.py

app.py

@@ -0,0 +1,109 @@
+import gradio as gr
+import torch
+import torch.nn as nn
+from torchvision import models, transforms
+from PIL import Image
+
+# ===============================================
+# 1. CONFIGURACIÓN Y CARGA DEL MODELO
+# ===============================================
+
+# Definición de las 8 clases de salida
+# NOTA: El modelo DENSENET201 DEBE haber sido entrenado con estas 8 clases.
+CLASS_NAMES = [
+    "Normal",
+    "Infarto Agudo",
+    "Infarto Antiguo",
+    "Fibrilación Auricular",
+    "Bloqueo de Rama Izquierda",
+    "Bloqueo de Rama Derecha",
+    "Extrasístole Auricular",
+    "Extrasístole Ventricular"
+]
+
+# Configuración del dispositivo (GPU si está disponible, sino CPU)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def load_model(model_path: str = "densenet_model.pth"):
+    """Carga los pesos del modelo DenseNet201 entrenado."""
+    
+    # 1. Definir la arquitectura base (DenseNet201)
+    # Se usa weights=None porque los pesos serán cargados de nuestro .pth
+    model = models.densenet201(weights=None) 
+    
+    # 2. Reemplazar la capa final para que coincida con el número de clases (8)
+    num_ftrs = model.classifier.in_features
+    model.classifier = nn.Linear(num_ftrs, len(CLASS_NAMES)) 
+    
+    # 3. Cargar los pesos entrenados
+    try:
+        model.load_state_dict(torch.load(model_path, map_location=device))
+        model.to(device)
+        model.eval() # Poner el modelo en modo de evaluación
+        print(f"Modelo cargado con éxito desde {model_path} en {device}.")
+        return model
+    except FileNotFoundError:
+        print(f"Error: El archivo del modelo '{model_path}' no se encontró. ¡La aplicación no funcionará correctamente sin él!")
+        return None 
+
+# Cargar el modelo globalmente (¡solo una vez!)
+ecg_model = load_model()
+
+# ===============================================
+# 2. FUNCIÓN DE PREDICCIÓN
+# ===============================================
+
+# Definición de las transformaciones necesarias para el pre-procesamiento
+preprocess = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    # Valores de normalización estándar para modelos pre-entrenados
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+def predict(image: Image.Image):
+    """Realiza la inferencia en la imagen del ECG."""
+    if ecg_model is None:
+        # Devuelve un mensaje de error legible en la UI de Gradio
+        return {"ERROR: Modelo no cargado (Falta .pth)": 1.0} 
+
+    # 1. Preprocesar la imagen
+    img_tensor = preprocess(image).unsqueeze(0).to(device)
+
+    # 2. Inferir sin calcular gradientes (más rápido)
+    with torch.no_grad():
+        output = ecg_model(img_tensor)
+
+    # 3. Post-procesar (Convertir a probabilidades)
+    probabilities = torch.nn.functional.softmax(output[0], dim=0)
+    
+    # Crear el diccionario de resultados para Gradio
+    results = {
+        CLASS_NAMES[i]: float(probabilities[i])
+        for i in range(len(CLASS_NAMES))
+    }
+    return results
+
+# ===============================================
+# 3. INTERFAZ GRADIO
+# ===============================================
+
+# Asignamos la interfaz a la variable 'demo'
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil", label="Sube una imagen de ECG"),
+    # Muestra las 5 clases con mayor probabilidad
+    outputs=gr.Label(num_top_classes=5, label="Clasificación del Modelo"), 
+    title="Análisis de ECG con IA: 8 Clases de Diagnóstico",
+    description="Sube una imagen de tu electrocardiograma. El modelo clasifica en 8 condiciones cardíacas: Normal, Infarto Agudo, Infarto Antiguo, Fibrilación Auricular, Bloqueo de Rama Izquierda, Bloqueo de Rama Derecha, Extrasístole Auricular y Extrasístole Ventricular.",
+    allow_flagging="auto"
+)
+
+# ===============================================
+# 4. LANZAMIENTO
+# ===============================================
+
+if __name__ == "__main__":
+    # Usa demo.launch() para iniciar la aplicación web
+    # share=True proporciona un enlace público temporal
+    demo.launch(share=True)