import numpy as np
import gradio as gr
import tensorflow as tf
from tensorflow.keras.models import load_model
from tensorflow.keras.applications.efficientnet import preprocess_input
from tensorflow.keras.preprocessing import image

# Load model
model = load_model("best_model_finetuned.h5")

CLASS_NAMES = ['fire_disaster', 'land_disaster', 'not_disaster', 'water_disaster']

def predict(img):
    img = img.resize((224, 224))
    img_array = np.array(img)
    img_array = np.expand_dims(img_array, axis=0)
    img_array = preprocess_input(img_array)

    prediction = model.predict(img_array)
    predicted_index = np.argmax(prediction)
    confidence = float(np.max(prediction))

    return {
        CLASS_NAMES[i]: float(prediction[0][i])
        for i in range(len(CLASS_NAMES))
    }

interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(),
    title="Disaster Classification CNN",
    description="Upload an image to classify disaster type"
)

interface.launch()