Usando from_pretrained_fastai para carga limpia

app.py

@@ -1,64 +1,56 @@
 
-import gradio as gr
 import gradio as gr
 import torch
 import timm
 from PIL import Image
 from torchvision import transforms
+from fastai.vision.all import *
 
-# 1. Lista de categorías (Asegúrate de que el orden sea el mismo que en tu vocab de Colab)
+# 1. Categorías (Asegúrate de que el orden sea ALFABÉTICO, que es el defecto de Fastai)
 categories = ['Bug', 'Dark', 'Dragon', 'Electric', 'Fairy', 'Fighting', 
               'Fire', 'Flying', 'Ghost', 'Grass', 'Ground', 'Ice', 
               'Normal', 'Poison', 'Psychic', 'Rock', 'Steel', 'Water']
 
-# 2. Carga del modelo directamente en PyTorch
-def load_model(weights_path):
-    # Creamos la arquitectura exacta
-    model = timm.create_model('convnext_tiny', pretrained=False, num_classes=len(categories))
+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)
+    learn = vision_learner(dls, 'convnext_tiny', pretrained=False)
     
-    # Cargamos los pesos del archivo .pth
-    # weights_only=False es necesario para cargar el estado guardado por fastai
-    state_dict = torch.load(weights_path, map_location='cpu', weights_only=False)
+    # Cargamos el archivo .pth
+    # Usamos torch.load directamente para mayor control
+    state = torch.load(weights_path, map_location='cpu', weights_only=False)
     
-    # Si los pesos vienen de un Learner de fastai, suelen estar dentro de una llave 'model'
-    if 'model' in state_dict:
-        state_dict = state_dict['model']
+    # 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:
+        learn.model.load_state_dict(state)
         
-    # Cargamos los pesos en la arquitectura
-    model.load_state_dict(state_dict, strict=False)
-    model.eval() # Modo evaluación
-    return model
+    learn.model.eval()
+    return learn
 
-# Cargamos el modelo (el archivo debe estar en el repo como checkpoint_1.pth)
-model = load_model('checkpoint_1.pth')
+# Cargamos el modelo
+learn = load_model_fastai_style('checkpoint_1.pth')
 
-# 3. Función de predicción manual
 def predict(img):
-    # Preprocesamiento: Resize(126) + Convertir a Tensor + Normalización ImageNet
-    # Esto es exactamente lo que hace fastai por debajo
-    transform = transforms.Compose([
-        transforms.Resize((126, 126)),
-        transforms.ToTensor(),
-        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
-    ])
+    # Usamos PILImage de fastai para que el preprocesamiento sea IDÉNTICO a Colab
+    img = PILImage.create(img)
     
-    # Convertimos la imagen de Gradio a Tensor
-    img_tensor = transform(img).unsqueeze(0) # Añadir dimensión de batch
+    # IMPORTANTE: Forzamos el resize al tamaño de entrenamiento
+    img = img.resize((126, 126))
     
-    with torch.no_grad():
-        outputs = model(img_tensor)
-        probs = torch.nn.functional.softmax(outputs[0], dim=0)
+    # Usamos el método predict del learner, que ya sabe normalizar
+    pred, pred_idx, probs = learn.predict(img)
     
-    # Retornamos diccionario con las probabilidades
     return {categories[i]: float(probs[i]) for i in range(len(categories))}
 
-# 4. Interfaz de Gradio
 demo = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),
     outputs=gr.Label(num_top_classes=3),
-    title="Pokemon Type Classifier",
-    description="Inferencia directa usando PyTorch y ConvNeXt Tiny."
+    title="Pokemon Type Classifier (Sincronizado)",
+    description="Inferencia con arquitectura completa de Fastai."
 )
 
 demo.launch()