Spaces:
Sleeping
Sleeping
| """ | |
| Grad-CAM: Gradient-weighted Class Activation Mapping | |
| Paper: https://arxiv.org/abs/1610.02391 | |
| Implemented from scratch using PyTorch tensor-level hooks. | |
| Uses tensor hooks instead of module backward hooks to avoid | |
| inplace operation conflicts in VGG, EfficientNet, MobileNet. | |
| """ | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| from typing import Optional | |
| class GradCAM: | |
| """ | |
| Grad-CAM implementation using forward hook + tensor-level gradient hook. | |
| The algorithm: | |
| 1. Forward pass → capture feature maps at target conv layer | |
| 2. Backward pass from class score → capture gradients via tensor hook | |
| 3. Global average pool the gradients → per-channel importance weights | |
| 4. Weighted combination of feature maps → class activation map | |
| 5. ReLU (we only care about features with positive influence) | |
| 6. Bilinear upsample to input size | |
| """ | |
| def __init__(self, model: torch.nn.Module, target_layer: torch.nn.Module): | |
| self.model = model | |
| self.target_layer = target_layer | |
| self.model.eval() | |
| self._feature_maps: Optional[torch.Tensor] = None | |
| self._gradients: Optional[torch.Tensor] = None | |
| # Forward hook only — no backward hook | |
| self._forward_hook = target_layer.register_forward_hook(self._save_feature_maps) | |
| def _save_feature_maps(self, module, input, output): | |
| # No .detach() — tensor must stay in graph so register_hook fires | |
| self._feature_maps = output | |
| def __call__( | |
| self, | |
| input_tensor: torch.Tensor, | |
| class_idx: Optional[int] = None, | |
| ) -> np.ndarray: | |
| self._gradients = None | |
| def save_grad(grad): | |
| self._gradients = grad.detach().clone() | |
| input_tensor = input_tensor.requires_grad_(True) | |
| logits = self.model(input_tensor) | |
| if class_idx is None: | |
| class_idx = logits.argmax(dim=1).item() | |
| # Tensor-level hook on the feature map | |
| hook_handle = self._feature_maps.register_hook(save_grad) | |
| self.model.zero_grad() | |
| logits[0, class_idx].backward() | |
| hook_handle.remove() | |
| feature_maps = self._feature_maps.detach() | |
| gradients = self._gradients | |
| weights = gradients.mean(dim=(2, 3), keepdim=True) | |
| cam = F.relu((weights * feature_maps).sum(dim=1, keepdim=True)) | |
| h, w = input_tensor.shape[2], input_tensor.shape[3] | |
| cam = F.interpolate(cam, size=(h, w), mode="bilinear", align_corners=False) | |
| cam = cam.squeeze().cpu().numpy() | |
| return self._normalize(cam) | |
| def _normalize(arr: np.ndarray) -> np.ndarray: | |
| min_val, max_val = arr.min(), arr.max() | |
| if max_val - min_val < 1e-8: | |
| return np.zeros_like(arr) | |
| return (arr - min_val) / (max_val - min_val) | |
| def remove_hooks(self): | |
| self._forward_hook.remove() | |
| class GradCAMPlusPlus(GradCAM): | |
| """ | |
| Grad-CAM++ — improved version with better localization for multiple instances. | |
| Paper: https://arxiv.org/abs/1710.11063 | |
| Uses second-order gradient approximation for per-pixel importance weights. | |
| Inherits tensor-level hook approach from GradCAM for VGG compatibility. | |
| """ | |
| def __call__( | |
| self, | |
| input_tensor: torch.Tensor, | |
| class_idx: Optional[int] = None, | |
| ) -> np.ndarray: | |
| self._gradients = None | |
| def save_grad(grad): | |
| self._gradients = grad.detach().clone() | |
| input_tensor = input_tensor.requires_grad_(True) | |
| logits = self.model(input_tensor) | |
| if class_idx is None: | |
| class_idx = logits.argmax(dim=1).item() | |
| hook_handle = self._feature_maps.register_hook(save_grad) | |
| self.model.zero_grad() | |
| logits[0, class_idx].backward() | |
| hook_handle.remove() | |
| feature_maps = self._feature_maps.detach() | |
| gradients = self._gradients | |
| # Grad-CAM++ second-order weight computation | |
| grads_power_2 = gradients ** 2 | |
| grads_power_3 = gradients ** 3 | |
| sum_activations = feature_maps.sum(dim=(2, 3), keepdim=True) | |
| denom = 2 * grads_power_2 + sum_activations * grads_power_3 | |
| denom = torch.where(denom != 0, denom, torch.ones_like(denom)) | |
| alpha = grads_power_2 / denom | |
| weights = (alpha * F.relu(gradients)).sum(dim=(2, 3), keepdim=True) | |
| cam = F.relu((weights * feature_maps).sum(dim=1, keepdim=True)) | |
| h, w = input_tensor.shape[2], input_tensor.shape[3] | |
| cam = F.interpolate(cam, size=(h, w), mode="bilinear", align_corners=False) | |
| cam = cam.squeeze().cpu().numpy() | |
| return self._normalize(cam) | |