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

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

# Get the labels from the data loaders
labels = learn.dls.vocab

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

# Example images for demonstration
examples = ['image.jpg']

# Create the Gradio interface
interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(),
    outputs=gr.Label(num_top_classes=3)
)

# Enable the queue to handle POST requests
interface.queue(api_open=True)

# Launch the interface
interface.launch()