Spaces:
Sleeping
Sleeping
import os | |
from huggingface_hub import from_pretrained_fastai | |
import gradio as gr | |
from gradio.components import Image | |
HF_TOKEN = os.getenv("HF_TOKEN") | |
print(HF_TOKEN[-1:-5]) | |
repo_id = "artificeresearch/spiritvision" | |
learner = from_pretrained_fastai(repo_id) | |
labels = learner.dls.vocab | |
def predict_fn(img): | |
""" | |
:param img: img is a PIL image object | |
:return: prediction and probabilities | |
""" | |
pred, pred_idx, probs = learner.predict(img) | |
return {labels[i]: float(probs[i]) for i in range(len(labels))} | |
gr.Interface(predict_fn, | |
Image(shape=(512, 512)), | |
outputs=gr.outputs.Label(num_top_classes=3), | |
token=HF_TOKEN).launch() | |