import os
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 = "NathanSegers/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 it
model = tf.keras.models.load_model("./model")

# 3) simple 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)
).launch()