148 lines
5.4 KiB
Python
148 lines
5.4 KiB
Python
import argparse
|
|
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
|
|
|
|
|
|
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='./examples/both.png',
|
|
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 componenet'
|
|
'of cam_weights*activations')
|
|
parser.add_argument('--method', type=str, default='gradcam',
|
|
choices=['gradcam', 'hirescam', 'gradcam++',
|
|
'scorecam', 'xgradcam',
|
|
'ablationcam', 'eigencam',
|
|
'eigengradcam', 'layercam', 'fullgrad'],
|
|
help='Can be gradcam/gradcam++/scorecam/xgradcam'
|
|
'/ablationcam/eigencam/eigengradcam/layercam')
|
|
|
|
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
|
|
|
|
|
|
def api(image_path,method,model_name,**kwargs):
|
|
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(pretrained=True)
|
|
# model = eval('models.'+model_name+'(pretrained=True)')
|
|
print(model)
|
|
# Choose the target layer you want to compute the visualization for.
|
|
# Usually this will be the last convolutional layer in the model.
|
|
# Some common choices can be:
|
|
# 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])
|
|
#target_layers = [model.layer4]
|
|
target_layer=kwargs['target_layer']
|
|
target_layers = [eval(f'model.{target_layer}')]
|
|
|
|
rgb_img = cv2.imread(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])
|
|
|
|
# We have to specify the target we want to generate
|
|
# the Class Activation Maps for.
|
|
# If targets is None, the highest scoring category (for every member in the batch) will be used.
|
|
# You can target specific categories by
|
|
# targets = [e.g ClassifierOutputTarget(281)]
|
|
targets = None
|
|
|
|
# Using the with statement ensures the context is freed, and you can
|
|
# recreate different CAM objects in a loop.
|
|
cam_algorithm = methods[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 = 32
|
|
print(args.eigen_smooth)
|
|
aug_smooth=kwargs['aug_smooth']
|
|
grayscale_cam = cam(input_tensor=input_tensor,
|
|
targets=targets,
|
|
aug_smooth=aug_smooth,
|
|
eigen_smooth=args.eigen_smooth)
|
|
|
|
# Here grayscale_cam has only one image in the batch
|
|
grayscale_cam = grayscale_cam[0, :]
|
|
|
|
cam_image = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True)
|
|
|
|
# cam_image is RGB encoded whereas "cv2.imwrite" requires BGR encoding.
|
|
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)
|
|
|
|
cv2.imwrite(f'{method}_cam.jpg', cam_image)
|
|
cv2.imwrite(f'{method}_gb.jpg', gb)
|
|
cv2.imwrite(f'{method}_cam_gb.jpg', cam_gb)
|
|
return method+'_gb.jpg'
|
|
|
|
kwargs={"target_layer":'layer1',
|
|
"aug_smooth":True,
|
|
"eigen_smooth":True}
|
|
path=api('sample/both.png','fullgrad','resnet',**kwargs)
|
|
print(path) |