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

# Downloaded the image classification model we trained in lab 6.1  and uploaded it to the huggingface space
MODEL_PATH = "cats-vs-dogs-finetuned.keras"
# the input size the model expects to see
IMAGE_SIZE = (180, 180)  

# Load model
model = tf.keras.models.load_model(MODEL_PATH)

def predict_image(img: Image.Image):
    #if there is no input image, the probability of cat and dog is 50-50
    if img is None:
        return {"Cat": 0.5, "Dog": 0.5}

    # Preprocess
    x = img.convert("RGB").resize(IMAGE_SIZE)
    x = np.asarray(x, dtype=np.float32) / 255.0
    x = np.expand_dims(x, 0)  # (1, H, W, 3)

    # Predict: model outputs shape (1,1) with sigmoid for "Dog" probability
    p_dog = float(model.predict(x)[0, 0])  
    return {"Cat": 1.0 - p_dog, "Dog": p_dog}

demo = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil", label="Upload a Cat or Dog"),
    outputs=gr.Label(num_top_classes=2, label="Prediction"),
    title="Cats vs Dogs (Keras, single-logit)",
    description="keras image classification model for cat-vs-dog images"
)

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