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

# Load the exported fastai Learner
learn = load_learner('model.pkl')
labels = learn.dls.vocab  # ['bird', 'shark']

# Prediction function
def predict(img):
    pred, pred_idx, probs = learn.predict(img)
    return {labels[i]: float(probs[i]) for i in range(len(labels))}

# Interface metadata
title = "Shark vs Bird Classifier"
description = "A simple image classifier that distinguishes between birds and sharks using fastai and Gradio."

# Optional: Example images must exist in the same directory
examples = ['shark.jpg', 'bird.jpg']

# Launch the Gradio interface
gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=2),
    title=title,
    description=description,
    examples=examples,
    allow_flagging='never'
).launch()