""" gradcam.py — Grad-CAM implementation on the ResNet-50 CNN branch. Computes gradient-weighted class activation maps for the wildfire risk output. Highlights spatial regions that most influenced the model's fire risk prediction. Output: Semi-transparent heatmap overlay on the original satellite image using a jet colormap. Red/yellow = high fire risk regions. """ import logging from typing import Optional, Tuple import cv2 import numpy as np import torch import torch.nn.functional as F from src.training.config import PATCH_SIZE, DEVICE logger = logging.getLogger(__name__) class GradCAM: """ Grad-CAM visualization for the CNN branch of the MTL model. Hooks into the last convolutional layer of ResNet-50 (layer4) to compute gradient-weighted class activation maps. """ def __init__(self, model, target_layer=None): """ Args: model: MultiTaskFusionModel instance. target_layer: Specific layer to hook into. Defaults to the last conv layer of ResNet-50 (layer4). """ self.model = model self.feature_maps = None self.gradients = None # Hook into the target layer if target_layer is None: target_layer = model.cnn_branch.layer4 self._register_hooks(target_layer) logger.info("GradCAM initialized — hooked into target layer.") def _register_hooks(self, target_layer): """Register forward and backward hooks on the target layer.""" def forward_hook(module, input, output): self.feature_maps = output.detach() def backward_hook(module, grad_input, grad_output): self.gradients = grad_output[0].detach() target_layer.register_forward_hook(forward_hook) target_layer.register_full_backward_hook(backward_hook) def generate( self, image: torch.Tensor, timeseries: torch.Tensor, target_size: Tuple[int, int] = (PATCH_SIZE, PATCH_SIZE), ) -> np.ndarray: """ Generate Grad-CAM heatmap for wildfire risk prediction. Args: image: (1, 4, 128, 128) single image tensor. timeseries: (1, 7, 6) single time-series tensor. target_size: Output heatmap size (H, W). Returns: Heatmap array of shape (H, W) with values in [0, 1]. """ self.model.eval() # Need gradients for Grad-CAM image_input = image.clone().requires_grad_(True) # Forward pass output = self.model(image_input, timeseries) risk_map = output["risk_map"] # (1, 128, 128) # Backward pass — use the mean risk score as the target self.model.zero_grad() target = risk_map.mean() target.backward(retain_graph=True) if self.gradients is None or self.feature_maps is None: logger.warning("Grad-CAM: No gradients captured. Returning uniform map.") return np.ones(target_size, dtype=np.float32) * 0.5 # Compute channel-wise weights: global average pooling of gradients weights = torch.mean(self.gradients, dim=(2, 3)) # (1, C) # Weighted combination of feature maps cam = torch.zeros( self.feature_maps.shape[2:], device=self.feature_maps.device ) for i, w in enumerate(weights[0]): cam += w * self.feature_maps[0, i] # ReLU: keep only positive contributions cam = F.relu(cam) # Normalize to [0, 1] if cam.max() > 0: cam = cam / cam.max() # Resize to target size cam_np = cam.cpu().numpy() cam_resized = cv2.resize(cam_np, target_size, interpolation=cv2.INTER_LINEAR) return cam_resized def generate_overlay( self, image: torch.Tensor, timeseries: torch.Tensor, alpha: float = 0.5, ) -> np.ndarray: """ Generate Grad-CAM overlay on the original satellite image. Creates a composite image with: - Original RGB satellite image as background - Grad-CAM heatmap overlaid with jet colormap - Alpha blending at 50% opacity Args: image: (1, 4, 128, 128) image tensor. timeseries: (1, 7, 6) time-series tensor. alpha: Blending factor. Returns: BGR overlay image of shape (128, 128, 3), uint8. """ # Generate Grad-CAM heatmap heatmap = self.generate(image, timeseries) # Convert heatmap to color using jet colormap heatmap_uint8 = np.uint8(255 * heatmap) heatmap_color = cv2.applyColorMap(heatmap_uint8, cv2.COLORMAP_JET) # Extract RGB from original image and convert to uint8 img_np = image[0, :3].detach().cpu().numpy() # (3, H, W) img_np = np.transpose(img_np, (1, 2, 0)) # (H, W, 3) — RGB # Normalize to [0, 255] img_min, img_max = img_np.min(), img_np.max() if img_max - img_min > 0: img_np = (img_np - img_min) / (img_max - img_min) img_uint8 = np.uint8(255 * img_np) # Convert RGB to BGR for OpenCV img_bgr = cv2.cvtColor(img_uint8, cv2.COLOR_RGB2BGR) # Alpha-blend overlay overlay = cv2.addWeighted(img_bgr, 1 - alpha, heatmap_color, alpha, 0) return overlay def to_base64( self, image: torch.Tensor, timeseries: torch.Tensor, ) -> str: """ Generate Grad-CAM overlay and encode as base64 PNG. Args: image: (1, 4, 128, 128) image tensor. timeseries: (1, 7, 6) time-series tensor. Returns: Base64-encoded PNG string. """ import base64 overlay = self.generate_overlay(image, timeseries) # Encode as PNG success, buffer = cv2.imencode(".png", overlay) if not success: logger.error("Failed to encode Grad-CAM overlay as PNG.") return "" return base64.b64encode(buffer).decode("utf-8") def generate_gradcam_for_prediction( model, image: torch.Tensor, timeseries: torch.Tensor, ) -> Tuple[np.ndarray, np.ndarray, str]: """ Convenience function to generate all Grad-CAM outputs. Returns: Tuple of (raw_heatmap, overlay_image, base64_string). """ gradcam = GradCAM(model) heatmap = gradcam.generate(image, timeseries) overlay = gradcam.generate_overlay(image, timeseries) b64 = gradcam.to_base64(image, timeseries) return heatmap, overlay, b64 if __name__ == "__main__": logging.basicConfig(level=logging.INFO) from src.models.fusion_model import MultiTaskFusionModel model = MultiTaskFusionModel(pretrained_cnn=False) model.eval() img = torch.randn(1, 4, 128, 128) ts = torch.randn(1, 7, 6) heatmap, overlay, b64 = generate_gradcam_for_prediction(model, img, ts) print(f"Heatmap shape: {heatmap.shape}, range: [{heatmap.min():.3f}, {heatmap.max():.3f}]") print(f"Overlay shape: {overlay.shape}") print(f"Base64 length: {len(b64)} chars")