Spaces:
Sleeping
Sleeping
| import cv2 | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| from PIL import Image | |
| class MultiHeadGradCAM: | |
| """ | |
| Grad-CAM for Multi-Head ConvNeXt V2 models. | |
| Supports targeting a specific head (1, 2, or 3) and a specific class within that head. | |
| """ | |
| def __init__(self, model, target_layer): | |
| self.model = model | |
| self.target_layer = target_layer | |
| self.gradients = None | |
| self.activations = None | |
| # Register hooks | |
| self.target_layer.register_forward_hook(self.save_activation) | |
| self.target_layer.register_full_backward_hook(self.save_gradient) | |
| def save_activation(self, module, input, output): | |
| if torch.backends.mps.is_available(): | |
| torch.mps.synchronize() | |
| self.activations = output.clone().detach() | |
| def save_gradient(self, module, grad_input, grad_output): | |
| self.gradients = grad_output[0] | |
| def generate_cam(self, input_tensor, target_head=2, sub_head=None, target_class=None): | |
| """ | |
| Generates the Class Activation Map (CAM). | |
| Args: | |
| input_tensor (torch.Tensor): Preprocessed image tensor (1, C, H, W) | |
| target_head (int): Which head to explain (1, 2, or 3). Defaults to 2 (Pathology Routing). | |
| sub_head (str, optional): The severity family (e.g., 'macular') if target_head == 3. | |
| target_class (int, optional): The class within the target head. | |
| If None, uses the class with the highest score. | |
| """ | |
| self.model.zero_grad() | |
| # Forward pass returns a dict | |
| outputs = self.model(input_tensor) | |
| if target_head == 1: | |
| target_output = outputs['normal_abnormal'] | |
| elif target_head == 2: | |
| target_output = outputs['pathology'] | |
| elif target_head == 3: | |
| if sub_head is None: | |
| raise ValueError("sub_head must be provided when target_head=3 (e.g., 'macular')") | |
| target_output = outputs['severity'][sub_head] | |
| else: | |
| raise ValueError(f"Unknown target_head: {target_head}") | |
| # If binary (Head 1), target_output is [1, 1] | |
| if target_output.shape[1] == 1: | |
| score = target_output[0, 0] | |
| else: | |
| if target_class is None: | |
| target_class = target_output.argmax(dim=1).item() | |
| score = target_output[0, target_class] | |
| # Backward pass | |
| score.backward() | |
| gradients = self.gradients.detach().cpu().numpy()[0] # (C, H, W) | |
| activations = self.activations.detach().cpu().numpy()[0] # (C, H, W) | |
| weights = np.mean(gradients, axis=(1, 2)) # (C,) | |
| cam = np.zeros(activations.shape[1:], dtype=np.float32) | |
| for i, w in enumerate(weights): | |
| cam += w * activations[i, :, :] | |
| cam = np.maximum(cam, 0) | |
| cam = cv2.resize(cam, (input_tensor.shape[3], input_tensor.shape[2])) | |
| cam_max = np.max(cam) | |
| cam_min = np.min(cam) | |
| if cam_max - cam_min > 1e-8: | |
| cam = (cam - cam_min) / (cam_max - cam_min) | |
| else: | |
| cam = np.zeros_like(cam) | |
| return cam | |
| def overlay_cam(img_pil: Image.Image, cam: np.ndarray, alpha: float = 0.5) -> Image.Image: | |
| if cam.shape != img_pil.size[::-1]: | |
| cam = cv2.resize(cam, img_pil.size) | |
| img_cv2 = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR) | |
| heatmap = np.uint8(255 * cam) | |
| heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET) | |
| superimposed = cv2.addWeighted(img_cv2, 1 - alpha, heatmap, alpha, 0) | |
| superimposed = cv2.cvtColor(superimposed, cv2.COLOR_BGR2RGB) | |
| return Image.fromarray(superimposed) | |
| class HierarchicalUNetGradCAM: | |
| """ | |
| Grad-CAM for Hierarchical U-Net models (unified segmentation/classification). | |
| Supports targeting a specific head (1, 2, or 3) and a specific class within that head. | |
| """ | |
| def __init__(self, model, target_layer): | |
| self.model = model | |
| self.target_layer = target_layer | |
| self.gradients = None | |
| self.activations = None | |
| # Register hooks | |
| self.target_layer.register_forward_hook(self.save_activation) | |
| self.target_layer.register_full_backward_hook(self.save_gradient) | |
| def save_activation(self, module, input, output): | |
| if torch.backends.mps.is_available(): | |
| torch.mps.synchronize() | |
| self.activations = output.clone().detach() | |
| def save_gradient(self, module, grad_input, grad_output): | |
| self.gradients = grad_output[0] | |
| def generate_cam(self, input_tensor, target_head=2, sub_head=None, target_class=None): | |
| """ | |
| Generates the Class Activation Map (CAM). | |
| Args: | |
| input_tensor (torch.Tensor): Preprocessed image tensor (1, 1, H, W) | |
| target_head (int): Which head to explain (1, 2, or 3). Defaults to 2 (Pathology Routing). | |
| sub_head (str, optional): The severity family (e.g., 'macular') if target_head == 3. | |
| target_class (int, optional): The class within the target head. | |
| If None, uses the class with the highest score. | |
| """ | |
| self.model.zero_grad() | |
| # Forward pass returning just classification logits | |
| outputs = self.model(input_tensor, task="classification") | |
| if target_head == 1: | |
| target_output = outputs['normal_abnormal'] | |
| elif target_head == 2: | |
| target_output = outputs['pathology'] | |
| elif target_head == 3: | |
| if sub_head is None: | |
| raise ValueError("sub_head must be provided when target_head=3 (e.g., 'macular')") | |
| target_output = outputs['severity'][sub_head] | |
| else: | |
| raise ValueError(f"Unknown target_head: {target_head}") | |
| # If binary (Head 1), target_output is [1, 1] | |
| if target_output.shape[1] == 1: | |
| score = target_output[0, 0] | |
| else: | |
| if target_class is None: | |
| target_class = target_output.argmax(dim=1).item() | |
| score = target_output[0, target_class] | |
| # Backward pass | |
| score.backward() | |
| gradients = self.gradients.detach().cpu().numpy()[0] # (C, H, W) | |
| activations = self.activations.detach().cpu().numpy()[0] # (C, H, W) | |
| weights = np.mean(gradients, axis=(1, 2)) # (C,) | |
| cam = np.zeros(activations.shape[1:], dtype=np.float32) | |
| for i, w in enumerate(weights): | |
| cam += w * activations[i, :, :] | |
| cam = np.maximum(cam, 0) | |
| cam = cv2.resize(cam, (input_tensor.shape[3], input_tensor.shape[2])) | |
| cam_max = np.max(cam) | |
| cam_min = np.min(cam) | |
| if cam_max - cam_min > 1e-8: | |
| cam = (cam - cam_min) / (cam_max - cam_min) | |
| else: | |
| cam = np.zeros_like(cam) | |
| return cam | |
| def overlay_cam(img_pil: Image.Image, cam: np.ndarray, alpha: float = 0.5) -> Image.Image: | |
| if cam.shape != img_pil.size[::-1]: | |
| cam = cv2.resize(cam, img_pil.size) | |
| img_arr = np.array(img_pil) | |
| if len(img_arr.shape) == 2: | |
| # It's a grayscale image (H, W) | |
| img_cv2 = cv2.cvtColor(img_arr, cv2.COLOR_GRAY2BGR) | |
| else: | |
| # It's an RGB image (H, W, 3) | |
| img_cv2 = cv2.cvtColor(img_arr, cv2.COLOR_RGB2BGR) | |
| heatmap = np.uint8(255 * cam) | |
| heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET) | |
| superimposed = cv2.addWeighted(img_cv2, 1 - alpha, heatmap, alpha, 0) | |
| superimposed = cv2.cvtColor(superimposed, cv2.COLOR_BGR2RGB) | |
| return Image.fromarray(superimposed) | |