from pathlib import Path

import gradio as gr
import numpy as np

from fastai.learner import load_learner


learner = load_learner(Path("model.pkl"))
labels = learner.dls.vocab


def classify(image: np.array):
    results = learner.predict(image)
    probabilities = results[2]
    probabilities = map(float, probabilities)
    return dict(zip(labels, probabilities))


description = """<p>Give me a cactus picture and I'll try to guess the variety.</p>
<p>I can recognize the following types of cactus:</p>
<ul>
{list_items}
</ul>
<p>This model was trained <a href="https://www.kaggle.com/jdsleppy/cactus-classifier/">here</a> by <a href="https://www.joelsleppy.com">Joel</a>.</p>
""".format(
list_items='\n'.join([f'  <li>{label}</li>' for label in labels])
)


gr.Interface(
    fn=classify,
    inputs="image",
    outputs="label",
    examples=["examples/disco.jpg", "examples/scarlet_crown.jpg", "examples/silver_torch.webp"],
    title="Cactus Classifier",
    description=description,
    allow_flagging="never",
).launch()