from pathlib import Path

import gradio as gr
from fastai.vision.all import *
from huggingface_hub import from_pretrained_fastai

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)
    return {labels[i]: float(probs[i]) for i, _ in enumerate(labels)}


examples = [str(path) for path in EXAMPLES_PATH.iterdir()] if EXAMPLES_PATH.exists() else None

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", height=192, width=192),
    outputs=gr.Label(num_top_classes=3),
    title="Identify which flower it is?",
    description="Identify which flower variety it is by uploading an image.",
    examples=examples,
    flagging_mode="never",
)
demo.queue().launch(share=False)