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

# ----------------------------
# LOAD MODEL (LOCAL FILE)
# ----------------------------
model = tf.keras.models.load_model("CModel.h5")
print(model.input_shape)


IMG_SIZE = (224, 224)

CLASS_NAMES = [
    "Normal",
    "Monkeypox"
]

# ----------------------------
# PREDICTION FUNCTION
# ----------------------------
def predict_image(image):
    image = image.convert("RGB")
    image = image.resize(IMG_SIZE)

    img_array = np.array(image) / 255.0
    img_array = np.expand_dims(img_array, axis=0)

    pred = model.predict(img_array)

    if pred.shape[1] == 1:
        confidence = float(pred[0][0])
        label = CLASS_NAMES[1] if confidence > 0.5 else CLASS_NAMES[0]
        return label, confidence
    else:
        class_index = int(np.argmax(pred))
        confidence = float(pred[0][class_index])
        return CLASS_NAMES[class_index], confidence

# ----------------------------
# GRADIO UI
# ----------------------------
interface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Label(label="Prediction"),
        gr.Number(label="Confidence")
    ],
    title="Monkeypox Classification using CNN",
    description="Upload a skin image to classify Monkeypox using a CNN model."
)

interface.launch()