import numpy as np import torch import ttach as tta from typing import Callable, List, Tuple from einops import rearrange from pytorch_grad_cam.activations_and_gradients import ActivationsAndGradients from pytorch_grad_cam.utils.svd_on_activations import get_2d_projection from pytorch_grad_cam.utils.image import scale_cam_image from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget class BaseCAM: def __init__(self, model: torch.nn.Module, target_layers: List[torch.nn.Module], use_cuda: bool = False, reshape_transform: Callable = None, compute_input_gradient: bool = False, uses_gradients: bool = True) -> None: self.model = model.eval() self.target_layers = target_layers self.cuda = use_cuda if self.cuda: self.model = model.cuda() self.reshape_transform = reshape_transform self.compute_input_gradient = compute_input_gradient self.uses_gradients = uses_gradients self.activations_and_grads = ActivationsAndGradients( self.model, target_layers, reshape_transform) """ Get a vector of weights for every channel in the target layer. Methods that return weights channels, will typically need to only implement this function. """ def get_cam_weights(self, input_tensor: torch.Tensor, target_layers: List[torch.nn.Module], targets: List[torch.nn.Module], activations: torch.Tensor, grads: torch.Tensor) -> np.ndarray: raise Exception("Not Implemented") def get_cam_image(self, input_tensor: torch.Tensor, target_layer: torch.nn.Module, targets: List[torch.nn.Module], activations: torch.Tensor, grads: torch.Tensor, eigen_smooth: bool = False) -> np.ndarray: weights = self.get_cam_weights(input_tensor, target_layer, targets, activations, grads) weighted_activations = weights[:, None, :] * activations H = W = int(weighted_activations.shape[1] ** 0.5) weighted_activations = rearrange(weighted_activations, "b (h w) c -> b c h w", h=H, w=W) if eigen_smooth: cam = get_2d_projection(weighted_activations) else: cam = weighted_activations.sum(axis=1) return cam def forward(self, input_tensor: torch.Tensor, targets: List[torch.nn.Module], eigen_smooth: bool = False, return_probs: bool = False) -> np.ndarray: if self.cuda: input_tensor = input_tensor.cuda() if self.compute_input_gradient: input_tensor = torch.autograd.Variable(input_tensor, requires_grad=True) outputs = self.activations_and_grads(input_tensor) target_categories = np.argmax(outputs.cpu().data.numpy(), axis=-1) if targets is None: targets = [ClassifierOutputTarget( category) for category in target_categories] if self.uses_gradients: self.model.zero_grad() loss = sum([target(output) for target, output in zip(targets, outputs)]) loss.backward(retain_graph=True) # In most of the saliency attribution papers, the saliency is # computed with a single target layer. # Commonly it is the last convolutional layer. # Here we support passing a list with multiple target layers. # It will compute the saliency image for every image, # and then aggregate them (with a default mean aggregation). # This gives you more flexibility in case you just want to # use all conv layers for example, all Batchnorm layers, # or something else. cam_per_layer = self.compute_cam_per_layer(input_tensor, targets, eigen_smooth) if not return_probs: return self.aggregate_multi_layers(cam_per_layer), target_categories return self.aggregate_multi_layers(cam_per_layer), torch.nn.functional.softmax(outputs, dim=-1).detach().cpu().numpy() def get_target_width_height(self, input_tensor: torch.Tensor) -> Tuple[int, int]: width, height = input_tensor.size(-1), input_tensor.size(-2) return width, height def compute_cam_per_layer( self, input_tensor: torch.Tensor, targets: List[torch.nn.Module], eigen_smooth: bool) -> np.ndarray: activations_list = [a.cpu().data.numpy() for a in self.activations_and_grads.activations] grads_list = [g.cpu().data.numpy() for g in self.activations_and_grads.gradients] target_size = self.get_target_width_height(input_tensor) cam_per_target_layer = [] # Loop over the saliency image from every layer for i in range(len(self.target_layers)): target_layer = self.target_layers[i] layer_activations = None layer_grads = None if i < len(activations_list): layer_activations = activations_list[i] if i < len(grads_list): layer_grads = grads_list[i] cam = self.get_cam_image(input_tensor, target_layer, targets, layer_activations, layer_grads, eigen_smooth) cam = np.maximum(cam, 0) scaled = scale_cam_image(cam, target_size) cam_per_target_layer.append(scaled[:, None, :]) return cam_per_target_layer def aggregate_multi_layers( self, cam_per_target_layer: np.ndarray) -> np.ndarray: cam_per_target_layer = np.concatenate(cam_per_target_layer, axis=1) cam_per_target_layer = np.maximum(cam_per_target_layer, 0) result = np.mean(cam_per_target_layer, axis=1) return scale_cam_image(result) def forward_augmentation_smoothing(self, input_tensor: torch.Tensor, targets: List[torch.nn.Module], eigen_smooth: bool = False, return_probs: bool = False) -> np.ndarray: transforms = tta.Compose( [ tta.HorizontalFlip(), tta.Multiply(factors=[0.9, 1, 1.1]), ] ) cams = [] for transform in transforms: augmented_tensor = transform.augment_image(input_tensor) cam, b = self.forward(augmented_tensor, targets, eigen_smooth, return_probs=return_probs) # The ttach library expects a tensor of size BxCxHxW cam = cam[:, None, :, :] cam = torch.from_numpy(cam) cam = transform.deaugment_mask(cam) # Back to numpy float32, HxW cam = cam.numpy() cam = cam[:, 0, :, :] cams.append(cam) cam = np.mean(np.float32(cams), axis=0) return cam, b def __call__(self, input_tensor: torch.Tensor, targets: List[torch.nn.Module] = None, aug_smooth: bool = False, eigen_smooth: bool = False, return_probs: bool = False) -> np.ndarray: # Smooth the CAM result with test time augmentation if aug_smooth is True: return self.forward_augmentation_smoothing( input_tensor, targets, eigen_smooth, return_probs=return_probs) return self.forward(input_tensor, targets, eigen_smooth, return_probs=return_probs) def __del__(self): self.activations_and_grads.release() def __enter__(self): return self def __exit__(self, exc_type, exc_value, exc_tb): self.activations_and_grads.release() if isinstance(exc_value, IndexError): # Handle IndexError here... print( f"An exception occurred in CAM with block: {exc_type}. Message: {exc_value}") return True