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

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

IMG_SIZE = (224, 224)
class_names = ["buildings", "forest", "glacier", "mountain", "sea", "street"]


# ==== Predict function ====
def predict(img: Image.Image):
    img = img.resize(IMG_SIZE)
    arr = np.array(img) / 255.0
    arr = np.expand_dims(arr, axis=0)

    preds = model.predict(arr)[0]

    return {class_names[i]: float(preds[i]) for i in range(len(class_names))}


# ==== Gradio UI ====
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=gr.Label(num_top_classes=3),
    title="Scene Classification (DenseNet121)",
    description="Upload an image and the model predicts one of 6 scene categories.",
    allow_flagging="never"
)

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