import gradio as gr
from fastai.learner import load_learner
from fastai.vision.all import PILImage

# Load the model directly (since it will be in the same repository)
model = load_learner('model.pkl')

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

# Create the interface
demo = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=2),
    title="🐱 Cat Detector",
    description="Upload an image to check if it contains a cat!",
)

demo.launch()