from fastai.vision.all import *
from huggingface_hub import hf_hub_download
import gradio as gr

# Re-declare is_cat before load_learner
def is_cat(x): return x[0].isupper()

# Download model at runtime
model_path = hf_hub_download(repo_id="JackStennett/cat-dog", filename="model.pkl")
learn = load_learner(model_path)

categories = ('Dog', 'Cat')

def classify_image(img):
    pred, idx, probs = learn.predict(img)
    return {categories[idx]: float(probs[idx])}

intf = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(label="Upload an Image"),
    outputs=gr.Label(num_top_classes=2),
    title="🐶 vs 🐱 Classifier",
    description="Upload an image of a cat or a dog and the model will tell you which one it is."
)

if __name__ == "__main__":
    intf.launch()