# app.py
import gradio as gr
from fastai.learner import load_learner
from fastai.vision.all import PILImage
from huggingface_hub import hf_hub_download
import torch

# Define the is_cat function that was used during training
def is_cat(x): return x[0].isupper()

def load_model():
    # Download the model from your model repository
    model_path = hf_hub_download(
        repo_id="RamyKhorshed/Lesson2FastAi",
        filename="model.pkl",
        repo_type="model"
    )
    return load_learner(model_path)

print("Loading model...")
model = load_model()
print("Model loaded!")

def predict_image(image):
    # Convert to FastAI format
    img = PILImage.create(image)
    
    # Predict
    pred, pred_idx, probs = model.predict(img)
    
    # Format output
    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 interface
demo = gr.Interface(
    fn=predict_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!",
    article="Upload any image and the model will predict whether it contains a cat or not."
)

demo.launch()