import requests
from huggingface_hub import from_pretrained_fastai

import gradio as gr

repo_id = "artificeresearch/spiritvision"
learner = from_pretrained_fastai(repo_id)


def predict_fn(img):
    """
    :param img: img is a PIL image object
    :return: prediction and probabilities
    """
    img = img.convert('RGB')
    prediction, _, probs = learner.predict(img)
    # print(f'{max(100 * probs):.2f}% {prediction} - {img}')
    return f'{max(100 * probs):.2f}% {prediction} - {img}'


gr.Interface(predict_fn, gr.inputs.Image(type='pil'), outputs='label').launch()