from fastai.vision.all import *
import gradio as gr
import pathlib, os


classes = ['rock', 'paper', 'scissors'] # c0, c1, c2

def classify_image(img, model='rock-paper-scissors-resnet34.pkl'):
    if os.name == 'nt': # workaround for Windows
        pathlib.PosixPath = pathlib.WindowsPath 
    if os.name == 'posix': # workaround for Linux
            pathlib.WindowsPath = pathlib.PosixPath 

    learn = load_learner(model)
    pred,idx,probs = learn.predict(img)
    return dict(zip(classes, map(float, probs)))

models = ['rock-paper-scissors-squeezenet.pkl','rock-paper-scissors-resnet34.pkl']
model = gr.Dropdown(models, label="Select Model")

image = gr.inputs.Image(shape=(192,192))
label = gr.outputs.Label()
examples = [
    ['c0-rock-IMG_20230225_171937.jpg'],
    ['c0-rock-IMG_20230225_171940.jpg'],
    ['c1-paper-IMG_20230225_172010.jpg'],
    ['c1-paper-IMG_20230225_172018.jpg'],
    ['c2-scissors-IMG_20230225_172025.jpg'],
    ['c2-scissors-IMG_20230225_172033.jpg']
 ]

# iface = gr.Interface(fn=classify_image, inputs=image, outputs=label, examples=examples)
iface = gr.Interface(fn=classify_image, inputs=[image, model], outputs=label, examples=examples)
iface.launch()