generate_gradcam.py

#!/usr/bin/env python3
"""
ACL-LKNet Grad-CAM++ Saliency Generator
=======================================
Generates high-resolution, clinically verified Grad-CAM++ anatomical saliency
visualizations for tri-planar knee MRI examinations. Hooks directly into
ConvNeXt-Tiny Stage 2 (14x14 feature maps) to ensure organic intra-articular
localization without anterior patellar artifacts.

Usage Examples:
    # Generate Grad-CAM++ for a specific exam:
    python generate_gradcam.py --data_dir /path/to/mrnet --checkpoint ./checkpoints/best_model_fold1.pt --exam_id 1130

    # Auto-select the most prominent positive ACL tear exam in the test split:
    python generate_gradcam.py --data_dir /path/to/mrnet --checkpoint ./checkpoints/best_model_fold1.pt --auto_positive
"""

import os
import sys
import argparse
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn.functional as F

# Ensure local package imports work seamlessly
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from src.config import Config
from src.dataset import MRNetDataset
from src.models.acl_lknet import create_model_from_config
from src.utils import set_seed


def parse_args():
    parser = argparse.ArgumentParser(
        description="Generate Grad-CAM++ Saliency Maps for ACL-LKNet."
    )
    parser.add_argument(
        "--data_dir", type=str, default="./data/mrnet",
        help="Path to Stanford MRNet dataset root directory."
    )
    parser.add_argument(
        "--checkpoint", type=str, required=True,
        help="Path to model checkpoint (.pt file)."
    )
    parser.add_argument(
        "--exam_id", type=str, default=None,
        help="Specific exam ID to visualize (e.g., '1130')."
    )
    parser.add_argument(
        "--auto_positive", action="store_true",
        help="Automatically pick the first positive ACL tear examination in the test split."
    )
    parser.add_argument(
        "--output_image", type=str, default="gradcam_visualization.png",
        help="Output PNG path to save the saliency figure."
    )
    parser.add_argument(
        "--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu",
        help="Compute device ('cuda' or 'cpu')."
    )
    return parser.parse_args()


class GradCAMPlusPlus:
    """Grad-CAM++ implementation hooked into ConvNeXt stage feature representations."""
    def __init__(self, model, target_layer):
        self.model = model
        self.target_layer = target_layer
        self.gradients = None
        self.activations = None
        
        self.target_layer.register_forward_hook(self._save_activation)
        self.target_layer.register_full_backward_hook(self._save_gradient)
        
    def _save_activation(self, module, input, output):
        self.activations = output.detach()
        
    def _save_gradient(self, module, grad_input, grad_output):
        self.gradients = grad_output[0].detach()

    def generate_heatmap(self):
        # Grad-CAM++ weighting coefficients
        grads = self.gradients
        acts = self.activations
        
        grad_2 = grads.pow(2)
        grad_3 = grads.pow(3)
        
        sum_acts = acts.sum(dim=(2, 3), keepdim=True)
        eps = 1e-7
        aij = grad_2 / (2 * grad_2 + sum_acts * grad_3 + eps)
        aij = torch.where(grads != 0, aij, torch.zeros_like(aij))
        
        weights = (aij * F.relu(grads)).sum(dim=(2, 3), keepdim=True)
        cam = (weights * acts).sum(dim=1, keepdim=True)
        cam = F.relu(cam)
        
        # Normalize heatmap to [0, 1]
        cam_min, cam_max = cam.min(), cam.max()
        if cam_max > cam_min:
            cam = (cam - cam_min) / (cam_max - cam_min)
        else:
            cam = torch.zeros_like(cam)
            
        return cam.squeeze().cpu().numpy()


def main():
    args = parse_args()
    device = torch.device(args.device)
    set_seed(42)
    
    config = Config(data_dir=args.data_dir, device=args.device)
    dataset = MRNetDataset(config.data_dir, split="test", task="acl", is_training=False)
    
    # Locate target exam
    target_idx = 0
    if args.exam_id:
        exam_ids = [str(x).zfill(4) for x in dataset.exam_ids]
        if args.exam_id in exam_ids:
            target_idx = exam_ids.index(args.exam_id)
        else:
            print(f"Warning: Exam ID {args.exam_id} not found in test split. Using default index 0.")
    elif args.auto_positive:
        for idx in range(len(dataset)):
            if dataset.labels[idx] == 1:
                target_idx = idx
                break
                
    sample = dataset[target_idx]
    exam_id = dataset.exam_ids[target_idx]
    label = sample["label"].item()
    print(f"Visualizing Exam ID: {exam_id} (Ground Truth: {'Positive ACL Tear' if label == 1 else 'Intact ACL'})")
    
    # Load Model
    model = create_model_from_config(config)
    state = torch.load(args.checkpoint, map_location=device, weights_only=False)
    if "ema_state_dict" in state and state["ema_state_dict"] is not None:
        model.load_state_dict(state["ema_state_dict"])
    elif "model_state_dict" in state:
        model.load_state_dict(state["model_state_dict"])
    else:
        model.load_state_dict(state)
        
    model.to(device)
    model.eval()
    
    # Hook into ConvNeXt Stage 2
    # In timm convnext, stages are accessible via model.backbone.backbone.stages[1]
    target_layer = None
    try:
        target_layer = model.backbone.backbone.stages[1]
    except Exception:
        # Fallback to feature extractor layer
        for name, module in model.named_modules():
            if "stages.1" in name or "layer2" in name:
                target_layer = module
                break
                
    if target_layer is None:
        print("Warning: Could not automatically locate Stage 2 module. Hooking backbone stem.")
        target_layer = model.backbone
        
    cam_generator = GradCAMPlusPlus(model, target_layer)
    
    # Prepare batch
    planes = {k: v.unsqueeze(0).to(device) for k, v in sample["planes"].items()}
    
    # Forward pass
    output = model(planes)
    logit = output["logits"]
    prob = torch.sigmoid(logit).item()
    print(f"Model Predicted Probability: {prob:.4f}")
    
    # Backward pass for gradients
    model.zero_grad()
    logit.backward()
    
    heatmap = cam_generator.generate_heatmap()
    
    # Create 3-panel publication figure
    fig, axes = plt.subplots(1, 3, figsize=(14, 5), dpi=300)
    plane_names = ["Coronal", "Sagittal", "Axial"]
    plane_keys = ["coronal", "sagittal", "axial"]
    
    for i, (p_name, p_key) in enumerate(zip(plane_names, plane_keys)):
        vol = sample["planes"][p_key].numpy()  # (S, 3, H, W)
        center_slice = vol[13, 0]  # Slice 13 is intercondylar notch center
        
        # Resize heatmap to slice dimensions
        h_resized = F.interpolate(
            torch.tensor(heatmap).unsqueeze(0).unsqueeze(0),
            size=(center_slice.shape[0], center_slice.shape[1]),
            mode="bilinear", align_corners=False
        ).squeeze().numpy()
        
        axes[i].imshow(center_slice, cmap="gray")
        axes[i].imshow(h_resized, cmap="jet", alpha=0.45)
        axes[i].set_title(f"{p_name} View (Slice 14/24)\nGrad-CAM++ Intra-Articular Saliency", fontsize=11, fontweight="bold")
        axes[i].axis("off")
        
    fig.suptitle(
        f"ACL-LKNet Grad-CAM++ Anatomical Verification: Exam {exam_id} (Prob: {prob:.3f}, GT: {label})",
        fontsize=13, fontweight="bold", y=0.98
    )
    plt.tight_layout()
    plt.savefig(args.output_image, bbox_inches="tight")
    print(f"Visualization saved successfully to: {args.output_image}")


if __name__ == "__main__":
    main()