app.py

import os
import time
import warnings
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Optional, Tuple
import gradio as gr
import numpy as np
import timm
import torch
import torch.nn as nn
import torchvision.transforms as transforms
from huggingface_hub import hf_hub_download
from PIL import Image
from pytorch_grad_cam import EigenCAM
from pytorch_grad_cam.utils.image import show_cam_on_image

# configuración
warnings.filterwarnings("ignore", message=".*HF Hub.*")
warnings.filterwarnings("ignore", category=FutureWarning)

# constantes globales
CLASES = {0: "Buen estado", 1: "Defectuoso"}
COLORES_CLASES = {0: "#28a745", 1: "#dc3545"}
COLOR_BORDE = {0: "#FFD700"}
EXTENSIONES_VALIDAS = [".png", ".jpg", ".jpeg"]
TAMANIO_MAXIMO_MB = 20
RESOLUCION = (224, 224)
DISPOSITIVO = "cuda" if torch.cuda.is_available() else "cpu"

RUTAS_MODELOS = {
    "Modelo_FTP": "modelo_v3_hibrido/hibrido_v2/fine_tuning_parcial/vit_papas_pesos_ftp.pt",
    "Modelo_TL": "modelo_v3_hibrido/hibrido_v2/transfer_learning/vit_papas_pesos_tl.pt",
}


# transformaciones de imagen
TRANSFORMACION = transforms.Compose(
    [
        transforms.Resize(RESOLUCION),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ]
)


@dataclass
class ConfigModelo:
    repo_id: str = "Carlos012/vit_papas"
    nombre_modelo: str = ""
    num_clases: int = 2
    modelo_base: str = "vit_base_patch16_224"


class ClasificadorPapasViT(nn.Module):
    def __init__(self, num_clases: int = 2, modelo_base: str = "vit_base_patch16_224"):
        super().__init__()
        self.backbone = timm.create_model(
            modelo_base, pretrained=False, num_classes=num_clases
        )
        # inicialización Xavier para mejor convergencia
        nn.init.xavier_uniform_(self.backbone.head.weight)
        nn.init.zeros_(self.backbone.head.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.backbone(x)


class GestorModelo:
    def __init__(self, config: ConfigModelo):
        self.config = config
        self.modelo: Optional[ClasificadorPapasViT] = None
        self.cam: Optional[EigenCAM] = None

    def descargar_pesos(self) -> Optional[str]:
        try:
            nombre_modelo = self._resolver_nombre_modelo()

            ruta_modelo = hf_hub_download(
                repo_id=self.config.repo_id,
                filename=nombre_modelo,
                repo_type="model",
            )
            return ruta_modelo
        except Exception as e:
            print(f"Error al descargar modelo: {e}")
            return None

    def cargar_modelo(self) -> bool:
        if self.modelo is not None:
            return True

        ruta_pesos = self.descargar_pesos()
        if ruta_pesos is None:
            return False

        try:
            # inicializar modelo
            self.modelo = ClasificadorPapasViT(
                num_clases=self.config.num_clases, modelo_base=self.config.modelo_base
            )

            # cargar pesos
            pesos = torch.load(ruta_pesos, map_location=DISPOSITIVO, weights_only=True)

            # limpiar prefijos si existen
            if any(k.startswith("_orig_mod.") for k in pesos.keys()):
                pesos = {k.replace("_orig_mod.", ""): v for k, v in pesos.items()}

            self.modelo.load_state_dict(pesos, strict=False)
            self.modelo.to(DISPOSITIVO)
            self.modelo.eval()

            # configurar Eigen-CAM
            capas_objetivo = [self.modelo.backbone.blocks[-1].norm1]
            self.cam = EigenCAM(
                model=self.modelo,
                target_layers=capas_objetivo,
                reshape_transform=self._reshape_transform,
            )

            print(f"Modelo cargado exitosamente en {DISPOSITIVO}")
            return True

        except Exception as e:
            print(f"Error al cargar modelo: {e}")
            return False

    def _resolver_nombre_modelo(self) -> str:
        # utilizar el nombre del modelo para resolver la ruta en HF Hub
        ruta = RUTAS_MODELOS.get(self.config.nombre_modelo, self.config.nombre_modelo)
        # normalizar ruta para evitar problemas de formato
        return ruta.strip().lstrip("/\\").replace("\\", "/")

    @staticmethod
    def _reshape_transform(tensor: torch.Tensor) -> torch.Tensor:
        # eliminar token CLS
        resultado = tensor[:, 1:, :]

        # calcular tamaño de malla
        tamanio_malla = int(np.sqrt(resultado.size(1)))

        # reorganizar: [Batch, Grid, Grid, Canales]
        resultado = resultado.reshape(
            tensor.size(0), tamanio_malla, tamanio_malla, tensor.size(2)
        )

        # transponer para formato imagen: [Batch, Canales, Altura, Anchura]
        return resultado.permute(0, 3, 1, 2)


class ProcesadorImagen:
    def __init__(self, gestor_ftp: GestorModelo, gestor_tl: GestorModelo):
        self.gestor_ftp = gestor_ftp
        self.gestor_tl = gestor_tl

    def validar_imagen(self, ruta_imagen: str) -> Tuple[bool, str]:
        if ruta_imagen is None:
            return False, "No se ha subido ninguna imagen"

        # validar extensión
        extension = Path(ruta_imagen).suffix.lower()
        if extension not in EXTENSIONES_VALIDAS:
            return False, f"Formato no válido. Use: {', '.join(EXTENSIONES_VALIDAS)}"

        # validar tamaño
        tamanio_mb = Path(ruta_imagen).stat().st_size / (1024 * 1024)
        if tamanio_mb > TAMANIO_MAXIMO_MB:
            return (
                False,
                f"Imagen muy grande ({tamanio_mb:.2f}MB). Máximo: {TAMANIO_MAXIMO_MB}MB",
            )

        return True, "OK"

    # función para procesar la imagen con un modelo específico
    def procesar_con_modelo(
        self, ruta_imagen: str, gestor: GestorModelo, nombre_modelo: str
    ) -> Dict:
        # cargar modelo si es necesario
        if gestor.modelo is None:
            if not gestor.cargar_modelo():
                return {"error": f"No se pudo cargar {nombre_modelo}"}

        try:
            # cargar y preparar imagen
            imagen = Image.open(ruta_imagen).convert("RGB")
            imagen_redimensionada = imagen.resize(RESOLUCION)
            tensor_imagen = TRANSFORMACION(imagen).unsqueeze(0).to(DISPOSITIVO)

            # predicción
            tiempo_inicio = time.time()
            with torch.no_grad():
                salida = gestor.modelo(tensor_imagen)
                probabilidades = torch.softmax(salida, dim=1)[0].cpu().numpy()
            tiempo_inferencia = time.time() - tiempo_inicio

            id_prediccion = int(np.argmax(probabilidades))
            confianza = float(probabilidades[id_prediccion])

            # generar Eigen-CAM
            mapa_atencion = gestor.cam(input_tensor=tensor_imagen, targets=None)[0, :]

            # crear visualización
            imagen_normalizada = np.float32(imagen_redimensionada) / 255.0
            visualizacion_cam = show_cam_on_image(
                imagen_normalizada, mapa_atencion, use_rgb=True
            )

            return {
                "clase_id": id_prediccion,
                "clase_nombre": CLASES[id_prediccion],
                "confianza": confianza,
                "probabilidades": probabilidades,
                "imagen_original": imagen_redimensionada,
                "mapa_atencion": mapa_atencion,
                "visualizacion_cam": visualizacion_cam,
                "tiempo_procesamiento": tiempo_inferencia,
                "nombre_modelo": nombre_modelo,
            }
        except Exception as e:
            return {"error": f"Error en {nombre_modelo}: {str(e)}"}

    # función para procesar la imagen con ambos modelos y comparar resultados
    def procesar(self, ruta_imagen: str) -> Dict:
        tiempo_inicio = time.time()

        # validar imagen
        valido, mensaje = self.validar_imagen(ruta_imagen)
        if not valido:
            return {"error": mensaje}

        # procesar con ambos modelos
        resultado_ftp = self.procesar_con_modelo(
            ruta_imagen, self.gestor_ftp, "Modelo FTP"
        )
        resultado_tl = self.procesar_con_modelo(
            ruta_imagen, self.gestor_tl, "Modelo TL"
        )

        # verificar errores
        if "error" in resultado_ftp:
            return resultado_ftp
        if "error" in resultado_tl:
            return resultado_tl

        tiempo_total = time.time() - tiempo_inicio

        # determinar ganador
        ganador = (
            "ftp" if resultado_ftp["confianza"] >= resultado_tl["confianza"] else "tl"
        )

        return {
            "resultado_ftp": resultado_ftp,
            "resultado_tl": resultado_tl,
            "ganador": ganador,
            "tiempo_total": tiempo_total,
        }


# instancias globales de los gestores
config_ftp = ConfigModelo(
    nombre_modelo="Modelo_FTP",
)
config_tl = ConfigModelo(
    nombre_modelo="Modelo_TL",
)
gestor_ftp = GestorModelo(config_ftp)
gestor_tl = GestorModelo(config_tl)
procesador = ProcesadorImagen(gestor_ftp, gestor_tl)


def predecir(imagen) -> Tuple:
    if imagen is None:
        return (
            "Por favor, sube una imagen",
            None,
            None,
            crear_html_error("No hay imagen para procesar"),
            "",
        )

    # procesar imagen con ambos modelos
    resultado = procesador.procesar(imagen)

    # manejar errores
    if "error" in resultado:
        return (
            f"Error: {resultado['error']}",
            None,
            None,
            crear_html_error(resultado["error"]),
            "",
        )

    # extraer resultados de ambos modelos
    res_ftp = resultado["resultado_ftp"]
    res_tl = resultado["resultado_tl"]
    ganador = resultado["ganador"]
    tiempo_total = resultado["tiempo_total"]

    # HTML comparativo de resultados
    html_comparativo = crear_html_comparativo(res_ftp, res_tl, ganador)

    # Eigen-CAMs
    vis_cam_ftp = res_ftp["visualizacion_cam"]
    vis_cam_tl = res_tl["visualizacion_cam"]

    # HTML de confianzas comparativo
    html_confianzas = crear_html_confianzas_comparativo(res_ftp, res_tl, ganador)

    # HTML de métricas
    html_metricas = crear_html_metricas_comparativo(
        res_ftp, res_tl, tiempo_total, DISPOSITIVO
    )

    return (html_comparativo, vis_cam_ftp, vis_cam_tl, html_confianzas, html_metricas)


# función para crear HTML comparativo de los dos modelos con borde dorado para el ganador
def crear_html_comparativo(res_ftp: Dict, res_tl: Dict, ganador: str) -> str:
    # estilos para el ganador
    estilo_ganador = (
        "border: 4px solid {COLOR_BORDE[0]}; box-shadow: 0 0 20px rgba(255, 215, 0, 0.5);"
    )

    # modelo FTP
    estilo_ftp = estilo_ganador if ganador == "ftp" else ""

    # modelo TL
    estilo_tl = estilo_ganador if ganador == "tl" else ""

    html = f"""
    <div style="display: flex; gap: 20px; margin-top: 10px;">
        <div style="flex: 1; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); 
                    padding: 30px; border-radius: 15px; text-align: center; {estilo_ftp}">
            <div style="color: white; font-size: 18px; font-weight: 600; margin-bottom: 10px;">
                MODELO FTP (Fine-Tuning Parcial)
            </div>
            <div style="color: white; font-size: 40px; font-weight: bold; margin: 15px 0;">
                {res_ftp["clase_nombre"]}
            </div>
            <div style="color: {COLOR_BORDE[0]}; font-size: 32px; font-weight: bold;">
                {res_ftp["confianza"]*100:.2f}%
            </div>
            <div style="color: rgba(255, 255, 255, 0.8); font-size: 14px; margin-top: 10px;">
                Tiempo: {res_ftp["tiempo_procesamiento"]:.3f}s
            </div>
        </div>
        
        <div style="flex: 1; background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%); 
                    padding: 30px; border-radius: 15px; text-align: center; {estilo_tl}">
            <div style="color: white; font-size: 18px; font-weight: 600; margin-bottom: 10px;">
                MODELO TL (Transfer Learning)
            </div>
            <div style="color: white; font-size: 40px; font-weight: bold; margin: 15px 0;">
                {res_tl["clase_nombre"]}
            </div>
            <div style="color: {COLOR_BORDE[0]}; font-size: 32px; font-weight: bold;">
                {res_tl["confianza"]*100:.2f}%
            </div>
            <div style="color: rgba(255, 255, 255, 0.8); font-size: 14px; margin-top: 10px;">
                Tiempo: {res_tl["tiempo_procesamiento"]:.3f}s
            </div>
        </div>
    </div>
    """
    return html


# función para crear HTML de confianzas comparando ambos modelos
def crear_html_confianzas_comparativo(res_ftp: Dict, res_tl: Dict, ganador: str) -> str:

    estilo_ganador_ftp = f"border: 3px solid {COLOR_BORDE[0]};" if ganador == "ftp" else ""
    estilo_ganador_tl = f"border: 3px solid {COLOR_BORDE[0]};" if ganador == "tl" else ""

    prob_ftp_buen = res_ftp["probabilidades"][0] * 100
    prob_ftp_def = res_ftp["probabilidades"][1] * 100
    prob_tl_buen = res_tl["probabilidades"][0] * 100
    prob_tl_def = res_tl["probabilidades"][1] * 100

    color_buen = COLORES_CLASES[0]
    color_def = COLORES_CLASES[1]

    html = f"""
    <div style="display: flex; gap: 20px; margin-top: 20px;">
        <!-- Modelo FTP -->
        <div style="flex: 1; {estilo_ganador_ftp} border-radius: 10px; padding: 5px;">
            <div style="text-align: center; color: #667eea; font-weight: bold; margin-bottom: 10px;">
                Modelo FTP
            </div>
            
            <div style="background: linear-gradient(135deg, #d4edda 0%, #c3e6cb 100%); 
                        padding: 20px; border-radius: 10px; margin-bottom: 10px;">
                <div style="color: #155724; font-size: 12px; font-weight: 600; margin-bottom: 5px;">
                    ✅ Buen Estado
                </div>
                <div style="color: {color_buen}; font-size: 32px; font-weight: bold;">
                    {prob_ftp_buen:.1f}%
                </div>
                <div style="background-color: rgba(40, 167, 69, 0.2); height: 6px; border-radius: 3px; margin-top: 8px; overflow: hidden;">
                    <div style="background-color: {color_buen}; height: 100%; width: {prob_ftp_buen}%;"></div>
                </div>
            </div>
            
            <div style="background: linear-gradient(135deg, #f8d7da 0%, #f5c6cb 100%); 
                        padding: 20px; border-radius: 10px;">
                <div style="color: #721c24; font-size: 12px; font-weight: 600; margin-bottom: 5px;">
                    ⚠️ Defectuoso
                </div>
                <div style="color: {color_def}; font-size: 32px; font-weight: bold;">
                    {prob_ftp_def:.1f}%
                </div>
                <div style="background-color: rgba(220, 53, 69, 0.2); height: 6px; border-radius: 3px; margin-top: 8px; overflow: hidden;">
                    <div style="background-color: {color_def}; height: 100%; width: {prob_ftp_def}%;"></div>
                </div>
            </div>
        </div>
        
        <!-- Modelo TL -->
        <div style="flex: 1; {estilo_ganador_tl} border-radius: 10px; padding: 5px;">
            <div style="text-align: center; color: #f5576c; font-weight: bold; margin-bottom: 10px;">
                Modelo TL
            </div>
            
            <div style="background: linear-gradient(135deg, #d4edda 0%, #c3e6cb 100%); 
                        padding: 20px; border-radius: 10px; margin-bottom: 10px;">
                <div style="color: #155724; font-size: 12px; font-weight: 600; margin-bottom: 5px;">
                    ✅ Buen Estado
                </div>
                <div style="color: {color_buen}; font-size: 32px; font-weight: bold;">
                    {prob_tl_buen:.1f}%
                </div>
                <div style="background-color: rgba(40, 167, 69, 0.2); height: 6px; border-radius: 3px; margin-top: 8px; overflow: hidden;">
                    <div style="background-color: {color_buen}; height: 100%; width: {prob_tl_buen}%;"></div>
                </div>
            </div>
            
            <div style="background: linear-gradient(135deg, #f8d7da 0%, #f5c6cb 100%); 
                        padding: 20px; border-radius: 10px;">
                <div style="color: #721c24; font-size: 12px; font-weight: 600; margin-bottom: 5px;">
                    ⚠️ Defectuoso
                </div>
                <div style="color: {color_def}; font-size: 32px; font-weight: bold;">
                    {prob_tl_def:.1f}%
                </div>
                <div style="background-color: rgba(220, 53, 69, 0.2); height: 6px; border-radius: 3px; margin-top: 8px; overflow: hidden;">
                    <div style="background-color: {color_def}; height: 100%; width: {prob_tl_def}%;"></div>
                </div>
            </div>
        </div>
    </div>
    """
    return html


# función para crear HTML de métricas comparando ambos modelos
def crear_html_metricas_comparativo(
    res_ftp: Dict, res_tl: Dict, tiempo_total: float, dispositivo: str
) -> str:

    diferencia_confianza = abs(res_ftp["confianza"] - res_tl["confianza"]) * 100

    html = f"""
    <div style="display: flex; gap: 15px; margin-top: 15px; flex-wrap: wrap;">
        <div style="background: linear-gradient(135deg, #e3f2fd 0%, #bbdefb 100%); 
                    padding: 15px 20px; border-radius: 10px; flex: 1; min-width: 140px;
                    box-shadow: 0 2px 4px rgba(0,0,0,0.1);">
            <div style="color: #1565c0; font-size: 11px; font-weight: 600; margin-bottom: 5px;">
                TIEMPO TOTAL
            </div>
            <div style="color: #0d47a1; font-size: 22px; font-weight: bold;">
                {tiempo_total:.3f}s
            </div>
        </div>
        
        <div style="background: linear-gradient(135deg, #fff3e0 0%, #ffe0b2 100%); 
                    padding: 15px 20px; border-radius: 10px; flex: 1; min-width: 140px;
                    box-shadow: 0 2px 4px rgba(0,0,0,0.1);">
            <div style="color: #e65100; font-size: 11px; font-weight: 600; margin-bottom: 5px;">
                DIFERENCIA DE CONFIANZA
            </div>
            <div style="color: #bf360c; font-size: 22px; font-weight: bold;">
                {diferencia_confianza:.2f}%
            </div>
        </div>
        
        <div style="background: linear-gradient(135deg, #f3e5f5 0%, #e1bee7 100%); 
                    padding: 15px 20px; border-radius: 10px; flex: 1; min-width: 140px;
                    box-shadow: 0 2px 4px rgba(0,0,0,0.1);">
            <div style="color: #6a1b9a; font-size: 11px; font-weight: 600; margin-bottom: 5px;">
                DISPOSITIVO
            </div>
            <div style="color: #4a148c; font-size: 22px; font-weight: bold; text-transform: uppercase;">
                {dispositivo}
            </div>
        </div>
        
        <div style="background: linear-gradient(135deg, #e8f5e9 0%, #c8e6c9 100%); 
                    padding: 15px 20px; border-radius: 10px; flex: 1; min-width: 140px;
                    box-shadow: 0 2px 4px rgba(0,0,0,0.1);">
            <div style="color: #2e7d32; font-size: 11px; font-weight: 600; margin-bottom: 5px;">
                ARQUITECTURA
            </div>
            <div style="color: #1b5e20; font-size: 18px; font-weight: bold;">
                ViT-Base/16
            </div>
        </div>
    </div>
    """
    return html


def crear_html_error(mensaje: str) -> str:
    html = f"""
    <div style="background-color: #fff3cd; border: 2px solid #ffc107; 
                border-radius: 10px; padding: 20px; margin-top: 20px;">
        <div style="color: #856404; font-size: 16px; font-weight: 600; margin-bottom: 10px;">
            Notificación
        </div>
        <div style="color: #856404; font-size: 14px;">
            {mensaje}
        </div>
    </div>
    """
    return html


def crear_interfaz() -> Tuple[gr.Blocks, str]:
    # CSS personalizado
    css_personalizado = """
    .gradio-container {
        font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif !important;
    }
    .resultado-principal {
        font-size: 24px !important;
        font-weight: bold !important;
        padding: 20px !important;
        border-radius: 10px !important;
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
        color: white !important;
        text-align: center !important;
    }
    footer {
        display: none !important;
    }
    """

    with gr.Blocks(title="Clasificador de Papas (Chaucha y Chola) - ViT") as demo:

        # encabezado
        gr.Markdown("""
        # 🥔 Clasificador de Calidad de Papas
        ### Vision Transformer (ViT-Base/16) - Fine Tuning parcial y Transfer Learning
        
        Sube una imagen de un tubérculo de papa (Chaucha o Chola) para comparar el rendimiento de dos modelos entrenados mediante Transfer Learning y Fine Tuning parcial.
        El sistema muestra los resultados, destacando con **borde dorado** el modelo con mayor confianza.
        """)

        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### Imagen de Entrada")
                imagen_entrada = gr.Image(
                    label="Subir imagen de papa", type="filepath", height=400
                )

                gr.Markdown(f"""
                **Especificaciones:**
                - Formatos soportados: PNG, JPG, JPEG
                - Tamaño máximo: {TAMANIO_MAXIMO_MB} MB
                - Resolución: Se redimensiona a {RESOLUCION[0]}×{RESOLUCION[1]}
                """)

                boton_predecir = gr.Button("Clasificar", variant="primary", size="lg")
                boton_limpiar = gr.ClearButton(
                    components=[imagen_entrada], value="Limpiar", size="md"
                )

            with gr.Column(scale=2):
                gr.Markdown("### Comparación de Resultados")

                resultado_comparativo = gr.HTML(label="Comparación de Predicciones")

                html_confianzas = gr.HTML(label="Probabilidades Detalladas")

        # mapas de atención
        gr.Markdown("### Mapas de Atención (Eigen-CAM)")
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("#### Modelo FTP (Fine Tuning parcial)")
                imagen_gradcam_ftp = gr.Image(
                    label="Eigen-CAM - Modelo FTP", height=350
                )

            with gr.Column(scale=1):
                gr.Markdown("#### Modelo TL (Transfer Learning)")
                imagen_gradcam_tl = gr.Image(label="Eigen-CAM - Modelo TL", height=350)

        # métricas
        gr.Markdown("### Métricas de Rendimiento")
        html_metricas = gr.HTML(label="Métricas Comparativas")

        # información adicional
        with gr.Accordion("Información del Sistema", open=False):
            gr.Markdown(f"""
            ### Detalles Técnicos
            
            - **Modelos**: 
              - **Modelo FTP**: Fine Tuning parcial (6 últimas capas ajustadas) del ViT-Base/16.
              - **Modelo TL**: Transfer Learning (solo cabeza clasificadora ajustada).
            - **Framework**: PyTorch.
            - **Visualización**: EigenCAM (Técnica optimizada para Transformers).
            - **Clases**: Buen estado / Defectuoso.
            - **Arquitectura**: ViT-Base/16 (86M parámetros).
            
            ### Interpretación del Mapa de Atención
            
            El mapa de atención (Eigen-CAM) muestra en **amarillo/naranja** las regiones de la imagen 
            que más influyeron en la decisión del modelo, mientras que las regiones azules representan áreas menos relevantes. 
            
            ### Criterio del modelo "Ganador"
            
            El modelo con **mayor confianza** en su predicción se destaca con un **borde dorado**; esto no necesariamente significa que sea correcto, solo que está más seguro de su decisión.
            
            ### Entrenamiento
            
            Ambos modelos fueron entrenados con imágenes de tubérculos de papas Chaucha y Chola, 
            capturando 3 tipos de defectos (brotes, podridas y cortadas) en condiciones de iluminación natural variable.
            """)

        # conectar eventos
        boton_predecir.click(
            fn=predecir,
            inputs=[imagen_entrada],
            outputs=[
                resultado_comparativo,
                imagen_gradcam_ftp,
                imagen_gradcam_tl,
                html_confianzas,
                html_metricas,
            ],
        )

        # predecir al subir imagen
        imagen_entrada.change(
            fn=predecir,
            inputs=[imagen_entrada],
            outputs=[
                resultado_comparativo,
                imagen_gradcam_ftp,
                imagen_gradcam_tl,
                html_confianzas,
                html_metricas,
            ],
        )

        # footer
        gr.Markdown("""
        ---
        **Nota**: Este sistema compara dos enfoques de entrenamiento (FTP y TL). Para decisiones críticas de calidad, consulte con un especialista en agronomía.
        """)

    return demo, css_personalizado


def main():
    print("=" * 70)
    print("🥔 Clasificador de Papas - Comparación de Modelos ViT")
    print("=" * 70)
    print(f"Dispositivo: {DISPOSITIVO.upper()}")
    print(f"Resolución: {RESOLUCION}")
    print("=" * 70)

    # Pre-cargar ambos modelos
    print("\nCargando modelos...")

    print("Cargando Modelo FTP (Fine-Tuning parcial)...")
    if gestor_ftp.cargar_modelo():
        print("✅ Modelo FTP cargado exitosamente")
    else:
        print("El Modelo FTP se cargará cuando sea necesario")

    print("\nCargando Modelo TL (Transfer Learning)...")
    if gestor_tl.cargar_modelo():
        print("✅ Modelo TL cargado exitosamente")
    else:
        print("El Modelo TL se cargará cuando sea necesario")

    print("\n" + "=" * 70)

    # Crear y lanzar interfaz
    demo, css_personalizado = crear_interfaz()

    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
        quiet=False,
        theme=gr.themes.Soft(primary_hue="blue", secondary_hue="green"),
        css=css_personalizado,
    )


if __name__ == "__main__":
    main()