Spaces:
Sleeping
Sleeping
File size: 3,484 Bytes
5bfab10 |
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 |
from typing import List
import torch
import numpy as np
import utils
from pytorch_grad_cam.base_cam import BaseCAM
from pytorch_grad_cam.utils import get_2d_projection
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
class YoloGradCAM(BaseCAM):
def __init__(self,
model,
target_layers,
scaled_anchors,
use_cuda=False,
reshape_transform=None):
super(YoloGradCAM, self).__init__(model,
target_layers,
use_cuda,
reshape_transform,
uses_gradients=False)
self.scaled_anchors = scaled_anchors
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:
return get_2d_projection(activations)
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:
bboxes = [[] for _ in range(1)]
for i in range(3):
batch_size, A, S, _, _ = outputs[i].shape
anchor = self.scaled_anchors[i]
boxes_scale_i = utils.cells_to_bboxes(
outputs[i], anchor, S=S, is_preds=True
)
for idx, (box) in enumerate(boxes_scale_i):
bboxes[idx] += box
nms_boxes = utils.non_max_suppression(
bboxes[0], iou_threshold=0.5, threshold=0.4, box_format="midpoint",
)
# target_categories = np.argmax(outputs.cpu().data.numpy(), axis=-1)
target_categories = [box[0] for box in nms_boxes]
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) |