import gradio as gr
import numpy as np
from PIL import Image
from tensorflow.keras.models import load_model

# Load the trained CNN model
model = load_model("model.h5")

# Define class labels
class_names = ["Monkeypox", "Not Monkeypox"]

def predict(img):
    # Resize & preprocess image
    img_resized = img.resize((224, 224))
    img_array = np.array(img_resized) / 255.0  # normalize
    img_array = np.expand_dims(img_array, axis=0)

    # Predict
    preds = model.predict(img_array)

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

# -------- Gradio Interface --------
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        <div style="text-align:center; padding: 15px; background: linear-gradient(90deg, #ff6f61, #ffcc70); border-radius: 12px;">
            <h1 style="color:white;">🐵 Monkeypox Classifier</h1>
            <p style="color:white; font-size:18px;">Upload or capture an image, and the model will classify it as <b>Monkeypox</b> or <b>Not Monkeypox</b>.</p>
        </div>
        """
    )

    with gr.Row():
        with gr.Column():
            input_img = gr.Image(type="pil", label="📸 Upload or Capture Image", sources=["upload", "webcam"])
            predict_btn = gr.Button("🔍 Predict", elem_id="predict-btn")
        with gr.Column():
            output_label = gr.Label(num_top_classes=2, label="Prediction")

    # Add custom CSS
    demo.load(
        lambda: None,
        None,
        None,
        _js="""
        () => {
            let btn = document.getElementById("predict-btn");
            if(btn){
                btn.style.background = "linear-gradient(45deg, #36d1dc, #5b86e5)";
                btn.style.color = "white";
                btn.style.fontWeight = "bold";
                btn.style.padding = "10px 20px";
                btn.style.borderRadius = "12px";
            }
        }
        """
    )

    predict_btn.click(fn=predict, inputs=input_img, outputs=output_label)

# Launch App
if __name__ == "__main__":
    demo.launch()