Spaces:
Sleeping
Sleeping
File size: 1,984 Bytes
65eeb0e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
import numpy as np
import cv2
import torchvision.transforms as transforms
class GradCAM:
def __init__(self, model, target_layer):
self.model = model
self.target_layer = target_layer
self.gradients = None
self.activations = None
self._register_hooks()
def _register_hooks(self):
def forward_hook(module, input, output):
self.activations = output
def backward_hook(module, grad_in, grad_out):
self.gradients = grad_out[0]
self.target_layer.register_forward_hook(forward_hook)
self.target_layer.register_backward_hook(backward_hook)
def generate_cam(self, input_image, target_class):
self.model.zero_grad()
output = self.model(input_image)
loss = output[:, target_class].sum()
loss.backward()
weights = self.gradients.mean(dim=(2, 3), keepdim=True)
cam = (weights * self.activations).sum(dim=1, keepdim=True)
cam = cam.detach().cpu().numpy()
cam = np.maximum(cam, 0)
cam = cam / cam.max()
cam = cam.squeeze()
return cam
def visualize_cam(self, cam, input_image):
cam = cv2.resize(cam, (input_image.shape[2], input_image.shape[3]))
heatmap = cv2.applyColorMap(np.uint8(255 * cam), cv2.COLORMAP_JET)
heatmap = np.float32(heatmap) / 255
input_image = np.moveaxis(input_image.cpu().numpy()[0], 0, -1)
input_image = np.float32(input_image)
cam_image = heatmap + input_image
cam_image = cam_image / cam_image.max()
return cam_image
def preprocess_image(image):
preprocess = transforms.Compose([
transforms.Grayscale(num_output_channels=1),
transforms.Resize((512, 512)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5], std=[0.5])
])
image = preprocess(image)
image = image.unsqueeze(0)
return image
|