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"""
MODELO FTP (Fine-Tuning Parcial)
{res_ftp["clase_nombre"]}
{res_ftp["confianza"]*100:.2f}%
Tiempo: {res_ftp["tiempo_procesamiento"]:.3f}s
MODELO TL (Transfer Learning)
{res_tl["clase_nombre"]}
{res_tl["confianza"]*100:.2f}%
Tiempo: {res_tl["tiempo_procesamiento"]:.3f}s
""" 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"""
Modelo FTP
✅ Buen Estado
{prob_ftp_buen:.1f}%
⚠️ Defectuoso
{prob_ftp_def:.1f}%
Modelo TL
✅ Buen Estado
{prob_tl_buen:.1f}%
⚠️ Defectuoso
{prob_tl_def:.1f}%
""" 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"""
TIEMPO TOTAL
{tiempo_total:.3f}s
DIFERENCIA DE CONFIANZA
{diferencia_confianza:.2f}%
DISPOSITIVO
{dispositivo}
ARQUITECTURA
ViT-Base/16
""" return html def crear_html_error(mensaje: str) -> str: html = f"""
Notificación
{mensaje}
""" 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()