File size: 3,346 Bytes
a639402 | 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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | import cv2
import torch
import numpy as np
class InfoHolder():
def __init__(self, heatmap_layer):
self.gradient = None
self.activation = None
self.heatmap_layer = heatmap_layer
def get_gradient(self, grad):
self.gradient = grad
def hook(self, model, input, output):
output.register_hook(self.get_gradient)
self.activation = output.detach()
def generate_heatmap(weighted_activation):
raw_heatmap = torch.mean(weighted_activation, 0)
heatmap = np.maximum(raw_heatmap.detach().cpu(), 0)
heatmap /= torch.max(heatmap) + 1e-10
return heatmap.numpy()
def superimpose(input_img, heatmap):
img = cv2.cvtColor(input_img,cv2.COLOR_BGR2RGB)
heatmap = cv2.resize(heatmap, (img.shape[1], img.shape[0]))
heatmap = np.uint8(255 * heatmap)
heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
superimposed_img = np.uint8(heatmap * 0.6 + img * 0.4)
pil_img = cv2.cvtColor(superimposed_img,cv2.COLOR_BGR2RGB)
return pil_img
def superimpose_to_mask(input_img, heatmap):
img = cv2.cvtColor(input_img,cv2.COLOR_BGR2RGB)
heatmap = cv2.resize(heatmap, (img.shape[1], img.shape[0]))
heatmap = np.uint8(255 * heatmap)
heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_BONE)
superimposed_img = np.uint8(heatmap * 1 + img * 0.0)
pil_img = cv2.cvtColor(heatmap,cv2.COLOR_BGR2GRAY)
(thresh, blackAndWhiteImage) = cv2.threshold(pil_img, 127, 255, cv2.THRESH_BINARY)
return blackAndWhiteImage
def to_RGB(tensor):
tensor = (tensor - tensor.min())
tensor = tensor/(tensor.max() + 1e-10)
image_binary = np.transpose(tensor.cpu().numpy(), (1, 2, 0))
image = np.uint8(255 * image_binary)
return image
def grad_cam(model, input_tensor, heatmap_layer, truelabel=None):
info = InfoHolder(heatmap_layer)
heatmap_layer.register_forward_hook(info.hook)
output = model(input_tensor.unsqueeze(0))[0]
truelabel = truelabel if truelabel else torch.argmax(output)
output[truelabel].backward()
weights = torch.mean(info.gradient, [0, 2, 3])
activation = info.activation.squeeze(0)
weighted_activation = torch.zeros(activation.shape)
for idx, (weight, activation) in enumerate(zip(weights, activation)):
weighted_activation[idx] = weight * activation
heatmap = generate_heatmap(weighted_activation)
input_image = to_RGB(input_tensor)
return superimpose(input_image, heatmap)
def save_grad_grad_cam(model, input_tensor, heatmap_layer, truelabel=None):
info = InfoHolder(heatmap_layer)
heatmap_layer.register_forward_hook(info.hook)
output = model(input_tensor.unsqueeze(0))[0]
truelabel = truelabel if truelabel else torch.argmax(output)
output[truelabel].backward()
weights = torch.mean(info.gradient, [0, 2, 3])
activation = info.activation.squeeze(0)
weighted_activation = torch.zeros(activation.shape)
for idx, (weight, activation) in enumerate(zip(weights, activation)):
weighted_activation[idx] = weight * activation
heatmap = generate_heatmap(weighted_activation)
input_image = to_RGB(input_tensor)
grad_image = superimpose_to_mask(input_image, heatmap)
return input_image, grad_image |