File size: 2,716 Bytes
37fbec9 | 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 | """
explainability/gradcam.py
-------------------------
Implements Grad-CAM for ViT using pytorch-grad-cam.
"""
import numpy as np
import torch
import cv2
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image
def reshape_transform(tensor, height=14, width=14):
"""
Reshape transform specifically for ViT models.
Removes the CLS token and reshapes the sequence of patch tokens
back into a 2D spatial feature map.
"""
# ViT output tensor shape: (B, num_tokens, D)
result = tensor[:, 1:, :].reshape(tensor.size(0),
height, width, tensor.size(2))
# Bring the channels to the first dimension like in CNNs
result = result.transpose(2, 3).transpose(1, 2)
return result
def generate_gradcam_heatmap(
model,
image_tensor: torch.Tensor,
target_class: int,
original_image: np.ndarray = None,
image_size: int = 224,
) -> np.ndarray:
"""
Generate a Grad-CAM heatmap for a given target class.
Args:
model: Trained ChestViT model
image_tensor: Normalized (1, 3, 224, 224) torch tensor
target_class: Index of the disease to explain (0-13)
original_image: Original RGB image in [0, 255] for overlay
Returns:
heatmap_overlay: (224, 224, 3) RGB image with Grad-CAM overlay
"""
model.eval()
# The target layer in Hugging Face ViT models is typically the last norm layer
# before the classification head, or the final layer of the encoder.
target_layers = [model.vit.encoder.layer[-1].layernorm_before]
# Initialize GradCAM
cam = GradCAM(model=model,
target_layers=target_layers,
reshape_transform=reshape_transform)
# Generate the heatmap for the target class
# target_category can be None to target the highest scoring class,
# but we specify it for exact disease visualization
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
targets = [ClassifierOutputTarget(target_class)]
# You can also pass aug_smooth=True and eigen_smooth=True
grayscale_cam = cam(input_tensor=image_tensor,
targets=targets)
grayscale_cam = grayscale_cam[0, :]
if original_image is not None:
# pytorch-grad-cam expects image float32 in [0, 1]
rgb_img = np.float32(original_image) / 255
# Resize to match output if needed
if rgb_img.shape[:2] != (image_size, image_size):
rgb_img = cv2.resize(rgb_img, (image_size, image_size))
visualization = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True)
return visualization
return grayscale_cam
|