import gradio as gr
from huggingface_hub import from_pretrained_fastai
from fastai.vision.all import *
repo_id = "hugginglearners/flowers_101_convnext_model"

learn = from_pretrained_fastai(repo_id)
labels = learn.dls.vocab
EXAMPLES_PATH = Path('./examples')


def predict(img):
    img = PILImage.create(img)
    _pred, _pred_w_idx, probs = learn.predict(img)
    # gradio doesn't support tensors, so converting to float
    labels_probs = {labels[i]: float(probs[i]) for i, _ in enumerate(labels)}
    return labels_probs

interface_options = {
    "title": "Identify which flower it is?",
    "description": "I am terribly bad at remembering names of flowers and trees and it's often difficult to fathom how diverse our natural world is.There are over 5,000 species of mammals, 10,000 species of birds, 30,000 species of fish – and astonishingly, over 400,000 different types of flowers.\n Identify which flower variety it is by uploading your images of flowers.",
    "interpretation": "default",
    "layout": "horizontal",
    "allow_flagging": "never",
    "examples": [f'{EXAMPLES_PATH}/{f.name}' for f in EXAMPLES_PATH.iterdir()],
}

demo = gr.Interface(
    fn=predict,
    inputs=gr.inputs.Image(shape=(192, 192)),
    outputs=gr.outputs.Label(num_top_classes=3),
    **interface_options,
)

launch_options = {
    "enable_queue": True,
    "share": True,
}

demo.launch(**launch_options)