Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import torch | |
| import ttach as tta | |
| from typing import Callable, List, Tuple | |
| 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, None] * activations | |
| 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) -> 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) | |
| if targets is None: | |
| target_categories = np.argmax(outputs.cpu().data.numpy(), axis=-1) | |
| 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) | |
| return self.aggregate_multi_layers(cam_per_layer) | |
| 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) -> 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 = self.forward(augmented_tensor, | |
| targets, | |
| eigen_smooth) | |
| # 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 | |
| def __call__(self, | |
| input_tensor: torch.Tensor, | |
| targets: List[torch.nn.Module] = None, | |
| aug_smooth: bool = False, | |
| eigen_smooth: 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 self.forward(input_tensor, | |
| targets, eigen_smooth) | |
| 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 | |