File size: 5,872 Bytes
1ea7ba6 | 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 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | """
Grad-CAM visualization for SpiceFusionNet.
Hooks into EfficientNet-B4's last convolutional layer.
Bug fixes:
- Display image now matches the *exact* image the model sees
(resize -> center crop), so heatmap pixels align correctly.
- Optional fusion-mode Grad-CAM: computes gradients w.r.t. the
fusion_head logits (matches the actual prediction pathway).
- Uses last MBConv block features (slightly richer than conv_head).
"""
import numpy as np
import cv2
import torch
import torch.nn.functional as F
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from pathlib import Path
from PIL import Image
import albumentations as A
from albumentations.pytorch import ToTensorV2
import config
from src.features import extract_all as _extract_all
# Resize-then-centercrop dimensions (must match get_val_transform exactly)
_RESIZE = int(config.IMG_SIZE * 256 / 224)
_CROP = config.IMG_SIZE
class GradCAM:
def __init__(self, model, target_layer=None, mode: str = "fusion"):
"""mode = "fusion" or "image" — which forward pass to backprop through."""
self.model = model
self.mode = mode
self.gradients = None
self.activations = None
# Last conv block of EfficientNet-B4. `conv_head` is a 1x1 conv;
# `blocks[-1]` carries the spatial features that flow into it.
if target_layer is None:
target_layer = model.backbone.conv_head
self._fwd = target_layer.register_forward_hook(self._save_activation)
self._bwd = target_layer.register_full_backward_hook(self._save_gradient)
def _save_activation(self, module, inp, out):
self.activations = out.detach()
def _save_gradient(self, module, grad_in, grad_out):
self.gradients = grad_out[0].detach()
def remove(self):
self._fwd.remove()
self._bwd.remove()
def __call__(self, x: torch.Tensor, tex=None, col=None, class_idx: int = None):
"""Returns (cam, pred_idx). cam is H×W float32 in [0,1]."""
self.model.eval()
self.model.zero_grad()
if self.mode == "fusion" and tex is not None and col is not None:
logits, _ = self.model.forward_fusion(x, tex, col)
else:
logits = self.model.forward_image(x)
if class_idx is None:
class_idx = logits.argmax(dim=1).item()
logits[0, class_idx].backward()
weights = self.gradients.mean(dim=(2, 3), keepdim=True)
cam = (weights * self.activations).sum(dim=1, keepdim=True)
cam = F.relu(cam).squeeze().cpu().numpy()
cam = cam - cam.min()
if cam.max() > 0:
cam = cam / cam.max()
return cam, class_idx
def _model_view(img_path: str):
"""Return BOTH the tensor the model sees AND the matching uint8 display
image, guaranteed to be the exact same pixels (resize -> center crop)."""
img = np.array(Image.open(img_path).convert("RGB"))
# Step 1: resize + center crop to produce display image (uint8)
crop_tf = A.Compose([
A.Resize(_RESIZE, _RESIZE),
A.CenterCrop(_CROP, _CROP),
])
display = crop_tf(image=img)["image"] # (224, 224, 3) uint8
# Step 2: normalize + to-tensor for the model
norm_tf = A.Compose([
A.Normalize(mean=config.IMG_MEAN, std=config.IMG_STD),
ToTensorV2(),
])
tensor = norm_tf(image=display)["image"].unsqueeze(0) # (1, 3, 224, 224)
# Step 3: hand-crafted features on the ORIGINAL image (matches dataset.py)
tex_np, col_np = _extract_all(img)
tex = torch.from_numpy(tex_np).unsqueeze(0).float()
col = torch.from_numpy(col_np).unsqueeze(0).float()
return tensor, display, tex, col
def visualize_gradcam(
model,
img_paths: list,
true_labels: list,
classes: list,
device: torch.device,
output_path: Path,
n_samples: int = 16,
mode: str = "fusion",
):
"""Plot Grad-CAM for up to n_samples images.
mode='fusion' computes gradients through the fusion_head (matches the
actual prediction). mode='image' uses the Phase-1 img_head.
"""
model.eval()
gcam = GradCAM(model, mode=mode)
cols = 4
rows = (n_samples + cols - 1) // cols
fig, axes = plt.subplots(rows, cols * 2, figsize=(cols * 6, rows * 3))
axes = axes.flatten()
plotted = 0
for img_path, true_label in zip(img_paths, true_labels):
if plotted >= n_samples:
break
tensor, display, tex, col = _model_view(img_path)
tensor = tensor.to(device); tex = tex.to(device); col = col.to(device)
cam, pred_idx = gcam(tensor, tex=tex, col=col)
# Overlay — display & cam now share identical 224x224 spatial grid
heatmap = cv2.resize(cam, (display.shape[1], display.shape[0]))
heatmap = cv2.applyColorMap(np.uint8(255 * heatmap), cv2.COLORMAP_JET)
heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
overlay = cv2.addWeighted(display, 0.5, heatmap, 0.5, 0)
ax_img = axes[plotted * 2]
ax_cam = axes[plotted * 2 + 1]
ax_img.imshow(display)
ax_img.set_title(f"True: {classes[true_label]}", fontsize=8)
ax_img.axis("off")
ax_cam.imshow(overlay)
color = "green" if pred_idx == true_label else "red"
ax_cam.set_title(f"Pred: {classes[pred_idx]}", fontsize=8, color=color)
ax_cam.axis("off")
plotted += 1
for i in range(plotted * 2, len(axes)):
axes[i].axis("off")
plt.suptitle(f"Grad-CAM ({mode} mode)", fontsize=13)
plt.tight_layout()
plt.savefig(output_path, dpi=150)
plt.close()
gcam.remove()
print(f"Grad-CAM saved to {output_path}")
|