app.py

import gradio as gr
import torch
import torch.nn as nn
import numpy as np
from PIL import Image
import matplotlib.cm as cm
from transformers import AutoModelForImageClassification, AutoConfig
import torchvision.transforms as T

# ── Configuración ──────────────────────────────────────────────────────────────
CLASES    = ["Glioma", "Meningioma", "No Tumor", "Pituitario"]
N_CLASES  = 4
IMG_SIZE  = 224
MEDIA     = [0.485, 0.456, 0.406]
DESV      = [0.229, 0.224, 0.225]

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Repositorios del usuario en Hugging Face
HF_REPO_VIT      = "webiacademic/brain-tumor-vit"
HF_REPO_CONVNEXT = "webiacademic/brain-tumor-convnext"

# ── Transformación de entrada ──────────────────────────────────────────────────
transform = T.Compose([
    T.Resize((IMG_SIZE, IMG_SIZE)),
    T.ToTensor(),
    T.Normalize(mean=MEDIA, std=DESV),
])

# ── Carga de modelos desde Hugging Face ─────────────────────────────────────────
print("Cargando modelos desde el Hub...")

# Cargamos ViT directamente desde tu repositorio
# Cargamos ViT directamente desde tu repositorio
modelo_vit = AutoModelForImageClassification.from_pretrained(
    HF_REPO_VIT,
    attn_implementation="eager"  # <-- ¡ESTA ES LA CLAVE!
)
modelo_vit = modelo_vit.to(DEVICE).eval()
# Cargamos ConvNeXt directamente desde tu repositorio
modelo_convnext = AutoModelForImageClassification.from_pretrained(HF_REPO_CONVNEXT)
modelo_convnext = modelo_convnext.to(DEVICE).eval()

print("Modelos cargados con éxito.")

# ── Grad-CAM ───────────────────────────────────────────────────────────────────
class GradCAM:
    def __init__(self, modelo, capa):
        self.activaciones = None
        self.gradientes   = None
        self._fwd = capa.register_forward_hook(
            lambda m, i, o: setattr(self, "activaciones", o.detach())
        )
        self._bwd = capa.register_full_backward_hook(
            lambda m, gi, go: setattr(self, "gradientes", go[0].detach())
        )

    def calcular(self, tensor, clase_idx):
        pesos = self.gradientes.mean(dim=[0, 2, 3])
        mapa  = sum(p * a for p, a in zip(pesos, self.activaciones[0]))
        mapa  = torch.relu(mapa).cpu().numpy()
        if mapa.max() > 0:
            mapa = (mapa - mapa.min()) / (mapa.max() - mapa.min() + 1e-8)
        return mapa

    def liberar(self):
        self._fwd.remove()
        self._bwd.remove()


def attention_rollout(modelo, imagen_tensor):
    """
    Calcula el mapa de atención por Attention Rollout de forma segura
    y limpia para evitar mapas vacíos o grises.
    """
    try:
        modelo.eval()
        # Forzar al modelo a devolver las matrices de atención
        with torch.no_grad():
            outputs = modelo(
                imagen_tensor.unsqueeze(0).to(DEVICE),
                output_attentions=True
            )
        
        atenciones = outputs.attentions
        if atenciones is None:
            print("El modelo no devolvió las atenciones.")
            return None
        
        # Clonamos y pasamos a CPU inmediatamente para asegurar estabilidad numérica
        n_tokens = atenciones[0].shape[-1]
        rollout = torch.eye(n_tokens, dtype=torch.float32)

        # Replicamos el algoritmo exacto de tu entrenamiento en CPU
        for att in atenciones:
            # Quitamos gradientes, enviamos a cpu y promediamos las cabezas (dim=1)
            att_media = att.detach().cpu().mean(dim=1)[0]
            # Añadir conexiones residuales
            att_media = att_media + torch.eye(n_tokens)
            # Renormalizar filas
            att_media = att_media / att_media.sum(dim=-1, keepdim=True)
            # Multiplicación acumulativa (Rollout)
            rollout = torch.matmul(att_media, rollout)

        # Extraer atención del token CLS (0) hacia los parches (1 en adelante)
        mapa = rollout[0, 1:]
        
        # Cambiar el tamaño lineal a la cuadrícula de 14x14
        mapa = mapa.reshape(14, 14).numpy()
        
        # Evitar división por cero si el mapa es uniforme
        rango = mapa.max() - mapa.min()
        if rango < 1e-8:
            return np.zeros((14, 14))
            
        # Normalizar estrictamente entre 0 y 1
        mapa = (mapa - mapa.min()) / rango
        return mapa

    except Exception as e:
        print(f"Error crítico en Attention Rollout: {e}")
        return None

def superponer_mapa(img_np, mapa_raw, size):
    mapa_pil    = Image.fromarray((mapa_raw * 255).astype(np.uint8)).resize(
        (size, size), Image.BILINEAR
    )
    mapa_rgb    = cm.jet(np.array(mapa_pil) / 255.0)[:, :, :3]
    overlay     = (0.55 * img_np + 0.45 * mapa_rgb).clip(0, 1)
    return (overlay * 255).astype(np.uint8)


# ── Función principal de inferencia ───────────────────────────────────────────
def clasificar(imagen_pil: Image.Image):
    if imagen_pil is None:
        return {}, {}, None, None

    imagen_pil = imagen_pil.convert("RGB")
    tensor     = transform(imagen_pil)
    img_np     = np.array(imagen_pil.resize((IMG_SIZE, IMG_SIZE))) / 255.0

    # ── ViT: inferencia + Attention Rollout ───────────────────────────────────
    with torch.no_grad():
        logits_vit = modelo_vit(tensor.unsqueeze(0).to(DEVICE)).logits
    probs_vit  = torch.softmax(logits_vit, dim=1)[0].cpu().numpy()

    mapa_vit   = attention_rollout(modelo_vit, tensor)
    img_vit    = superponer_mapa(img_np, mapa_vit, IMG_SIZE) if mapa_vit is not None else (img_np * 255).astype(np.uint8)

    # ── ConvNeXt: inferencia + Grad-CAM ───────────────────────────────────────
    capa_obj   = modelo_convnext.convnext.encoder.stages[-1].layers[-1]
    gcam       = GradCAM(modelo_convnext, capa_obj)

    tensor_req = tensor.unsqueeze(0).to(DEVICE).requires_grad_(True)
    logits_cnx = modelo_convnext(tensor_req).logits
    probs_cnx  = torch.softmax(logits_cnx, dim=1)[0].detach().cpu().numpy()
    pred_cnx   = int(probs_cnx.argmax())

    modelo_convnext.zero_grad()
    logits_cnx[0, pred_cnx].backward()
    mapa_cnx   = gcam.calcular(tensor_req, pred_cnx)
    gcam.liberar()

    img_cnx    = superponer_mapa(img_np, mapa_cnx, IMG_SIZE)

    # ── Formatear salidas para Gradio ─────────────────────────────────────────
    out_vit = {CLASES[i]: float(probs_vit[i]) for i in range(N_CLASES)}
    out_cnx = {CLASES[i]: float(probs_cnx[i]) for i in range(N_CLASES)}

    return out_vit, out_cnx, Image.fromarray(img_vit), Image.fromarray(img_cnx)


# ── Interfaz Gradio ────────────────────────────────────────────────────────────
CSS = """
#titulo { text-align: center; font-size: 1.3em; font-weight: 700; color: #1a3a5c; margin-bottom: 4px; }
#subtitulo { text-align: center; color: #555; margin-bottom: 16px; font-size: 0.95em; }
.label-vit  { color: #4f81bd !important; font-weight: 600; }
.label-cnx  { color: #1a3a5c !important; font-weight: 600; }
"""

DESCRIPCION = """
Sube una imagen de resonancia magnética (MRI) cerebral y compara en tiempo real
la predicción de **ViT** (Vision Transformer) y **ConvNeXt**, junto con sus
mapas de interpretabilidad (Attention Rollout y Grad-CAM).

**Clases:** Glioma · Meningioma · No Tumor · Tumor Pituitario
"""

EJEMPLOS_INFO = """
> **Nota:** Para mejores resultados usa imágenes MRI en escala de grises,
> preferiblemente con contraste (CE-MRI). El modelo fue entrenado con el
> [Brain Tumor MRI Dataset](https://www.kaggle.com/datasets/masoudnickparvar/brain-tumor-mri-dataset).
"""

with gr.Blocks(css=CSS, theme=gr.themes.Soft()) as demo:

    gr.Markdown("# Clasificación de Tumores Cerebrales en MRI", elem_id="titulo")
    gr.Markdown(DESCRIPCION, elem_id="subtitulo")

    with gr.Row():
        with gr.Column(scale=1):
            imagen_entrada = gr.Image(
                type="pil",
                label="Imagen MRI de entrada",
                height=280,
            )
            btn = gr.Button("🔍 Clasificar", variant="primary", size="lg")
            gr.Markdown(EJEMPLOS_INFO)

        with gr.Column(scale=2):
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### ViT — Vision Transformer", elem_classes="label-vit")
                    probs_vit  = gr.Label(num_top_classes=4, label="Probabilidades por clase")
                    mapa_vit   = gr.Image(label="Attention Rollout", height=220)

                with gr.Column():
                    gr.Markdown("### ConvNeXt", elem_classes="label-cnx")
                    probs_cnx  = gr.Label(num_top_classes=4, label="Probabilidades por clase")
                    mapa_cnx   = gr.Image(label="Grad-CAM", height=220)

    btn.click(
        fn=clasificar,
        inputs=[imagen_entrada],
        outputs=[probs_vit, probs_cnx, mapa_vit, mapa_cnx],
    )

    gr.Markdown("""
---
**Proyecto final — Deep Learning (1INF52) · PUCP · 2026-1**
Modelos: `google/vit-base-patch16-224` · `facebook/convnext-base-224` · Fine-tuning
""")

demo.launch()