Grad-CAM可视化特征图

安装pytorch-grad-cam

pip install grad-cam

Demo

import cv2
import torch
from torchvision import transforms
from PIL import Image
from pytorch_grad_cam import GradCAM, HiResCAM, ScoreCAM, GradCAMPlusPlus, AblationCAM, XGradCAM, EigenCAM, FullGrad
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
from pytorch_grad_cam.utils.image import show_cam_on_image
from torchvision.models import resnet50
import numpy as np
from pytorch_grad_cam.utils.image import show_cam_on_image, \
    deprocess_image, \
    preprocess_image
image_path = "frame_3840.jpg"

# 加载图片
image = Image.open(image_path).convert('RGB')

rgb_img = cv2.imread(image_path, 1)[:, :, ::-1]  # 1是读取rgb
# imread返回从指定路径加载的图像
rgb_img = cv2.imread(image_path, 1)  # imread()读取的是BGR格式
rgb_img = np.float32(rgb_img) / 255

# preprocess_image作用：归一化图像，并转成tensor
input_tensor = preprocess_image(rgb_img, mean=[0.485, 0.456, 0.406],
                                std=[0.229, 0.224, 0.225])  # torch.Size([1, 3, 224, 224])
# Create an input tensor image for your model..
# Note: input_tensor can be a batch tensor with several images!

model = resnet50(pretrained=True)
target_layer = [model.layer4[-1]]
# input_tensor = # Create an input tensor image for your model..
# Note: input_tensor can be a batch tensor with several images!

# # Construct the CAM object once, and then re-use it on many images:
# cam = GradCAM(model=model, target_layers=target_layers)
#
# # You can also use it within a with statement, to make sure it is freed,
# # In case you need to re-create it inside an outer loop:
# # with GradCAM(model=model, target_layers=target_layers) as cam:
# #   ...
#
# # We have to specify the target we want to generate
# # the Class Activation Maps for.
# # If targets is None, the highest scoring category
# # will be used for every image in the batch.
# # Here we use ClassifierOutputTarget, but you can define your own custom targets
# # That are, for example, combinations of categories, or specific outputs in a non standard model.
#
# targets = [ClassifierOutputTarget(751)]
#
# # You can also pass aug_smooth=True and eigen_smooth=True, to apply smoothing.
# grayscale_cam = cam(input_tensor=input_tensor, targets=targets)
#
# # In this example grayscale_cam has only one image in the batch:
# grayscale_cam = grayscale_cam[0, :]
# rgb_img = np.array(image.resize((224, 224)))
# visualization = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True)
#
# # You can also get the model outputs without having to re-inference
# model_outputs = cam.outputs

'''
3)初始化CAM对象，包括模型，目标层以及是否使用cuda等
'''
# Construct the CAM object once, and then re-use it on many images:
cam = GradCAM(model=model, target_layers=target_layer)
'''
4)选定目标类别，如果不设置，则默认为分数最高的那一类
'''
# If target_category is None, the highest scoring category
# will be used for every image in the batch.
# target_category can also be an integer, or a list of different integers
# for every image in the batch.
# target_category = None
target_category = None

'''
5)计算cam
'''
# You can also pass aug_smooth=True and eigen_smooth=True, to apply smoothing.
# grayscale_cam = cam(input_tensor=input_tensor, targets=target_category)  # [batch, 224,224]
grayscale_cam = cam(input_tensor=input_tensor, targets=target_category)  # [batch, 224,224]
# ----------------------------------
'''
6)展示热力图并保存
'''
# In this example grayscale_cam has only one image in the batch:
# 7.展示热力图并保存, grayscale_cam是一个batch的结果，只能选择一张进行展示
grayscale_cam = grayscale_cam[0]
visualization = show_cam_on_image(rgb_img, grayscale_cam)  # (224, 224, 3)
cv2.imwrite(f'first_try.jpg', visualization)