# %%
from typing import Dict, List
import gradio as gr  # type: ignore
from fastai.vision.learner import load_learner, Learner  # type: ignore
from PIL import Image
from numpy import shape  # type: ignore

# %%
learn: Learner = load_learner("bears.pkl")
# %%


def categories() -> List[str]:
    return learn.dls.vocab


def classify_image(img: Image.Image) -> Dict[str, float]:
    pred, idx, probs = learn.predict(img)
    print(f"Predicted the image was a {pred} bear with {probs[idx].item()*100:.02f}%")
    probs = map(float, probs)
    return dict(zip(categories(), probs))


# %%

image_input = gr.Image(width=224, height=224)
label = gr.Label()
iface = gr.Interface(fn=classify_image, inputs=[image_input], outputs=[label])
iface.launch()

# %%