-
Notifications
You must be signed in to change notification settings - Fork 3
/
Copy pathgcam.py
131 lines (103 loc) · 4.33 KB
/
gcam.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
from PIL import Image
from imageio import imread
from skimage.io import imsave
from imgaug import imshow
from pytorch_grad_cam import GradCAM, ScoreCAM, GradCAMPlusPlus, AblationCAM, XGradCAM, EigenCAM
from pytorch_grad_cam.utils.image import show_cam_on_image
from skimage.io import show
from skimage.transform import resize
from torch.autograd import Variable
from torchvision import transforms
from torchvision.models import resnet50, mobilenetv3, mnasnet1_0, resnet18, resnet34, resnet152
from torch import from_numpy
from pathlib import Path
import numpy as np
from skimage.data import chelsea
from tqdm import tqdm
gradcam_dict = {
"gradcam": GradCAM,
"gradcam_plusplus": GradCAMPlusPlus,
"eigengradcam": EigenCAM
}
def gradcam_img(model, image, target_layer, cam_type, target_category):
#rgb_img = imread(image) / 255.
imsize = (224, 224)
loader = transforms.Compose([transforms.Scale(imsize), transforms.ToTensor()])
rgb_img = resize(imread(image) / 255., (224, 224))
def image_loader(image_name):
"""load image, returns cuda tensor"""
image = Image.open(image_name)
image = loader(image).float()
image = Variable(image, requires_grad=True)
image = image.unsqueeze(0) #this is for VGG, may not be needed for ResNet
return image.cuda() #assumes that you're using GPU
input_tensor = image_loader(image)
# 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_dict[cam_type](model=model, target_layer=target_layer, use_cuda=True)
# 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.
# You can also pass aug_smooth=True and eigen_smooth=True, to apply smoothing.
grayscale_cam = cam(input_tensor=input_tensor, target_category=target_category)
# In this example grayscale_cam has only one image in the batch:
grayscale_cam = grayscale_cam[0, :]
visualization = show_cam_on_image(rgb_img.astype(np.float32), grayscale_cam)
return visualization, resize(imread(image), (224, 224))
names = {
"dome": 897,
"altar": 677,
"church": 690,
"monastery": 686,
"palace": 685,
}
names = {value: key for key, value in names.items()}
target_categories = [None] + list(names.keys())
image = "./felix.jpeg"
savepath = Path("./heatmaps/")
cam_type = ["gradcam", "gradcam_plusplus", "eigengradcam"]
images = [str(file) for file in Path("./heatmap_src/").glob("**/*")]
model_names = ["ResNet18", "ResNet34", "ResNet50", "ResNet152", "MobileNetV3"]
def res18():
rn50 = resnet18(pretrained=True)
rn_target_layer = rn50.layer4[-1]
return rn50, rn_target_layer
def res34():
rn50 = resnet34(pretrained=True)
rn_target_layer = rn50.layer4[-1]
return rn50, rn_target_layer
def res50():
rn50 = resnet50(pretrained=True)
rn_target_layer = rn50.layer4[-1]
return rn50, rn_target_layer
def res152():
rn50 = resnet152(pretrained=True)
rn_target_layer = rn50.layer4[-1]
return rn50, rn_target_layer
def mobilenet():
model = mobilenetv3.mobilenet_v3_large(pretrained=True, progress=True)
target_layer = model.features[-1]
return model, target_layer
models = {
"ResNet18": res18,
"ResNet34": res34,
"ResNet50": res50,
"ResNet152": res152,
"MobileNetV3": mobilenet
}
from skimage import img_as_ubyte
for target_category in tqdm(target_categories):
for cam in cam_type:
for model_name in model_names:
sp = (savepath / cam / model_name)
sp.mkdir(exist_ok=True, parents=True)
model, target_layer = models[model_name]()
for i, img in enumerate(images):
r1 = sp / f"heatmap_{Path(img).name}" if target_category is None else sp / f"heatmap_{names[target_category]}_{Path(img).name}"
r2 = sp / f"original{Path(img).name}" if target_category is None else sp / f"original_{names[target_category]}_{Path(img).name}"
viz, im = gradcam_img(model, img, target_layer, cam, target_category)
#imshow(viz)
imsave(r1, img_as_ubyte(viz))
imsave(r2, img_as_ubyte(im))