app.py

import gradio as gr
import numpy as np
import tensorflow as tf
from tensorflow.keras.preprocessing.image import array_to_img
import matplotlib.cm as cm
import matplotlib.pyplot as plt
import os
from PIL import Image # Ensure PIL is imported

# --- Global Variables ---
# Define model path assuming it's in the same directory or a 'models' subdirectory
# For simplicity, we'll assume 'best_model.h5' is uploaded alongside app.py
MODEL_PATH = 'best_model.h5'

# Ensure these variables are correctly defined and not defaulted
INPUT_SHAPE = (224, 224, 3)
NUM_CLASSES = 4
class_labels = ['Glioma', 'Meningioma', 'Normal', 'Pituitary'] # Use the updated class labels
last_conv_layer_name = 'conv5_block16_2_conv' # From previous cell's finding

# --- Load Model (outside predict function for efficiency) ---
try:
    best_model = tf.keras.models.load_model(MODEL_PATH)
    print(f"Model loaded successfully from: {MODEL_PATH}")
except Exception as e:
    print(f"Error loading model from {MODEL_PATH}: {e}")
    best_model = None # Set to None if loading fails

# --- Explainability Helper Functions ---
# Helper function to display Grad-CAM
def display_gradcam(img_pil, heatmap, alpha=0.4):
    img_array = tf.keras.preprocessing.image.img_to_array(img_pil)

    # Rescale heatmap to a range 0-255
    heatmap = np.uint8(255 * heatmap)

    # Use jet colormap to colorize heatmap
    jet = plt.colormaps["jet"]
    jet_colors = jet(np.arange(256))
    jet_heatmap = jet_colors[heatmap]

    # Create an image with RGB colorized heatmap
    jet_heatmap_pil = Image.fromarray(np.uint8(jet_heatmap * 255))
    jet_heatmap_pil = jet_heatmap_pil.resize((img_pil.width, img_pil.height))
    jet_heatmap_array = tf.keras.preprocessing.image.img_to_array(jet_heatmap_pil)
    jet_heatmap_array = jet_heatmap_array[:, :, :3] # Convert RGBA to RGB

    # Superimpose the heatmap on original image
    superimposed_img_array = jet_heatmap_array * alpha + img_array
    superimposed_img_array = np.clip(superimposed_img_array, 0, 255).astype(np.uint8)
    superimposed_img_pil = Image.fromarray(superimposed_img_array)

    return superimposed_img_pil, Image.fromarray(np.uint8(jet_heatmap_array))

# Helper function to get Grad-CAM heatmap
def make_gradcam_heatmap(img_array, model, last_conv_layer_name, pred_index=None):
    grad_model = tf.keras.models.Model(
        [model.inputs], [model.get_layer(last_conv_layer_name).output, model.output]
    )

    with tf.GradientTape() as tape:
        last_conv_layer_output, preds = grad_model([img_array])
        if pred_index is None:
            pred_index = tf.argmax(preds[0])
        class_channel = preds[:, pred_index]

    grads = tape.gradient(class_channel, last_conv_layer_output)
    if grads is None:
        return np.zeros(last_conv_layer_output.shape[1:-1])

    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
    last_conv_layer_output = last_conv_layer_output[0]
    heatmap = last_conv_layer_output @ pooled_grads[..., tf.newaxis]
    heatmap = tf.squeeze(heatmap)
    heatmap = tf.maximum(heatmap, 0) / (tf.reduce_max(heatmap) + 1e-7)
    return heatmap.numpy()

# Grad-CAM++
def make_gradcam_plus_plus_heatmap(img_array, model, last_conv_layer_name, pred_index=None):
    grad_model = tf.keras.models.Model(
        model.inputs, [model.get_layer(last_conv_layer_name).output, model.output]
    )

    with tf.GradientTape(persistent=True) as tape:
        last_conv_layer_output, preds = grad_model([img_array])
        if pred_index is None:
            pred_index = tf.argmax(preds[0])
        class_channel = preds[:, pred_index]

        first_grad_tensor = tape.gradient(class_channel, last_conv_layer_output)
        if first_grad_tensor is None:
            del tape
            return np.zeros(last_conv_layer_output.shape[1:-1])

        first_grad_tensor = first_grad_tensor[0]
        second_grad_tensor = tape.gradient(first_grad_tensor, last_conv_layer_output)[0]
        third_grad_tensor = tape.gradient(second_grad_tensor, last_conv_layer_output)[0]

    del tape

    last_conv_layer_output_nobatch = last_conv_layer_output[0]
    first_grad_nobatch = first_grad_tensor[0]
    second_grad_nobatch = second_grad_tensor[0]
    third_grad_nobatch = third_grad_tensor[0]

    pooled_second_grad = tf.reduce_mean(second_grad_nobatch, axis=(0, 1))
    pooled_third_grad = tf.reduce_mean(third_grad_nobatch, axis=(0, 1))

    sum_activations_per_channel = tf.reduce_sum(last_conv_layer_output_nobatch, axis=(0, 1))

    eps = 1e-7
    alpha_num = pooled_second_grad
    alpha_den = pooled_second_grad * 2 + pooled_third_grad * sum_activations_per_channel
    alpha_den = tf.where(tf.equal(alpha_den, 0.0), eps, alpha_den) # Handle division by zero for tensors
    alphas = alpha_num / alpha_den

    alphas = alphas[tf.newaxis, tf.newaxis, :] # (1, 1, C)

    weights = tf.maximum(first_grad_nobatch, 0.0)

    deep_insights = alphas * weights * last_conv_layer_output_nobatch
    heatmap = tf.reduce_sum(deep_insights, axis=-1) # (H, W)

    heatmap = tf.maximum(heatmap, 0) / (tf.reduce_max(heatmap) + 1e-7)
    return heatmap.numpy()

# LayerCAM
def make_layercam_heatmap(img_array, model, last_conv_layer_name, pred_index=None):
    grad_model = tf.keras.models.Model(
        model.inputs, [model.get_layer(last_conv_layer_name).output, model.output]
    )

    with tf.GradientTape() as tape:
        last_conv_layer_output, preds = grad_model([img_array])
        if pred_index is None:
            pred_index = tf.argmax(preds[0])
        class_channel = preds[:, pred_index]

    grads = tape.gradient(class_channel, last_conv_layer_output)
    if grads is None:
        return np.zeros(last_conv_layer_output.shape[1:-1])

    heatmap = tf.reduce_sum(tf.abs(grads[0][0]) * last_conv_layer_output[0], axis=-1)

    heatmap = tf.maximum(heatmap, 0) / (tf.reduce_max(heatmap) + 1e-7)
    return heatmap.numpy()

# ScoreCAM
def make_scorecam_heatmap(img_array, model, last_conv_layer_name, pred_index=None):
    intermediate_model = tf.keras.models.Model(inputs=model.inputs, outputs=model.get_layer(last_conv_layer_name).output)

    activations_batch = intermediate_model.predict(img_array, verbose=0)
    activations = activations_batch[0]

    original_h, original_w = img_array.shape[1:3]

    upsampled_activations = []
    for i in range(activations.shape[-1]):
        channel_activation = activations[:, :, i]
        upsampled_channel = tf.image.resize(tf.expand_dims(channel_activation, -1),
                                            (original_h, original_w),
                                            method='bilinear').numpy()[:, :, 0]
        upsampled_activations.append(upsampled_channel)

    upsampled_activations_stacked = np.stack(upsampled_activations, axis=-1)

    if pred_index is None:
        preds = model.predict(img_array, verbose=0)
        pred_index = tf.argmax(preds[0])

    masked_images_list = []
    for i in range(activations.shape[-1]):
        mask = upsampled_activations_stacked[:, :, i]
        normalized_mask = mask / (np.max(mask) + 1e-7) if np.max(mask) > 0 else mask
        masked_img_unprocessed = img_array[0] * normalized_mask[:, :, np.newaxis]
        masked_img_processed = np.expand_dims(masked_img_unprocessed, axis=0)
        masked_images_list.append(masked_img_processed)

    if not masked_images_list:
        return np.zeros((original_h, original_w))

    batched_masked_images = np.vstack(masked_images_list)

    preds_masked = model.predict(batched_masked_images, verbose=0)
    scores_for_target_class = preds_masked[:, pred_index]

    min_score = np.min(scores_for_target_class)
    max_score = np.max(scores_for_target_class)

    if (max_score - min_score) == 0:
        weights = np.zeros_like(scores_for_target_class)
    else:
        weights = (scores_for_target_class - min_score) / (max_score - min_score + 1e-7)

    heatmap_scorecam = np.sum(upsampled_activations_stacked * weights[np.newaxis, np.newaxis, :], axis=-1)

    heatmap_scorecam = np.maximum(heatmap_scorecam, 0)
    if np.max(heatmap_scorecam) > 0:
        heatmap_scorecam /= (np.max(heatmap_scorecam) + 1e-7)
    else:
        heatmap_scorecam = np.zeros_like(heatmap_scorecam)

    return heatmap_scorecam


# --- Prediction and Explainability Function for Gradio ---
def predict_and_explain(image_pil):
    if best_model is None:
        return "Error: Model not loaded.", "N/A", None, None, None, None, None, None, None, None, None

    # Preprocess the image
    img_resized = image_pil.resize((INPUT_SHAPE[0], INPUT_SHAPE[1]))
    img_array = tf.keras.preprocessing.image.img_to_array(img_resized)
    img_array = np.expand_dims(img_array, axis=0) # Add batch dimension
    processed_img = tf.keras.applications.densenet.preprocess_input(img_array)

    # Make prediction
    preds = best_model.predict(processed_img, verbose=0)
    predicted_class_idx = np.argmax(preds[0])
    predicted_class_name = class_labels[predicted_class_idx]
    confidence = preds[0][predicted_class_idx] * 100

    # Generate heatmaps
    grad_cam_heatmap = make_gradcam_heatmap(processed_img, best_model, last_conv_layer_name, pred_index=predicted_class_idx)
    grad_cam_plus_plus_heatmap = make_gradcam_plus_plus_heatmap(processed_img, best_model, last_conv_layer_name, pred_index=predicted_class_idx)
    layercam_heatmap = make_layercam_heatmap(processed_img, best_model, last_conv_layer_name, pred_index=predicted_class_idx)
    scorecam_heatmap = make_scorecam_heatmap(processed_img, best_model, last_conv_layer_name, pred_index=predicted_class_idx)

    # Superimpose heatmaps
    original_img_display = img_resized # Use the resized PIL image

    superimposed_grad_cam, grad_cam_heatmap_img = display_gradcam(img_resized, grad_cam_heatmap)
    superimposed_grad_cam_plus_plus, grad_cam_plus_plus_heatmap_img = display_gradcam(img_resized, grad_cam_plus_plus_heatmap)
    superimposed_layercam, layercam_heatmap_img = display_gradcam(img_resized, layercam_heatmap)
    superimposed_scorecam, scorecam_heatmap_img = display_gradcam(img_resized, scorecam_heatmap)

    return (
        predicted_class_name,
        f"{confidence:.2f}%",
        original_img_display,
        grad_cam_heatmap_img,
        superimposed_grad_cam,
        grad_cam_plus_plus_heatmap_img,
        superimposed_grad_cam_plus_plus,
        layercam_heatmap_img,
        superimposed_layercam,
        scorecam_heatmap_img,
        superimposed_scorecam
    )

# --- Set up Gradio Interface ---
if best_model is not None:
    interface = gr.Interface(
        fn=predict_and_explain,
        inputs=gr.Image(type="pil", label="Upload MRI Scan"),
        outputs=[
            gr.Label(label="Predicted Class"),
            gr.Textbox(label="Confidence"),
            gr.Image(label="Original Image"),
            gr.Image(label="Grad-CAM Heatmap"),
            gr.Image(label="Grad-CAM Superimposed"),
            gr.Image(label="Grad-CAM++ Heatmap"),
            gr.Image(label="Grad-CAM++ Superimposed"),
            gr.Image(label="LayerCAM Heatmap"),
            gr.Image(label="LayerCAM Superimposed"),
            gr.Image(label="ScoreCAM Heatmap"),
            gr.Image(label="ScoreCAM Superimposed")
        ],
        title="Brain Tumor Classification with Explainability",
        description="Upload an MRI image to classify brain tumor types and visualize model's focus using Grad-CAM, Grad-CAM++, LayerCAM, and ScoreCAM."
    )

    # To run on Hugging Face Spaces, do not use debug=True or share=True
    # interface.launch(debug=True, share=True)
    interface.launch()
else:
    print("Gradio interface could not be launched because the model failed to load.")