from pathlib import Path

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

EXAMPLES_PATH = Path("./examples")
repo_id = "hugginglearners/multi-object-classification"
learner = from_pretrained_fastai(repo_id)
labels = learner.dls.vocab


def infer(img):
    img = PILImage.create(img)
    _pred, _pred_w_idx, probs = learner.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(
    infer,
    gr.Image(type="pil", height=192, width=192),
    gr.Label(num_top_classes=3),
    examples=examples,
    flagging_mode="never",
    title="Multilabel Image classification",
    description="Detect which type of object appears in the image.",
    article='Author: <a href="https://huggingface.co/geninhu">Nhu Hoang</a>.',
    live=False,
)
demo.queue().launch(debug=False, inbrowser=False)