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

# Load your trained model
learn = load_learner('export.pkl')

# Define the labels (in this case, just 'cat' and 'dog')
labels = ['cat', 'dog']

# Define the prediction function
def predict(img):
    img = PILImage.create(img)
    pred, pred_idx, probs = learn.predict(img)
    # Return a dictionary of probabilities
    return {labels[i]: float(probs[i]) for i in range(len(labels))}

# Update the title, description, and other details for your cat vs. dog classifier
title = "Cat vs Dog Classifier"
description = "A classifier to distinguish between cats and dogs. Trained with fastai on a relevant dataset."
article = "<p style='text-align: center'><a href='https://yourlinkhere.com' target='_blank'>Blog post or additional information</a></p>"
examples = ['/path/to/example_image.jpg'] # Update this path to your example images

# Create and launch the Gradio interface
gr.Interface(
    fn=predict,
    inputs=gr.inputs.Image(shape=(512, 512)),
    outputs=gr.outputs.Label(num_top_classes=2),
    title=title,
    description=description,
    article=article,
    examples=examples,
    interpretation='default',
    enable_queue=True
).launch()