Update app.py

app.py

@@ -1,16 +1,19 @@
 import gradio as gr
-from huggingface_hub import hf_hub_download
-import torch
-import pickle
+from fastai.vision.all import *
+from huggingface_hub import from_pretrained_fastai
+modelo = from_pretrained_fastai("pagaliv/futurama-character-classifier")
 
-# Workaround para cargar el modelo
-model_path = hf_hub_download(repo_id="pagaliv/futurama-character-classifier", filename="model.pkl")
-modelo = torch.load(model_path, pickle_module=pickle, map_location='cpu')
+def clasificar_imagen(img):
+    pred, _, prob = modelo.predict(img)
+    return {modelo.dls.vocab[i]: float(prob[i]) for i in range(len(prob))}
 
-def clasificar(img):
-    # Preprocesa la imagen como lo hacías durante el entrenamiento
-    # (Aquí necesitarás tu lógica personalizada)
-    pred = modelo.predict(img)
-    return pred
+# Interfaz
+iface = gr.Interface(
+    fn=clasificar_imagen,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Label(num_top_classes=3),
+    examples=["BENDER.jpg","BENDER2.jpg", "FRY.jpg","FRY2.jpg","LEELA.jpg","LEELA2.jpg"],  # Asegúrate de subir estos archivos
+    title="Clasificador de Futurama"
+)
 
-gr.Interface(fn=clasificar, inputs="image", outputs="label").launch()
+iface.launch()