Usando from_pretrained_fastai para carga limpia

app.py

@@ -6,22 +6,23 @@ from PIL import Image
 from torchvision import transforms
 from fastai.vision.all import *
 
-# 1. Categorías (Asegúrate de que el orden sea ALFABÉTICO, que es el defecto de Fastai)
+# 1. Categorías (Asegúrate de que el orden sea el mismo que dls.vocab en Colab)
 categories = ['Bug', 'Dark', 'Dragon', 'Electric', 'Fairy', 'Fighting', 
               'Fire', 'Flying', 'Ghost', 'Grass', 'Ground', 'Ice', 
               'Normal', 'Poison', 'Psychic', 'Rock', 'Steel', 'Water']
 
 def load_model_fastai_style(weights_path):
-    # Creamos un Learner vacío para que Fastai construya la arquitectura COMPLETA
-    # (Cuerpo de ConvNeXt + Cabeza de clasificación de Fastai)
-    dls = DataLoaders.from_empty(categories)
+    # En lugar de from_empty, creamos dls mínimos con Datasets vacíos
+    # Esto es universal y no depende de métodos de clase que cambian
+    empty_ds = Datasets([None], [[]])
+    dls = DataLoaders.from_dsets(empty_ds, empty_ds, path='.', bs=1, device='cpu')
+    dls.vocab = categories
+    
+    # Construimos el Learner
     learn = vision_learner(dls, 'convnext_tiny', pretrained=False)
     
-    # Cargamos el archivo .pth
-    # Usamos torch.load directamente para mayor control
+    # Cargar pesos con control de errores
     state = torch.load(weights_path, map_location='cpu', weights_only=False)
-    
-    # Si guardaste con learn.save, los pesos están en state['model']
     if isinstance(state, dict) and 'model' in state:
         learn.model.load_state_dict(state['model'])
     else:
@@ -34,23 +35,24 @@ def load_model_fastai_style(weights_path):
 learn = load_model_fastai_style('checkpoint_1.pth')
 
 def predict(img):
-    # Usamos PILImage de fastai para que el preprocesamiento sea IDÉNTICO a Colab
+    # 1. Convertir a PILImage de Fastai
     img = PILImage.create(img)
-    
-    # IMPORTANTE: Forzamos el resize al tamaño de entrenamiento
+    # 2. Resize manual al tamaño de entrenamiento
     img = img.resize((126, 126))
     
-    # Usamos el método predict del learner, que ya sabe normalizar
+    # 3. Predicción
+    # El método predict gestiona internamente la normalización que el modelo espera
     pred, pred_idx, probs = learn.predict(img)
     
     return {categories[i]: float(probs[i]) for i in range(len(categories))}
 
+# Interfaz
 demo = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),
     outputs=gr.Label(num_top_classes=3),
-    title="Pokemon Type Classifier (Sincronizado)",
-    description="Inferencia con arquitectura completa de Fastai."
+    title="Pokemon Type Classifier",
+    description="Identifica el tipo principal de tu Pokémon."
 )
 
 demo.launch()