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

# Load your model
model = tf.keras.models.load_model("disaster_mobilenet.keras")

# Your class names
CLASS_NAMES = ["collapsed_building", "fire", "flooded_areas", "normal", "traffic_incident"]

def predict(image):
    # Preprocess
    image = image.resize((96,96))
    img_array = np.array(image)/255.0
    img_array = np.expand_dims(img_array, 0)  # batch
    
    # Predict
    preds = model.predict(img_array)
    class_idx = np.argmax(preds[0])
    return {CLASS_NAMES[i]: float(preds[0][i]) for i in range(len(CLASS_NAMES))}

# Gradio interface
iface = gr.Interface(fn=predict, inputs=gr.Image(type="pil"), outputs=gr.Label())
iface.launch()