gradcam.py

"""
Grad-CAM Heatmap Generator for EfficientNetB3 Malaria Detector

Generates class activation maps showing which regions of the blood smear
the model focuses on when making predictions.
"""

import os
import time
import glob
import numpy as np
import cv2
import tensorflow as tf


# Cache the gradient model per base model to avoid rebuilding each request
_grad_model_cache = {}


def _get_grad_model(model, last_conv_layer_name="top_conv"):
    """Build a model that outputs the conv layer activations + final prediction."""
    if id(model) in _grad_model_cache:
        return _grad_model_cache[id(model)]

    # EfficientNetB3's last conv layer is named 'top_conv'
    last_conv_layer = model.get_layer(last_conv_layer_name)
    grad_model = tf.keras.Model(
        inputs=model.input,
        outputs=[last_conv_layer.output, model.output],
    )
    _grad_model_cache[id(model)] = grad_model
    return grad_model


def generate_gradcam(model, img_array, output_path, original_image_path,
                     last_conv_layer_name="top_conv", alpha=0.4):
    """
    Generate a Grad-CAM heatmap overlay and save to disk.

    Args:
        model: Loaded Keras model.
        img_array: Preprocessed image array (1, H, W, 3) in [0, 255].
        output_path: Path to save the heatmap overlay image.
        original_image_path: Path to the original uploaded image for overlay.
        last_conv_layer_name: Name of the last convolutional layer.
        alpha: Opacity of the heatmap overlay.

    Returns:
        True if successful, False otherwise.
    """
    try:
        grad_model = _get_grad_model(model, last_conv_layer_name)

        # Compute gradient of prediction w.r.t. conv layer output
        with tf.GradientTape() as tape:
            conv_outputs, predictions = grad_model(img_array)
            # Binary classification: single sigmoid output
            loss = predictions[:, 0]

        grads = tape.gradient(loss, conv_outputs)

        # Global average pooling of gradients
        pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))

        # Weight the conv outputs by the pooled gradients
        conv_outputs = conv_outputs[0]
        heatmap = conv_outputs @ pooled_grads[..., tf.newaxis]
        heatmap = tf.squeeze(heatmap)

        # ReLU and normalize
        heatmap = tf.maximum(heatmap, 0) / (tf.math.reduce_max(heatmap) + 1e-8)
        heatmap = heatmap.numpy()

        # Load original image for overlay
        original = cv2.imread(original_image_path)
        if original is None:
            return False

        h, w = original.shape[:2]

        # Resize heatmap to match original image
        heatmap_resized = cv2.resize(heatmap, (w, h))
        heatmap_uint8 = np.uint8(255 * heatmap_resized)
        heatmap_colored = cv2.applyColorMap(heatmap_uint8, cv2.COLORMAP_JET)

        # Overlay
        overlay = cv2.addWeighted(original, 1 - alpha, heatmap_colored, alpha, 0)

        # Save
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        cv2.imwrite(output_path, overlay)
        return True

    except Exception as e:
        print(f"Grad-CAM generation failed: {e}")
        return False


def cleanup_old_gradcam(gradcam_dir, max_age_seconds=3600):
    """Remove Grad-CAM images older than max_age_seconds (default: 1 hour)."""
    if not os.path.exists(gradcam_dir):
        return
    now = time.time()
    for filepath in glob.glob(os.path.join(gradcam_dir, "*.png")):
        try:
            if now - os.path.getmtime(filepath) > max_age_seconds:
                os.remove(filepath)
        except OSError:
            pass