import tensorflow as tf
import gradio as gr
import numpy as np
from huggingface_hub import hf_hub_download

# 1) Download your SavedModel from the Hub
repo_id = "Agent-37/masterclass-2025"
hf_hub_download(repo_id, filename="config.json", repo_type="model", local_dir="./model")
hf_hub_download(repo_id, filename="metadata.json", repo_type="model", local_dir="./model")
hf_hub_download(repo_id, filename="model.weights.h5", repo_type="model", local_dir="./model")

# 2) Load the model
model = tf.keras.models.load_model("./model")

# 3) Preprocess + predict
CLASS_NAMES = ["cat", "dog", "panda"]
def predict(image):
    resized_image = tf.image.resize(image, (64,64))
    images_to_predict = np.expand_dims(np.array(resized_image), axis=0)
    probs = model.predict(images_to_predict)[0]
    return {c: float(p) for c,p in zip(CLASS_NAMES, probs)}

# 4) Launch Gradio
gr.Interface(
    fn=predict,
    inputs=gr.Image(),
    outputs=gr.Label(num_top_classes=3),
    title="🐾 Animal Classifier",
    description="Upload an image of a cat, dog, or panda."
).launch()