Explainable Computer Vision: Implementing Grad-CAM Over Deep CNNs for Retinal Lesion Localization

The Black-Box Dilemma in Diagnostic Imaging

In medical imaging domains—such as our work developing RetinaScan AI and the IDRP retinal screening portals—predictive accuracy alone is insufficient. When an ophthalmologist or clinical screener evaluates an automated system that classifies a fundus scan as Moderate Non-Proliferative Diabetic Retinopathy (NPDR), their immediate question is: "What visual biomarkers led to this prediction?"

If the model classified the image based on genuine pathology (such as cotton-wool spots, blot hemorrhages, or microaneurysms), the prediction has clinical credibility. If, however, the model learned a spurious correlation—such as an artifact border on the camera lens, illumination falloff, or clinic watermark stamps—the system is dangerously flawed despite having high benchmark test accuracy.

How Grad-CAM Computes Spatial Saliency

Gradient-weighted Class Activation Mapping (Grad-CAM) addresses this opacity by using the gradients of the target concept (in our case, the score for class \(c\)) flowing into the final convolutional layer of a CNN.

[ Input Image (224x224x3) ]
             │
             ▼
┌─────────────────────────────────────────────────────────────┐
│ Convolutional Feature Extractor (e.g. EfficientNet-B4)      │
│ Feature maps preserve spatial geometry of anatomical lesions│
└────────────────────────────┬────────────────────────────────┘
                             │ Final Conv Layer Activations: A^k
                             ▼
┌─────────────────────────────────────────────────────────────┐
│ Global Average Pooling & Linear Classification Head         │
│ Computes Class Score: y^c                                   │
└────────────────────────────┬────────────────────────────────┘
                             │
                             ▼  [ Backpropagate Gradient: dy^c / dA^k ]
┌─────────────────────────────────────────────────────────────┐
│ Neuron Importance Weights: lpha_k^c                       │
│ lpha_k^c = (1/Z) * \sum_i \sum_j (dy^c / dA_{i,j}^k)       │
└────────────────────────────┬────────────────────────────────┘
                             │
                             ▼  [ Linear Combination & ReLU ]
┌─────────────────────────────────────────────────────────────┐
│ Localization Heatmap: L_{Grad-CAM}^c = ReLU(\sum_k lpha_k^c A^k)│
│ Positive gradients highlight features that increase score c │
└─────────────────────────────────────────────────────────────┘

Full PyTorch Grad-CAM Implementation

Here is the clean, self-contained PyTorch implementation registering forward and backward hooks on the target convolutional layer:

import torch
import torch.nn.functional as F
import numpy as np
import cv2

class GradCAM:
    def __init__(self, model: torch.nn.Module, target_layer: torch.nn.Module):
        self.model = model
        self.target_layer = target_layer
        self.gradients = None
        self.activations = None
        
        # Register PyTorch execution hooks
        self.target_layer.register_forward_hook(self._save_activations)
        self.target_layer.register_full_backward_hook(self._save_gradients)

    def _save_activations(self, module, input, output):
        self.activations = output.detach()

    def _save_gradients(self, module, grad_input, grad_output):
        self.gradients = grad_output[0].detach()

    def generate_heatmap(self, input_tensor: torch.Tensor, class_idx: int = None) -> np.ndarray:
        self.model.zero_grad()
        output = self.model(input_tensor)

        if class_idx is None:
            class_idx = torch.argmax(output, dim=1).item()

        score = output[0, class_idx]
        score.backward(retain_graph=True)

        # 1. Global Average Pooling of gradients: alpha_k
        weights = torch.mean(self.gradients, dim=(2, 3), keepdim=True)

        # 2. Linear combination of weighted activations
        cam = torch.sum(weights * self.activations, dim=1).squeeze()

        # 3. Apply ReLU (we only care about features with positive impact)
        cam = F.relu(cam)

        # 4. Normalize between 0.0 and 1.0
        cam_np = cam.cpu().numpy()
        cam_min, cam_max = np.min(cam_np), np.max(cam_np)
        if cam_max > cam_min:
            cam_norm = (cam_np - cam_min) / (cam_max - cam_min)
        else:
            cam_norm = np.zeros_like(cam_np)

        return cam_norm

def overlay_heatmap_on_fundus(rgb_image: np.ndarray, heatmap: np.ndarray, alpha: float = 0.45) -> np.ndarray:
    '''Resizes heatmap to original resolution and blends using OpenCV colormap.'''
    h, w = rgb_image.shape[:2]
    # Resize heatmap to match original fundus dimensions
    resized_cam = cv2.resize(heatmap, (w, h))
    heatmap_uint8 = np.uint8(255 * resized_cam)
    
    # Apply JET or TURBO colormap for thermal biomarker visibility
    colored_heatmap = cv2.applyColorMap(heatmap_uint8, cv2.COLORMAP_JET)
    colored_heatmap = cv2.cvtColor(colored_heatmap, cv2.COLOR_BGR2RGB)

    blended = np.uint8(alpha * colored_heatmap + (1.0 - alpha) * rgb_image)
    return blended

Validating Localization Against Clinical Ground Truth

Grad-CAM heatmaps must be verified against actual clinical annotations. In our testing protocols:

1. Pixel Intersection over Union (IoU): We measure bounding box overlap between high-activation thermal peaks (top 20% heatmap intensity) and ophthalmologist-marked microaneurysms.
2. Sanity Checks (Model Parameter Randomization): We perform the Adebayo sanity test—randomizing model weights layer-by-layer to confirm that the heatmap collapses. If a saliency map remains unchanged after randomizing weights, it is acting as an edge detector rather than reflecting learned model features.