Arquitectura modular con funciones get_x/y reales

app.py

@@ -1,49 +1,31 @@
 
 import gradio as gr
-import torch
-import timm
-from PIL import Image
 from fastai.vision.all import *
+import __main__
 
-# 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']
+# 1. Importar las funciones reales del archivo utils
+from utils import get_x, get_y, ds_combined
 
-def load_model_fastai_style(weights_path):
-    # Creamos dls mínimos
-    empty_ds = Datasets([None], [[]])
-    dls = DataLoaders.from_dsets(empty_ds, empty_ds, path='.', bs=1, device='cpu')
-    dls.vocab = categories
-    
-    # Especificamos loss_func explícitamente para evitar el AssertionError
-    learn = vision_learner(dls, 'convnext_tiny', pretrained=False, loss_func=CrossEntropyLossFlat())
-    
-    # Cargar pesos
-    state = torch.load(weights_path, map_location='cpu', weights_only=False)
-    if isinstance(state, dict) and 'model' in state:
-        learn.model.load_state_dict(state['model'])
-    else:
-        learn.model.load_state_dict(state)
-        
-    learn.model.eval()
-    return learn
+# 2. Inyectarlas en el Namespace principal
+__main__.get_x = get_x
+__main__.get_y = get_y
+__main__.ds_combined = ds_combined
 
-# Carga
-learn = load_model_fastai_style('checkpoint_1.pth')
+# 3. Carga del modelo exportado (.pkl)
+# Asegúrate de que el archivo se llame 'modelo_pokemon.pkl' en tu repo
+learn = load_learner('modelo_pokemon.pkl')
+
+labels = learn.dls.vocab
 
 def predict(img):
-    img = PILImage.create(img).resize((126, 126))
+    img = PILImage.create(img)
+    # Al usar el predict del learner cargado, usará las funciones get_x/y que inyectamos
     pred, pred_idx, probs = learn.predict(img)
-    return {categories[i]: float(probs[i]) for i in range(len(categories))}
+    return {labels[i]: float(probs[i]) for i in range(len(labels))}
 
-# Interfaz
-demo = gr.Interface(
+gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="pil"),
+    inputs=gr.Image(),
     outputs=gr.Label(num_top_classes=3),
-    title="Pokemon Type Classifier",
-    description="Identifica el tipo principal de tu Pokémon."
-)
-
-demo.launch()
+    title="Detector de Tipos Pokémon (Espejo de Colab)"
+).launch()

utils.py

@@ -0,0 +1,17 @@
+
+from fastai.vision.all import *
+
+# Simulamos el objeto que tenías en Colab para que las funciones no den NameError
+# Pero lo dejamos vacío porque en la App predeciremos imágenes nuevas, no del dataset
+ds_combined = {} 
+
+def get_x(i): 
+    # En entrenamiento usabas: ds_combined[i]['image_data'].convert('RGB')
+    # En la App, el 'i' que recibirá será la imagen que subas directamente
+    if isinstance(i, (PILImage, Image.Image)): return i.convert('RGB')
+    return ds_combined[i]['image_data'].convert('RGB')
+
+def get_y(i): 
+    # En entrenamiento: ds_combined[i]['Type 1']
+    if isinstance(ds_combined, dict) and len(ds_combined) == 0: return "Unknown"
+    return ds_combined[i]['Type 1']