|
|
import torch
|
|
|
import torch.nn.functional as F
|
|
|
import cv2
|
|
|
import numpy as np
|
|
|
from src.model import get_model
|
|
|
|
|
|
class GradCAM:
|
|
|
def __init__(self, model_path, target_layer='layer4'):
|
|
|
self.model = get_model(pretrained=False)
|
|
|
self.model.load_state_dict(torch.load(model_path, map_location='cpu'))
|
|
|
self.model.eval()
|
|
|
self.target_layer = target_layer
|
|
|
self.gradients = None
|
|
|
self.activations = None
|
|
|
self.hooks = []
|
|
|
self._register_hooks()
|
|
|
|
|
|
def _register_hooks(self):
|
|
|
def backward_hook(module, grad_input, grad_output):
|
|
|
self.gradients = grad_output[0]
|
|
|
|
|
|
def forward_hook(module, input, output):
|
|
|
self.activations = output
|
|
|
|
|
|
for name, module in self.model.named_modules():
|
|
|
if target_layer in name:
|
|
|
self.hooks.append(module.register_forward_hook(forward_hook))
|
|
|
self.hooks.append(module.register_backward_hook(backward_hook))
|
|
|
|
|
|
def generate(self, input_tensor, class_idx=None):
|
|
|
self.model.zero_grad()
|
|
|
output = self.model(input_tensor)
|
|
|
if class_idx is None:
|
|
|
class_idx = output.argmax().item()
|
|
|
score = output[0, class_idx]
|
|
|
score.backward()
|
|
|
|
|
|
gradients = self.gradients[0]
|
|
|
activations = self.activations[0]
|
|
|
weights = torch.mean(gradients, dim=(1, 2), keepdim=True)
|
|
|
|
|
|
cam = torch.sum(weights * activations, dim=1, keepdim=True)
|
|
|
cam = F.relu(cam)
|
|
|
cam = cam.squeeze().detach().cpu().numpy()
|
|
|
cam = (cam - np.min(cam)) / (np.max(cam) - np.min(cam) + 1e-8)
|
|
|
|
|
|
|
|
|
for hook in self.hooks:
|
|
|
hook.remove()
|
|
|
|
|
|
return cam
|
|
|
|
|
|
def __del__(self):
|
|
|
for hook in self.hooks:
|
|
|
hook.remove()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
