239 lines
9.0 KiB
Python
239 lines
9.0 KiB
Python
|
import argparse
|
|||
|
|
import cv2
|
|||
|
|
import numpy as np
|
|||
|
|
import torch
|
|||
|
|
import json
|
|||
|
|
from torchvision import models
|
|||
|
|
from importlib import import_module
|
|||
|
|
from utils import LogHelper
|
|||
|
|
from pytorch_grad_cam import GradCAM, \
|
|||
|
|
HiResCAM, \
|
|||
|
|
ScoreCAM, \
|
|||
|
|
GradCAMPlusPlus, \
|
|||
|
|
AblationCAM, \
|
|||
|
|
XGradCAM, \
|
|||
|
|
EigenCAM, \
|
|||
|
|
EigenGradCAM, \
|
|||
|
|
LayerCAM, \
|
|||
|
|
FullGrad, \
|
|||
|
|
GradCAMElementWise
|
|||
|
|
from pytorch_grad_cam.utils.find_layers import find_layer_types_recursive
|
|||
|
|
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
|
|||
|
|
from PIL import Image
|
|||
|
|
from torchvision import transforms
|
|||
|
|
import urllib.request
|
|||
|
|
import json
|
|||
|
|
|
|||
|
|
|
|||
|
|
# Load labels
|
|||
|
|
with open("imagenet_1000.json") as f:
|
|||
|
|
labels = json.load(f)
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
available_device = 'cuda'
|
|||
|
|
def get_args():
|
|||
|
|
parser = argparse.ArgumentParser()
|
|||
|
|
parser.add_argument('--use-cuda', action='store_true', default=True,
|
|||
|
|
help='Use NVIDIA GPU acceleration')
|
|||
|
|
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')
|
|||
|
|
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
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ImageInterpretability():
|
|||
|
|
def __init__(self):
|
|||
|
|
super().__init__()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def perform(self,image_path: str, method: str, model_info: dict,output_path: str,log_path=None,**kwargs):
|
|||
|
|
# aug_smooth, eigen_smooth,
|
|||
|
|
args = get_args()
|
|||
|
|
'''
|
|||
|
|
图片输入地址
|
|||
|
|
:param image_path: (type=str, required=True) value= （图片地址）
|
|||
|
|
|
|||
|
|
可解释性算法
|
|||
|
|
:param method: (type=str, required=True) value=['gradcam', 'hirescam', 'scorecam', 'gradcam++',
|
|||
|
|
'ablationcam', 'xgradcam', 'eigencam', 'eigengradcam', 'layercam', 'gradcamelementwise','fullgrad'] （可解释性方法）
|
|||
|
|
|
|||
|
|
:param model_name: (type=str, required=True) value=[resnet, vgg, densenet, mnasnet] （模型名称）
|
|||
|
|
|
|||
|
|
:param model: (type=nn.Module, required=True) value=[resnet18, resnet50, vgg11, vgg13, vgg16, vgg19, densenet161, mnasnet1_0] （模型）
|
|||
|
|
|
|||
|
|
kwargs：非必要传入参数，在特定要求下传入
|
|||
|
|
{
|
|||
|
|
aug_smooth:默认采用数据增强技术来改善cam质量
|
|||
|
|
(type=bool) value=[Ture, False]
|
|||
|
|
|
|||
|
|
eigen_smooth：计算CAM（类激活映射）权重和激活之间的矩阵乘积，然后提取该结果的第一个主成分来减少结果中的噪音。
|
|||
|
|
(type=bool) value=[Ture, False]
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
'''
|
|||
|
|
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_class = getattr(models,model_info.model_name)
|
|||
|
|
# model = model_class(pretrained=True)
|
|||
|
|
model = self.get_model(info=model_info, device=available_device)
|
|||
|
|
# attack_log = LogHelper(log_path=log_path, root_log_name='aitest').build_new_log()
|
|||
|
|
|
|||
|
|
if 'resnet' in model_info.get('model_name').lower():
|
|||
|
|
target_layers = [model.layer4]
|
|||
|
|
elif 'vgg' in model_info.get('model_name').lower():
|
|||
|
|
target_layers = [model.features[-1]]
|
|||
|
|
elif 'densenet' in model_info.get('model_name').lower():
|
|||
|
|
target_layers = [model.features[-1]]
|
|||
|
|
elif 'mnasnet' in model_info.get('model_name').lower():
|
|||
|
|
target_layers = [model.layers[-1]]
|
|||
|
|
else:
|
|||
|
|
target_layers = find_layer_types_recursive(model, [torch.nn.ReLU])
|
|||
|
|
|
|||
|
|
|
|||
|
|
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])
|
|||
|
|
|
|||
|
|
|
|||
|
|
targets = None
|
|||
|
|
|
|||
|
|
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
|
|||
|
|
aug_smooth = kwargs['aug_smooth']
|
|||
|
|
eigen_smooth = kwargs['eigen_smooth']
|
|||
|
|
grayscale_cam = cam(input_tensor=input_tensor,
|
|||
|
|
targets=targets,
|
|||
|
|
aug_smooth=aug_smooth,
|
|||
|
|
eigen_smooth=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)
|
|||
|
|
|
|||
|
|
model.eval()
|
|||
|
|
image = Image.open(image_path)
|
|||
|
|
transform = transforms.Compose([
|
|||
|
|
transforms.Resize(256),
|
|||
|
|
transforms.CenterCrop(224),
|
|||
|
|
transforms.ToTensor(),
|
|||
|
|
transforms.Normalize(
|
|||
|
|
mean=[0.485, 0.456, 0.406],
|
|||
|
|
std=[0.229, 0.224, 0.225]
|
|||
|
|
)
|
|||
|
|
])
|
|||
|
|
image = transform(image)
|
|||
|
|
# Perform forward pass
|
|||
|
|
with torch.no_grad():
|
|||
|
|
output = model(image.unsqueeze(0))
|
|||
|
|
|
|||
|
|
# Convert output to probabilities
|
|||
|
|
probabilities = torch.softmax(output[0], dim=0)
|
|||
|
|
# Get predicted label and probability
|
|||
|
|
predicted_class = torch.argmax(probabilities).item()
|
|||
|
|
probability = probabilities[predicted_class].item()
|
|||
|
|
|
|||
|
|
# Get predicted label name
|
|||
|
|
predicted_class = torch.argmax(probabilities).item()
|
|||
|
|
label_name = labels[str(predicted_class)]
|
|||
|
|
#print(f"Predicted label: {label_name}")
|
|||
|
|
|
|||
|
|
#imwrite不支持带有中文路径的地址
|
|||
|
|
cv2.imwrite(output_path+label_name+'-'+method+"_cam3.jpg", cam_image)
|
|||
|
|
cv2.imwrite(output_path+label_name+'-'+method+"_gb.jpg", gb)
|
|||
|
|
cv2.imwrite(output_path+label_name+'-'+method+"_cam_gb.jpg",cam_gb)
|
|||
|
|
|
|||
|
|
return {'output_path':output_path,'Predicted_class': label_name, 'Probability': probability}
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def get_model(info, device):
|
|||
|
|
# if isinstance(info, dict):
|
|||
|
|
# info = argparse.Namespace(**info)
|
|||
|
|
#
|
|||
|
|
# if isinstance(info.path, dict):
|
|||
|
|
# if isinstance(info.path, str):
|
|||
|
|
# info.path = json.loads(info.path)
|
|||
|
|
|
|||
|
|
# load pytorch model from user upload files
|
|||
|
|
# if info.ownership != 'aitest' and ('upload' in info.type and 'pytorch' in info.type):
|
|||
|
|
# return load_pytorch_model(state_dict_path=info.path['parameter_file'], device=device, net_file_path=info.path['structure_file'])
|
|||
|
|
|
|||
|
|
# load built_in text model from textattack
|
|||
|
|
if 'torchvision' == info.get('source'):
|
|||
|
|
model_class = getattr(models, info.get('model_name'))
|
|||
|
|
return model_class(pretrained=True)
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
def load_pytorch_model(state_dict_path, model_class_name, device, net_file_path=None):
|
|||
|
|
"""
|
|||
|
|
model = load_user_model('word_cnn_for_classification',
|
|||
|
|
'C:\\Users\\pcl\\.cache\\textattack\\models\\classification\\cnn\\rotten-tomatoes\\pytorch_model.bin',
|
|||
|
|
'WordCNNForClassification')
|
|||
|
|
"""
|
|||
|
|
if net_file_path:
|
|||
|
|
net_file_path = str(net_file_path).replace('\\', '/')
|
|||
|
|
net_file_path = net_file_path.replace('./', '').replace('/', '.')
|
|||
|
|
model_class = getattr(import_module(net_file_path), model_class_name)
|
|||
|
|
model = model_class()
|
|||
|
|
model.load(state_dict_path, device)
|
|||
|
|
# tokenizer = model.tokenizer
|
|||
|
|
# model = textattack.models.wrappers.PyTorchModelWrapper(
|
|||
|
|
# model, tokenizer
|
|||
|
|
# )
|
|||
|
|
return model
|
|||
|
|
else:
|
|||
|
|
return torch.load(state_dict_path, device)
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|