Spaces:
Sleeping
Sleeping
| """ | |
| Grad-CAM — Gradient-weighted Class Activation Mapping | |
| Generates heatmap overlays showing which facial regions influenced the prediction. | |
| Essential for IEEE paper explainability section. | |
| """ | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| import torchvision.transforms as T | |
| from PIL import Image | |
| MEAN = [0.485, 0.456, 0.406] | |
| STD = [0.229, 0.224, 0.225] | |
| transform = T.Compose([T.ToTensor(), T.Normalize(mean=MEAN, std=STD)]) | |
| class GradCAM: | |
| """ | |
| Hooks into the last convolutional layer of the spatial branch | |
| to produce a class activation map. | |
| """ | |
| def __init__(self, model): | |
| self.model = model | |
| self.device = next(model.parameters()).device | |
| self._fmaps = None | |
| self._grads = None | |
| self._hook_fwd = None | |
| self._hook_bwd = None | |
| self._register_hooks() | |
| def _register_hooks(self): | |
| # Target: last conv block of EfficientNet spatial branch | |
| try: | |
| if hasattr(self.model.spatial, "backbone"): | |
| target = self.model.spatial.backbone | |
| # timm efficientnet: last block | |
| if hasattr(target, "blocks"): | |
| target_layer = target.blocks[-1] | |
| elif hasattr(target, "features"): | |
| target_layer = target.features[-1] | |
| else: | |
| target_layer = list(target.children())[-2] | |
| else: | |
| target_layer = list(self.model.spatial.children())[-2] | |
| self._hook_fwd = target_layer.register_forward_hook(self._save_fmaps) | |
| self._hook_bwd = target_layer.register_full_backward_hook(self._save_grads) | |
| except Exception: | |
| pass # graceful fallback — still saves plain overlay | |
| def _save_fmaps(self, module, inp, out): | |
| self._fmaps = out.detach() | |
| def _save_grads(self, module, grad_in, grad_out): | |
| self._grads = grad_out[0].detach() | |
| def _compute_cam(self, face_rgb: np.ndarray) -> np.ndarray: | |
| """Returns cam as uint8 (224,224,3) or plain red overlay on failure.""" | |
| if self._hook_fwd is None: | |
| return self._plain_overlay(face_rgb) | |
| tensor = transform(face_rgb).unsqueeze(0).to(self.device) | |
| self.model.zero_grad() | |
| logits = self.model(tensor) | |
| score = torch.sigmoid(logits) | |
| score.backward() | |
| if self._fmaps is None or self._grads is None: | |
| return self._plain_overlay(face_rgb) | |
| weights = self._grads.mean(dim=[2, 3], keepdim=True) # (1, C, 1, 1) | |
| cam = (weights * self._fmaps).sum(dim=1, keepdim=True) # (1, 1, H, W) | |
| cam = F.relu(cam) | |
| cam = cam.squeeze().cpu().numpy() | |
| # Normalise | |
| if cam.max() > 0: | |
| cam = cam / cam.max() | |
| cam_resized = cv2.resize(cam, (224, 224)) | |
| heatmap = cv2.applyColorMap(np.uint8(255 * cam_resized), cv2.COLORMAP_JET) | |
| heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB) | |
| # Overlay on original face | |
| overlay = cv2.addWeighted(face_rgb, 0.6, heatmap, 0.4, 0) | |
| return overlay | |
| def _plain_overlay(face_rgb: np.ndarray) -> np.ndarray: | |
| """Fallback: red-tinted overlay when hooks unavailable.""" | |
| tint = np.zeros_like(face_rgb) | |
| tint[:, :, 0] = 100 | |
| return cv2.addWeighted(face_rgb, 0.8, tint, 0.2, 0) | |
| def generate(self, face_rgb: np.ndarray, save_path: str): | |
| """Compute Grad-CAM and save to disk.""" | |
| try: | |
| overlay = self._compute_cam(face_rgb) | |
| except Exception: | |
| overlay = face_rgb | |
| img = Image.fromarray(overlay.astype(np.uint8)) | |
| img.save(save_path, quality=85) | |
| def remove_hooks(self): | |
| if self._hook_fwd: | |
| self._hook_fwd.remove() | |
| if self._hook_bwd: | |
| self._hook_bwd.remove() | |