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()