import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image

# Load your saved model (.keras file uploaded to your HF repo)
model = tf.keras.models.load_model("my_model.keras")

# Define class labels (update if your order is reversed)
class_names = ["Cat", "Dog"]

def predict(image):
    # Preprocess image (resize to your model's input size)
    img = image.resize((150, 150))   # change if your model used a different size
    img = np.array(img) / 255.0      # normalize to [0,1]
    img = np.expand_dims(img, axis=0)

    # Run prediction
    pred = model.predict(img)[0][0]  # get scalar value
    
    if pred < 0.5:
        return {"Cat": 1 - float(pred), "Dog": float(pred)}
    else:
        return {"Cat": 1 - float(pred), "Dog": float(pred)}

# Build Gradio interface
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=2),
    title="Dog vs Cat Classifier 🐶🐱",
    description="Upload an image of a dog or cat, and the model will predict which one it is."
)

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