Spaces:
Sleeping
Sleeping
File size: 4,458 Bytes
c04bc97 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
from PIL import Image
class GradCAM:
"""
Custom hook-based Grad-CAM implementation for PyTorch CNNs (e.g., ResNet-18).
Attributes predictions to image regions by weighting convolutional activations with backpropagated gradients.
"""
def __init__(self, model: nn.Module, target_layer: nn.Module):
self.model = model
self.target_layer = target_layer
self.activations = None
self.gradients = None
# Register hooks
self.forward_hook_handle = target_layer.register_forward_hook(self._forward_hook)
# register_full_backward_hook is the modern PyTorch standard replacing register_backward_hook
self.backward_hook_handle = target_layer.register_full_backward_hook(self._backward_hook)
def _forward_hook(self, module, input, output):
self.activations = output.detach()
def _backward_hook(self, module, grad_input, grad_output):
# grad_output is a tuple containing gradients with respect to output feature map
self.gradients = grad_output[0].detach()
def generate_heatmap(self, input_tensor: torch.Tensor, target_class: int = None):
"""
Generates the 2D Grad-CAM heatmap for a given input tensor and target class.
input_tensor: shape (1, C, H, W)
"""
self.model.eval()
self.model.zero_grad()
# Forward pass
output = self.model(input_tensor)
# If target class is not provided, use the class with the highest probability
if target_class is None:
target_class = torch.argmax(output, dim=1).item()
score = output[0, target_class]
# Backward pass to calculate gradients
score.backward()
# Extract activations and gradients
# Shape of activations & gradients: (1, channels, h_feat, w_feat)
acts = self.activations
grads = self.gradients
if acts is None or grads is None:
raise RuntimeError("Hook activations or gradients are empty. Verify forward/backward pass execution.")
# Global Average Pooling of gradients to compute channel weights
weights = torch.mean(grads, dim=(2, 3), keepdim=True) # Shape: (1, channels, 1, 1)
# Weighted sum of feature map activations
cam = torch.sum(weights * acts, dim=1, keepdim=True) # Shape: (1, 1, h_feat, w_feat)
# Apply ReLU to retain only positive influences on target class
cam = torch.clamp(cam, min=0)
# 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)
# Remove batch and channel dims, convert to numpy
heatmap = cam.squeeze().cpu().numpy()
return heatmap, target_class
def close(self):
"""Removes the hooks to prevent memory leaks."""
self.forward_hook_handle.remove()
self.backward_hook_handle.remove()
def overlay_gradcam_on_image(original_img: np.ndarray, heatmap: np.ndarray, alpha: float = 0.5):
"""
Overlays a Grad-CAM heatmap onto a raw image using matplotlib colormap.
original_img: numpy array of shape (H, W, 3), values in range [0, 1] or [0, 255]
heatmap: 2D numpy array of shape (H, W), values in range [0, 1]
"""
# Ensure original_img is in range [0, 255] and uint8
if original_img.dtype != np.uint8:
if original_img.max() <= 1.0:
original_img = (original_img * 255).astype(np.uint8)
else:
original_img = original_img.astype(np.uint8)
# Resize heatmap to match image dimensions
pil_heat = Image.fromarray((heatmap * 255).astype(np.uint8))
pil_heat = pil_heat.resize((original_img.shape[1], original_img.shape[0]), Image.Resampling.BILINEAR)
heatmap_resized = np.array(pil_heat) / 255.0
# Get colormap jet
cmap = plt.get_cmap('jet')
color_heatmap = cmap(heatmap_resized)[:, :, :3] # Drop alpha channel
color_heatmap = (color_heatmap * 255).astype(np.uint8)
# Blend images
blended = (alpha * color_heatmap + (1 - alpha) * original_img).astype(np.uint8)
return blended, color_heatmap
|