导入相关库
import argparse
import os
import cv2
import numpy as np
import torch
from torchvision import models
from pytorch_grad_cam import (
GradCAM, HiResCAM, ScoreCAM, GradCAMPlusPlus,
AblationCAM, XGradCAM, EigenCAM, EigenGradCAM,
LayerCAM, FullGrad, GradCAMElementWise
)
from pytorch_grad_cam import GuidedBackpropReLUModel
from pytorch_grad_cam.utils.image import (
show_cam_on_image, deprocess_image, preprocess_image
)
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
导入自己的模型
这里以ResNet为例
import torchvision.models as models
#自己模型可以这样导入
# from xxx import xxx
加载相关参数
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument('--use-cuda', action='store_true', default=False,
help='Use NVIDIA GPU acceleration')
parser.add_argument(
'--image-path',
type=str,
default='autodl-tmp/mydata/cam/pytorch-grad-cam-master/result/storagetank.jpg',
help='Input image path')
parser.add_argument('--aug-smooth', action='store_true',
help='Apply test time augmentation to smooth the CAM')
parser.add_argument(
'--eigen-smooth',
action='store_true',
help='Reduce noise by taking the first principle component'
'of cam_weights*activations')
parser.add_argument('--method', type=str, default='gradcam',
choices=[
'gradcam', 'hirescam', 'gradcam++',
'scorecam', 'xgradcam', 'ablationcam',
'eigencam', 'eigengradcam', 'layercam',
'fullgrad', 'gradcamelementwise'
],
help='CAM method')
parser.add_argument('--output-dir', type=str, default='autodl-tmp/mydata/cam/pytorch-grad-cam-master/result',
help='Output directory to save the images')
args = parser.parse_args()
args.use_cuda = args.use_cuda and torch.cuda.is_available()
if args.use_cuda:
print('Using GPU for acceleration')
else:
print('Using CPU for computation')
return args
代码中的需要可视化的图片的地址以及可视化后生成图片的地址需要重新填写。
主函数
以ResNet50为例:
if __name__ == '__main__':
""" python cam.py -image-path <path_to_image>
Example usage of loading an image and computing:
1. CAM
2. Guided Back Propagation
3. Combining both
"""
args = get_args()
methods = {
"gradcam": GradCAM,
"hirescam": HiResCAM,
"scorecam": ScoreCAM,
"gradcam++": GradCAMPlusPlus,
"ablationcam": AblationCAM,
"xgradcam": XGradCAM,
"eigencam": EigenCAM,
"eigengradcam": EigenGradCAM,
"layercam": LayerCAM,
"fullgrad": FullGrad,
"gradcamelementwise": GradCAMElementWise
}
model = models.resnet50(weights = models.ResNet50_Weights.DEFAULT)
需要换成自己模型的可以将上面的模型代码换成:
model = XXXX()
并使用下面代码加载自己模型的训练文件
weights_dict=torch.load("mydata/abs/parameters/UCMerced_checkpoint.pth", map_location='cpu')
model.load_state_dict(weights_dict, strict=False)
选择模型需要可视化的层
target_layers = [model.layer4]
#Resnet18 and 50: model.layer4
# VGG, densenet161: model.features[-1]
# mnasnet1_0: model.layers[-1]
# You can print the model to help chose the layer
# You can pass a list with several target layers,
# in that case the CAMs will be computed per layer and then aggregated.
# You can also try selecting all layers of a certain type, with e.g:
# from pytorch_grad_cam.utils.find_layers import find_layer_types_recursive
# find_layer_types_recursive(model, [torch.nn.ReLU])
处理图片
rgb_img = cv2.imread(args.image_path, 1)[:, :, ::-1]
rgb_img = np.float32(rgb_img) / 255
input_tensor = preprocess_image(rgb_img,
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
绘制图片并存储
targets = None
# Using the with statement ensures the context is freed, and you can
# recreate different CAM objects in a loop.
cam_algorithm = methods[args.method]
with cam_algorithm(model=model,
target_layers=target_layers,
use_cuda=args.use_cuda) as cam:
# AblationCAM and ScoreCAM have batched implementations.
# You can override the internal batch size for faster computation.
cam.batch_size = 4
grayscale_cam = cam(input_tensor=input_tensor,
targets=targets,
aug_smooth=args.aug_smooth,
eigen_smooth=args.eigen_smooth)
grayscale_cam = grayscale_cam[0, :]
cam_image = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True)
cam_image = cv2.cvtColor(cam_image, cv2.COLOR_RGB2BGR)
gb_model = GuidedBackpropReLUModel(model=model, use_cuda=args.use_cuda)
gb = gb_model(input_tensor, target_category=None)
cam_mask = cv2.merge([grayscale_cam, grayscale_cam, grayscale_cam])
cam_gb = deprocess_image(cam_mask * gb)
gb = deprocess_image(gb)
os.makedirs(args.output_dir, exist_ok=True)
cam_output_path = os.path.join(args.output_dir, f'{args.method}_resnet50cam.jpg')
gb_output_path = os.path.join(args.output_dir, f'{args.method}_resnet50gb.jpg')
cam_gb_output_path = os.path.join(args.output_dir, f'{args.method}_cam_resnet50gb.jpg')
cv2.imwrite(cam_output_path, cam_image)
cv2.imwrite(gb_output_path, gb)
cv2.imwrite(cam_gb_output_path, cam_gb)
一些可视化结果如下: