from fastai.vision.all import *
import gradio as gr
import glob

examples = glob.glob("examples/*")
learn = load_learner("model.pkl")

labels = learn.dls.vocab
def classify_image(img):
    pred, idx, probs = learn.predict(img)
    return {labels[i]: float(probs[i]) for i in range(len(labels))}
image = gr.inputs.Image(shape=(192,192))
label = gr.outputs.Label()

title = "Pasta Recognition"

description = """
Pasta is delicious, but there are too many pasta types to remember. This app will help you
recall the name of the pasta whose image you'll upload.

Under the hood it's a pasta classifier, built using Fast.ai by fine-tuning the
`convnext_small_in22k` model. It supports only the followig pasta types:

* Campanelle
* Ditalini
* Fettuccine
* Fusilli
* Linguine
* Macaroni
* Pappardelle

The model's accuracy is ~87.3%.

The notebook used to train the model can be seen [in Google Colab](https://colab.research.google.com/drive/1NM6BvRkV7xDlnncFqVRGxxg0GFtqpVgT?usp=sharing).
"""

iface = gr.Interface(fn=classify_image, inputs=image, outputs=gr.outputs.Label(), examples=examples, title=title, description=description)
iface.launch(inline=False, share=False)