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

# Define categories
categories = ['dog', 'cat']

# Define the function that was used during training
def is_cat(x): 
    return x[0].isupper()

# Load the model
learn = load_learner('model.pkl')

# Gradio prediction function
def predict_image(img):
    pred, idx, probs = learn.predict(img)
    return {categories[i]: float(probs[i]) for i in range(len(categories))}

# Create Gradio interface components
image = gr.components.Image()
label = gr.components.Label()
examples = [['dog.jpg'], ['cat.jpg']]

# Create and launch the interface
interface = gr.Interface(
    fn=predict_image,
    inputs=image,
    outputs=label,
    examples=examples,
    title="Cat vs Dog Classifier"
)

if __name__ == "__main__":
    interface.launch()