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

import platform
import pathlib
plt = platform.system() 
if plt == 'Linux': 
    pathlib.WindowsPath = pathlib.PosixPath

pasta_shape_labels = (
    "bucatini",
    "cannelloni",
    "cavatappi",
    "conchiglie",
    "farfalle",
    "fettuccine",
    "fusilli",
    "gemelli",
    "lasagna",
    "linguine",
    "macaroni",
    "orecchiette",
    "orzo",
    "penne",
    "ravioli",
    "rigatoni",
    "rotini",
    "spaghetti",
    "tagliatelle",
    "tortellini"
)

model = load_learner('pasta_shape_recognizer_v2.pkl')

def recognize_image(image):
    pred, idx, probs = model.predict(image)
    return dict(zip(pasta_shape_labels, map(float, probs)))

image = gr.inputs.Image(shape=(192,192))
label = gr.outputs.Label(num_top_classes=5)
examples = [
    'unknown00.png',
    # 'unknown01.png',
    'unknown02.png',
    'unknown03.png',
    'unknown04.png',
    'unknown05.png',
    'unknown06.png',
    'unknown07.png',
    'unknown08.png',
    'unknown09.png',
    'unknown10.png',
    'unknown11.png',
    'unknown12.png',
    'unknown13.png'
]

iface = gr.Interface(fn=recognize_image, inputs=image, outputs=label, examples=examples)
iface.launch(inline=False)