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

# =========================
# 1. Load Model
# =========================
model = tf.keras.models.load_model("cnn_model_transfer.keras")

# Class labels (match training order)
class_names = ["Cat", "Dog"]

# =========================
# 2. Preprocess Function
# =========================
def preprocess_image(img):
    img = img.resize((224, 224))   # MobileNetV2 size
    img = np.array(img)

    # Apply same preprocessing used during training
    img = tf.keras.applications.mobilenet_v2.preprocess_input(img)

    img = np.expand_dims(img, axis=0)
    return img

# =========================
# 3. Prediction Function
# =========================
def predict(img):
    img_array = preprocess_image(img)

    pred = model.predict(img_array)[0][0]

    if pred > 0.5:
        label = class_names[1]
        confidence = float(pred)
    else:
        label = class_names[0]
        confidence = float(1 - pred)

    return {
        "Cat": 1 - pred,
        "Dog": pred
    }

# =========================
# 4. Gradio Interface
# =========================
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=2),
    title="🐱🐶 Cat vs Dog Classifier",
    description="Upload an image and the model will predict whether it is a Cat or Dog using MobileNetV2 Transfer Learning."
)

# =========================
# 5. LauncH
# =========================
demo.launch()