import gradio as gr
from huggingface_hub import hf_hub_download
from fastai.learner import load_learner
import os

TOKEN = os.environ["token"]
REPO_ID = "42digital/deepfashion_classification_vit-large-patch14-clip-336"
FILENAME = "model.pkl"

learner = load_learner(
    hf_hub_download(repo_id=REPO_ID, filename=FILENAME, token=TOKEN)
)

def predict(img):
    _, _, probs = learner.predict(img)
    probs = [float(p) for p in probs.detach()]
    preds = {k: v for k, v in zip(learner.dls.vocab, probs)}
    return preds

gr.Interface(fn=predict, 
             inputs=gr.Image(type="numpy"),
             outputs=gr.Label(num_top_classes=4),
             examples=[]).launch()