GradCAM

• CAM package

# -*- coding: utf-8 -*-
"""
Created on Fri Sep  2 15:25:33 2022@author: Lenovo
"""
import cv2
import numpy as npclass ActivationsAndGradients:""" Class for extracting activations andregistering gradients from targeted intermediate layers """def __init__(self, model, target_layers, reshape_transform):self.model = modelself.gradients = []self.activations = []self.reshape_transform = reshape_transformself.handles = []for target_layer in target_layers:self.handles.append(target_layer.register_forward_hook(self.save_activation))# Backward compatibility with older pytorch versions:if hasattr(target_layer, 'register_full_backward_hook'):self.handles.append(target_layer.register_full_backward_hook(self.save_gradient))else:self.handles.append(target_layer.register_backward_hook(self.save_gradient))def save_activation(self, module, input, output):activation = outputif self.reshape_transform is not None:activation = self.reshape_transform(activation)self.activations.append(activation.cpu().detach())def save_gradient(self, module, grad_input, grad_output):# Gradients are computed in reverse ordergrad = grad_output[0]if self.reshape_transform is not None:grad = self.reshape_transform(grad)self.gradients = [grad.cpu().detach()] + self.gradientsdef __call__(self, x):self.gradients = []self.activations = []return self.model(x)def release(self):for handle in self.handles:handle.remove()class GradCAM:def __init__(self,model,target_layers,reshape_transform=None,use_cuda=False):self.model = model.eval()self.target_layers = target_layersself.reshape_transform = reshape_transformself.cuda = use_cudaif self.cuda:self.model = model.cuda()self.activations_and_grads = ActivationsAndGradients(self.model, target_layers, reshape_transform)""" Get a vector of weights for every channel in the target layer.Methods that return weights channels,will typically need to only implement this function. """@staticmethoddef get_cam_weights(grads):return np.mean(grads, axis=(2, 3), keepdims=True)@staticmethoddef get_loss(output, target_category):loss = 0for i in range(len(target_category)):loss = loss + output[i, target_category[i]]return lossdef get_cam_image(self, activations, grads):weights = self.get_cam_weights(grads)weighted_activations = weights * activationscam = weighted_activations.sum(axis=1)return cam@staticmethoddef get_target_width_height(input_tensor):width, height = input_tensor.size(-1), input_tensor.size(-2)return width, heightdef compute_cam_per_layer(self, input_tensor):activations_list = [a.cpu().data.numpy()for a in self.activations_and_grads.activations]grads_list = [g.cpu().data.numpy()for g in self.activations_and_grads.gradients]target_size = self.get_target_width_height(input_tensor)cam_per_target_layer = []# Loop over the saliency image from every layerfor layer_activations, layer_grads in zip(activations_list, grads_list):cam = self.get_cam_image(layer_activations, layer_grads)cam[cam < 0] = 0  # works like mute the min-max scale in the function of scale_cam_imagescaled = self.scale_cam_image(cam, target_size)cam_per_target_layer.append(scaled[:, None, :])return cam_per_target_layerdef aggregate_multi_layers(self, cam_per_target_layer):cam_per_target_layer = np.concatenate(cam_per_target_layer, axis=1)cam_per_target_layer = np.maximum(cam_per_target_layer, 0)result = np.mean(cam_per_target_layer, axis=1)return self.scale_cam_image(result)@staticmethoddef scale_cam_image(cam, target_size=None):result = []for img in cam:img = img - np.min(img)img = img / (1e-7 + np.max(img))if target_size is not None:img = cv2.resize(img, target_size)result.append(img)result = np.float32(result)return resultdef __call__(self, input_tensor, target_category=None):if self.cuda:input_tensor = input_tensor.cuda()# 正向传播得到网络输出logits(未经过softmax)output = self.activations_and_grads(input_tensor)if isinstance(target_category, int):target_category = [target_category] * input_tensor.size(0)if target_category is None:target_category = np.argmax(output.cpu().data.numpy(), axis=-1)print(f"category id: {target_category}")else:assert (len(target_category) == input_tensor.size(0))self.model.zero_grad()loss = self.get_loss(output, target_category)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)return self.aggregate_multi_layers(cam_per_layer)def __del__(self):self.activations_and_grads.release()def __enter__(self):return selfdef __exit__(self, exc_type, exc_value, exc_tb):self.activations_and_grads.release()if isinstance(exc_value, IndexError):# Handle IndexError here...print(f"An exception occurred in CAM with block: {exc_type}. Message: {exc_value}")return True

调用函数

import sys
from tqdm import tqdm
import torch
from torchvision import transforms
import numpy as np
import os
import cv2
from PIL import Image
from CAM import GradCAM
def show_cam_on_image(img: np.ndarray,mask: np.ndarray,use_rgb: bool = False,colormap: int = cv2.COLORMAP_JET) -> np.ndarray:""" This function overlays the cam mask on the image as an heatmap.By default the heatmap is in BGR format.:param img: The base image in RGB or BGR format.:param mask: The cam mask.:param use_rgb: Whether to use an RGB or BGR heatmap, this should be set to True if 'img' is in RGB format.:param colormap: The OpenCV colormap to be used.:returns: The default image with the cam overlay."""heatmap = cv2.applyColorMap(np.uint8(255 * mask), colormap)if use_rgb:heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)heatmap = np.float32(heatmap) / 255if np.max(img) > 1:raise Exception("The input image should np.float32 in the range [0, 1]")cam = heatmap + imgcam = cam / np.max(cam)return np.uint8(255 * cam)def save_cam_mask(cam_path,mask: np.ndarray,w: int = 224,h: int = 224,use_rgb: bool = True,colormap: int = cv2.COLORMAP_JET):'''只保存生成的CAM，其中的参数mask表示生成的CAM:cam_path cam保存的地址:mask 生成的CAM,此时是224*224的灰度图，需要转换成RGB:h 图像的高:w 图像的宽:ues_rgb 使用RGB格式:colormap ...'''heatmap = cv2.applyColorMap(np.uint8(255 * mask), colormap)if use_rgb:heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)heatmap = transforms.Resize([w, h])(Image.fromarray(heatmap))heatmap.save(cam_path)def get_cam(model, img_path, target_layers, data_transform):'''根据model,target_layer对图像img_path进行cam可视化:db  使用的model:img_path  需要可是化的图像:target_layer  目标可视化那一层的特征图:data_transform:cam_path  生成的cam图片:camAorig_path  生成cam叠加原图的图片'''assert os.path.exists(img_path), "file: '{}' dose not exist.".format(img_path)img = Image.open(img_path).convert('RGB')h, w = img.sizeorimg = img.resize((224, 224))orimg = np.array(orimg, dtype=np.uint8)img_tensor = data_transform(img)input_tensor = torch.unsqueeze(img_tensor, dim=0)cam = GradCAM(model=model, target_layers=target_layers, use_cuda=False)target_category = 1  # tabby, tabby cat# 生成cam灰度图grayscale_cam = cam(input_tensor=input_tensor, target_category=target_category)grayscale_cam = grayscale_cam[0, :]# 将CAM的灰度图转换成三通道并保存#    save_cam_mask(cam_path= cam_path,mask = grayscale_cam,w=w,h=h)# 将原图和cam进行叠加显示visualization = show_cam_on_image(orimg.astype(dtype=np.float32) / 255.,grayscale_cam,use_rgb=True)visualization = transforms.Resize([w, h])(Image.fromarray(visualization))#    visualization.save(camAorig_path)return visualization

使用函数

# Grad-CAM
import os
import cv2
import torch
import utils
from torchvision import transforms
from model import convnext_tiny as create_model
from PIL import Image
num_classes = 2
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
data_transform = transforms.Compose([transforms.Resize([224, 224]),transforms.ToTensor(),transforms.Normalize([0.456, 0.485, 0.406],[0.224, 0.229, 0.225])])    
img_dir='../test_data/img/' 
cam_dir='../test_data/cam/'
color_list=os.listdir(img_dir)
#%%
weights_path = '..t/weights/patch.pth'
model = create_model(num_classes=num_classes).to(device)
model.load_state_dict(torch.load(weights_path, map_location=device))
target_layers = model.stages[3]
for picture in color_list:img_path = os.path.join(img_dir,picture)cam_path = os.path.join(cam_dir,picture)cam = utils.get_cam(model, img_path, target_layers, data_transform)
#    print(type(cam))cam.save(cam_path)