import gradio as gr
import pickle
from PIL import Image

from fastai.learner import load_learner

model = load_learner('./mushrooms.pkl')


categories = [
    "Agaricus",
    "Amanita",
    "Boletus",
    "Cortinarius",
    "Entoloma",
    "Hygrocybe",
    "Lactarius",
    "Russula",
    "Suillus",
]

def classify_image(image):
    prediction, _, probs = model.predict(image)
    return dict(zip(categories, map(float, probs)))

iface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(),
    title="Image Classifier",
    description="WHAT IS THE MUSHROOM?"
)

# Launch app
if __name__ == "__main__":
    iface.launch()