| """Grad-CAM heatmaps for DenseNet-121, mapped back onto the original image.""" |
| from __future__ import annotations |
|
|
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| from PIL import Image |
|
|
| from .preprocess import to_tensor |
|
|
|
|
| def gradcam(model, img: Image.Image, class_idx: int) -> np.ndarray: |
| """Return a heatmap in [0, 1] with the same size as ``img``.""" |
| device = next(model.parameters()).device |
| x = to_tensor(img).to(device) |
| with torch.enable_grad(): |
| feats = model.features(x) |
| feats.retain_grad() |
| out = F.relu(feats) |
| logits = model.classifier(torch.flatten(F.adaptive_avg_pool2d(out, 1), 1)) |
| model.zero_grad(set_to_none=True) |
| logits[0, class_idx].backward() |
| w = feats.grad.mean(dim=(2, 3), keepdim=True) |
| cam = F.relu((w * feats).sum(1, keepdim=True)) |
| cam = F.interpolate(cam, size=(448, 448), mode="bilinear", align_corners=False)[0, 0] |
| cam = cam.detach().cpu().numpy() |
|
|
| |
| W, H = img.size |
| m = max(W, H) |
| left, top = int((m - W) / 2) / m * 448, int((m - H) / 2) / m * 448 |
| cam = cam[int(round(top)):int(round(448 - top)), int(round(left)):int(round(448 - left))] |
| cam = np.asarray(Image.fromarray(cam.astype(np.float32), mode="F").resize((W, H), Image.BILINEAR)) |
| rng = cam.max() - cam.min() |
| return (cam - cam.min()) / rng if rng > 0 else np.zeros_like(cam) |
|
|
|
|
| def overlay(img: Image.Image, cam: np.ndarray, alpha: float = 0.45) -> Image.Image: |
| """Blend a heatmap (red = most influential) onto the image.""" |
| base = np.asarray(img.convert("RGB"), dtype=np.float32) |
| c = np.clip(cam, 0, 1)[..., None] |
| heat = np.concatenate([np.clip(1.5 * c, 0, 1), np.clip(1.5 * c - 0.5, 0, 1) * 0.6, 0.2 * (1 - c)], axis=-1) * 255 |
| a = alpha * c |
| return Image.fromarray(np.clip(base * (1 - a) + heat * a, 0, 255).astype(np.uint8)) |
|
|
