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

# Import trained model
learn = load_learner("model.pkl")

# Define an object with labels (keys) and tensors (values)
categories = {"Dog", "Cat"}
def classify_image(img):
    pred,_,probs = learn.predict(img)
    return dict(zip(categories, map(float,probs)))

# Build the Gradio interface
image = gr.inputs.Image(shape=(192, 192))
label = gr.outputs.Label()
examples = ["examples/max.jpg", "examples/nymo.jpg"]
    
intf = gr.Interface(
    fn=classify_image, 
    inputs=image, 
    outputs=label, 
    examples=examples
)

# Start the server
intf.launch(inline=False)