return best_model

VisionEnsembleModel.py

@@ -5,14 +5,8 @@ import torch.nn as nn
 import timm
 
 class VisionEnsembleModel(nn.Module):
-    """
-    La misma clase de modelo que definiste en Colab.
-    """
     def __init__(self, num_classes, cnn_model_name='efficientnet_b2', vit_model_name='vit_small_patch16_224'):
         super().__init__()
-        # Usamos pretrained=False aquí porque cargaremos nuestros propios pesos.
-        # Timm cargará los pesos preentrenados si no encuentra un state_dict local,
-        # pero es más limpio ser explícito. Al final, los sobrescribiremos.
         self.cnn = timm.create_model(cnn_model_name, pretrained=False, num_classes=num_classes)
         cnn_features = self.cnn.get_classifier().in_features
         self.cnn.reset_classifier(0)

app.py

@@ -1,16 +1,19 @@
-# app.py (versión con Gradio)
+# app.py (Versión Corregida para Cargar el Modelo Ensamblado)
 
 import torch
 import torchvision.transforms as transforms
 from PIL import Image
 import json
-import timm
 import gradio as gr
 
-# --- 1. Carga del Modelo y Componentes (igual que antes) ---
+# --- ¡CAMBIO 1: Importar la clase de nuestro modelo! ---
+from VisionEnsembleModel import VisionEnsembleModel
 
-device = torch.device("cpu") # Es más seguro usar CPU en el plan gratuito
-MODEL_PATH = "model/best_vision_ensemble_model.pth" # Usaremos el modelo ligero que sabemos que funciona
+# --- 1. Carga del Modelo y Componentes ---
+device = torch.device("cpu")
+
+# --- ¡VERIFICACIÓN! Apuntar al archivo de modelo correcto ---
+MODEL_PATH = "model/best_vision_ensemble_model.pth"
 LABELS_PATH = "model/species_labels_map.json"
 NUM_CLASSES = 156
 
@@ -18,65 +21,44 @@ NUM_CLASSES = 156
 with open(LABELS_PATH) as f:
     labels_map = json.load(f)
 
-# Definir y cargar el modelo ligero
-model = timm.create_model('mobilenetv2_100', pretrained=False, num_classes=NUM_CLASSES)
+# --- ¡CAMBIO 2: Instanciar nuestro modelo personalizado! ---
+model = VisionEnsembleModel(num_classes=NUM_CLASSES)
 model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
 model.to(device)
 model.eval()
 
-print("Modelo MobileNetV2 cargado y listo.")
+print("Modelo VisionEnsembleModel cargado y listo.")
 
-# Definir las transformaciones de la imagen
+# --- 2. Definir la Función de Predicción para Gradio (SIN CAMBIOS) ---
 transforms_val = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 ])
 
-# --- 2. Definir la Función de Predicción para Gradio ---
-
 def predict(image):
-    """
-    Esta función toma una imagen (de Gradio) y devuelve un diccionario de predicciones.
-    """
-    # La imagen de Gradio viene como un array de Numpy, la convertimos a PIL Image
     pil_image = Image.fromarray(image.astype('uint8'), 'RGB')
-    
-    # Preprocesar la imagen
     input_tensor = transforms_val(pil_image).unsqueeze(0).to(device)
-
-    # Realizar la predicción
     with torch.no_grad():
         output = model(input_tensor)
         probabilities = torch.nn.functional.softmax(output[0], dim=0)
     
-    # Crear un diccionario de confianza para las 5 mejores predicciones
     top5_prob, top5_catid = torch.topk(probabilities, 5)
-    
     confidences = {}
     for i in range(top5_prob.size(0)):
         species_id = top5_catid[i].item()
         prob = top5_prob[i].item()
         species_name = labels_map.get(str(species_id), "Desconocido")
         confidences[species_name] = prob
-        
     return confidences
 
-# --- 3. Crear y Lanzar la Interfaz de Gradio ---
-
-# Creamos la interfaz
+# --- 3. Crear y Lanzar la Interfaz de Gradio (SIN CAMBIOS) ---
 iface = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="numpy", label="Sube una imagen de tu orquídea"),
     outputs=gr.Label(num_top_classes=5, label="Predicciones"),
-    title="Clasificador de Orquídeas",
-    description="Sube una foto de una orquídea y la IA intentará identificar la especie. Este modelo usa un MobileNetV2.",
-    examples=[
-        # Puedes añadir rutas a imágenes de ejemplo si las subes a tu repositorio
-        # ["ejemplo1.jpg"],
-        # ["ejemplo2.jpg"]
-    ]
+    title="Clasificador de Orquídeas (Modelo Ensamblado)",
+    description="Sube una foto de una orquídea y la IA intentará identificar la especie.",
 )
 
-# Lanzamos la aplicación
 iface.launch()