Usando from_pretrained_fastai para carga limpia

app.py

@@ -9,32 +9,29 @@ categories = ['Bug', 'Dark', 'Dragon', 'Electric', 'Fairy', 'Fighting',
               'Normal', 'Poison', 'Psychic', 'Rock', 'Steel', 'Water']
 
 def load_clean_model(weights_path):
-    # Creamos un DataBlock manualmente. 
-    # Esto define la estructura sin intentar leer ninguna lista.
-    dblock = DataBlock(
-        blocks=(ImageBlock, CategoryBlock(vocab=categories)),
-        item_tfms=Resize(126)
-    )
-    
-    # Creamos un DataLoader vacío que solo sirve para inicializar el Learner
-    # Usamos dls de un objeto dummy que no requiere archivos reales
-    dls = dblock.dataloaders([PILImage.create(np.zeros((126,126,3), dtype=np.uint8))], 
-                             [categories[0]], bs=1)
-    
     # Reconstruimos la arquitectura ConvNeXt Tiny
-    learn = vision_learner(dls, 'convnext_tiny', pretrained=False)
+    # num_classes debe coincidir con el número de categorías
+    model = timm.create_model('convnext_tiny', pretrained=False, num_classes=len(categories))
+    
+    # Creamos un Learner mínimo. Al usar dls.from_empty no hay procesos de carga de datos.
+    dls = DataLoaders.from_empty(categories)
+    learn = Learner(dls, model, loss_func=CrossEntropyLossFlat())
     
     # Cargamos los pesos (.pth)
-    # Fastai busca el archivo weights_path + '.pth'
+    # Fastai buscará el archivo weights_path + '.pth' (checkpoint_1.pth)
     learn.load(weights_path) 
     return learn
 
-# 2. Inicializar (Cargará checkpoint_1.pth)
+# 2. Inicializar el modelo
 learn = load_clean_model('checkpoint_1')
 
 def predict(img):
+    # PILImage.create prepara la imagen para Fastai
     img = PILImage.create(img).resize((126, 126))
+    
+    # El método predict de Learner aplicará las transformaciones necesarias
     pred, pred_idx, probs = learn.predict(img)
+    
     return {categories[i]: float(probs[i]) for i in range(len(categories))}
 
 # 3. Interfaz de Gradio