# app.py
import gradio as gr
from fastai.learner import load_learner
from fastai.vision.all import PILImage
from PIL import Image

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

def classify_image(image):
    # Convert to FastAI PILImage
    img = PILImage.create(image)
    
    # Get prediction
    pred, pred_idx, probs = model.predict(img)
    
    # Return prediction and confidence
    return {
        "cat": float(probs[pred_idx]) if str(pred) == "cat" else 1 - float(probs[pred_idx])
    }

# Create Gradio interface
iface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(),
    outputs=gr.Label(num_top_classes=2),
    title="Cat Classifier",
    description="Upload an image to check if it contains a cat!",
    examples=["example1.jpg", "example2.jpg"]  # Optional: Add example images if you have them
)

iface.launch()