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

# ===== CONFIG =====
IMG_SIZE = 300
CLASSES = ["BACTERIAL", "NORMAL", "VIRAL"]
MODEL_PATH = "DenseNet201Final.keras"      # adjust path if needed
LAST_CONV_LAYER = "conv5_block32_concat"       # must exist in model summary

# ===== Load Model =====
model = tf.keras.models.load_model(MODEL_PATH, compile=False)

#===== Grad-CAM Function =====
def grad_cam(model, img_array, layer_name):
    grad_model = tf.keras.models.Model(
        inputs=model.inputs,
        outputs=[
            model.get_layer(layer_name).output,
            model.output
        ]
    )

    with tf.GradientTape() as tape:
        conv_outputs, preds = grad_model(img_array)

        # FIX: ensure preds is a tensor, not a list
        if isinstance(preds, (list, tuple)):
            preds = preds[0]

        class_idx = tf.argmax(preds[0])
        loss = preds[0, class_idx]

    grads = tape.gradient(loss, conv_outputs)
    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))

    conv_outputs = conv_outputs[0]
    heatmap = tf.reduce_sum(pooled_grads * conv_outputs, axis=-1)

    heatmap = tf.maximum(heatmap, 0)
    heatmap /= tf.reduce_max(heatmap) + 1e-10

    return heatmap.numpy(), int(class_idx)


# ===== Prediction Function =====
def predict_xray(image):
    if image is None:
        return None, "No image uploaded"

    # Convert PIL → NumPy
    img = np.array(image.convert("RGB"))
    img_resized = cv2.resize(img, (IMG_SIZE, IMG_SIZE))
    img_norm = img_resized / 255.0
    input_img = tf.convert_to_tensor(img_norm[None, ...], dtype=tf.float32)


    # Prediction
    preds = model.predict(input_img)
    class_idx = np.argmax(preds)
    confidence = np.max(preds)

    label = f"{CLASSES[class_idx]} ({confidence:.2%})"

    #Grad-CAM
    heatmap, _ = grad_cam(model, input_img, LAST_CONV_LAYER)
    heatmap = cv2.resize(heatmap, (IMG_SIZE, IMG_SIZE))
    heatmap = np.uint8(255 * heatmap)
    colored_heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)

    superimposed = cv2.addWeighted(
        colored_heatmap, 0.5,
        np.uint8(img_resized), 0.5,
        0
    )

    return Image.fromarray(superimposed), label

# ===== Gradio UI =====
demo = gr.Interface(
    fn=predict_xray,
    inputs=gr.Image(type="pil", label="Upload Chest X-ray"),
    outputs=[
        gr.Image(label="Grad-CAM Heatmap"),
        gr.Textbox(label="Prediction")
    ],
    title="🫁 Chest X-ray Pneumonia Classifier",
    description="DenseNet201-based model with Grad-CAM explainability",
)

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