Spaces:
Running
Running
| import cv2 | |
| import torch | |
| import numpy as np | |
| from PIL import Image | |
| from typing import List, Callable, Optional | |
| from functools import partial | |
| from pytorch_grad_cam import GradCAM | |
| from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget | |
| from pytorch_grad_cam.utils.image import show_cam_on_image | |
| """ Model wrapper to return a tensor""" | |
| class HuggingfaceToTensorModelWrapper(torch.nn.Module): | |
| def __init__(self, model): | |
| super(HuggingfaceToTensorModelWrapper, self).__init__() | |
| self.model = model | |
| def forward(self, x): | |
| return self.model(x).logits | |
| class ClassActivationMap(object): | |
| def __init__(self, model, processor): | |
| self.model = HuggingfaceToTensorModelWrapper(model) | |
| target_layer = model.swinv2.layernorm | |
| self.target_layer = [target_layer] | |
| self.processor = processor | |
| def swinT_reshape_transform_huggingface(self, tensor, width, height): | |
| result = tensor.reshape(tensor.size(0), | |
| height, | |
| width, | |
| tensor.size(2)) | |
| result = result.transpose(2, 3).transpose(1, 2) | |
| return result | |
| def run_grad_cam_on_image(self, | |
| targets_for_gradcam: List[Callable], | |
| reshape_transform: Optional[Callable], | |
| input_tensor: torch.nn.Module, | |
| input_image: Image, | |
| method: Callable=GradCAM): | |
| with method(model=self.model, | |
| target_layers=self.target_layer, | |
| reshape_transform=reshape_transform) as cam: | |
| # Replicate the tensor for each of the categories we want to create Grad-CAM for: | |
| # print(input_tensor.size()) | |
| repeated_tensor = input_tensor[None, :].repeat(len(targets_for_gradcam), 1, 1, 1) | |
| # print(repeated_tensor.size()) | |
| batch_results = cam(input_tensor=repeated_tensor, | |
| targets=targets_for_gradcam) | |
| results = [] | |
| for grayscale_cam in batch_results: | |
| visualization = show_cam_on_image(np.float32(input_image) / 255, | |
| grayscale_cam, | |
| use_rgb=True) | |
| # Make it weight less in the notebook: | |
| visualization = cv2.resize(visualization, | |
| (visualization.shape[1] // 1, visualization.shape[0] // 1)) | |
| results.append(visualization) | |
| return np.hstack(results) | |
| def get_cam(self, image, category_id): | |
| image = Image.fromarray(image).resize((self.processor.size['height'], self.processor.size['width'])) | |
| img_tensor = self.processor(images=image, return_tensors="pt")['pixel_values'].squeeze() | |
| targets_for_gradcam = [ClassifierOutputTarget(category_id)] | |
| reshape_transform = partial(self.swinT_reshape_transform_huggingface, | |
| width=img_tensor.shape[2] // 32, | |
| height=img_tensor.shape[1] // 32) | |
| cam = self.run_grad_cam_on_image(input_tensor=img_tensor, | |
| input_image=image, | |
| targets_for_gradcam=targets_for_gradcam, | |
| reshape_transform=reshape_transform) | |
| return cam |