-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathfeature_smoothing.py
176 lines (150 loc) · 6.95 KB
/
feature_smoothing.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
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
import numpy as np
import torch
from torchvision import models
from medpy.filter import anisotropic_diffusion
from cv2.ximgproc import anisotropicDiffusion
from cv2 import bilateralFilter
import torch.nn as nn
import util
from util import load_image, load_pkl_image
from torchsummary import summary
from matplotlib import pyplot as plt
from scipy.ndimage.filters import median_filter, convolve, gaussian_filter
import getpass
import time
from modified_curvature_motion import modified_curvature_motion
import warnings
from anisotropic_diffusion import anisotropic_diffusion_3d
def smooth(inputs, mode="anisotropic-diffusion", param=None):
"""
Different diffusion mode
:param inputs: input of a layer that needs to be smoothed, either raw images (numpy) or inter-features (tensor)
(note: "input" is a reserved key word in python)
:param mode: smoothing method
:param param: parameters of the smoothing method
:return: smoothed input
"""
is_feature = type(inputs) != np.ndarray # whether the input is an intermediate feature (tensor)
if is_feature:
inputs = inputs.cpu().numpy() # convert features from tensor to array
if mode == "anisotropic-diffusion":
# TODO: need to take care of the two sets of arguments on iteration number, medpy need around 10, cv2 need 4
param = (0.1, 20, 4) if param is None else param
if is_feature:
inputs = anisotropic_diffusion(inputs, kappa=param[1], niter=int(param[2]))
else:
# inputs = anisotropicDiffusion(np.transpose((inputs * 255).astype(np.uint8), (1, 2, 0)),
# alpha=param[0], K=param[1], niters=int(param[2]))
# inputs = np.transpose(inputs, (2, 0, 1)).astype("float32") / 255
inputs = anisotropic_diffusion_3d(input=inputs, niter=8, gamma=0.2, kappa=50)
elif mode == "mean":
param = (1, 1, 3, 3) if param is None else param
if is_feature:
inputs = convolve(inputs, weights=np.full(param, 1.0/27))
else:
inputs = convolve(inputs, weights=np.full(param[1:], 1.0/9))
inputs = ((inputs * 255).astype(np.uint8) / 255.0).astype(np.float32)
elif mode == "median":
param = (1, 1, 3, 3) if param is None else param
if is_feature:
inputs = median_filter(inputs, size=param)
else:
inputs = median_filter(inputs, size=param[1:])
inputs = ((inputs * 255).astype(np.uint8) / 255.0).astype(np.float32)
elif mode == "modified-curvature-motion":
if is_feature:
param = (1, 0.1) if param is None else param
inputs = inputs.squeeze()
inputs = modified_curvature_motion(inputs, niter=int(param[0]), k=param[1])
inputs = np.expand_dims(inputs, axis=0)
else:
param = (20, 0.9) if param is None else param
inputs = modified_curvature_motion(inputs, niter=int(param[0]), k=param[1])
inputs = ((inputs * 255).astype(np.uint8) / 255.0).astype(np.float32)
elif mode == "bilateral":
if not is_feature: # only works for raw images, [0, 1] works better than [0, 255]
param = (9, 75, 75) if param is None else param
inputs = np.transpose(bilateralFilter(np.transpose(inputs, (1, 2, 0)),
d=int(param[0]),
sigmaColor=(param[1]),
sigmaSpace=(param[2])),
(2, 0, 1))
elif mode == "gaussian":
param = (3,) if param is None else param
inputs = gaussian_filter(inputs, param[0]).astype("float32")
elif mode == 'quantization':
inputs = ((inputs * 255).astype(np.uint8) / 255.0).astype(np.float32)
elif mode == "customized":
pass
else:
warnings.warn("no such smoothing method")
quit()
if is_feature:
inputs = (torch.from_numpy(inputs)).cuda() # convert numpy back to tensor, for features only
return inputs
def forward_and_smooth(image, model, smooth_list=[], visual_list=[], mode="anisotropic-diffusion", param=None):
"""
Layer by layer operation for forward and smooth
:param image: input image, numpy array with shape (3 x 224 x 224)
:param model: pre-trained nn model
:param smooth_list: list of layers whose input shall be smoothed
:param visual_list: list of layers whose feature maps will be visualized
:param mode: smoothing mode
:param visual_flag: visualization flag
:param param: parameters of the smoothing method
:return: probability
"""
if -1 in smooth_list:
image = smooth(image, mode=mode, param=param)
# Normalize and convert to tensor
image = torch.unsqueeze(util.normalize_image_to_tensor(image), 0).cuda()
for p in model.parameters():
p.requires_grad = False
modules = list(model.children())
if len(visual_list) != 0:
plt.figure("features")
features = 0
for i in range(0, len(modules)-1):
layer = modules[i:i+1]
layer = nn.Sequential(*layer)
if i == 0:
features = layer(image)
else:
features = layer(features)
if i in smooth_list:
features = smooth(features, mode=mode, param=param)
if i in visual_list and len(visual_list) != 0:
plt.subplot(2, 4, i+1)
plt.title("layer{}".format(i))
plt.imshow(features[0, 20, :].cpu().numpy(), cmap="gray")
if len(visual_list) != 0:
plt.show()
resnet_last = modules[-1:]
resnet_last = nn.Sequential(*resnet_last)
prob = resnet_last(features.view(features.size(0), -1))
return prob.cpu().numpy()
if __name__ == "__main__":
# All directories must end with '/'
if getpass.getuser() == 'fantasie': # user is Yifei
ROOT_DIR = '/media/fantasie/backup/data/ILSVRC2012/' # root directory for ImageNet
else: # user is Chao
ROOT_DIR = '/home/chao/PycharmProjects/data/ILSVRC2012/'
img_path = ROOT_DIR + 'val_correct_adv_resnet152_pgd-0.01-0.002/n01440764/ILSVRC2012_val_00000293.JPEG'
if img_path.endswith('pkl'):
image = load_pkl_image(img_path)
else:
if 'resnet152' in img_path: # adversarial images, already resized
image = load_image(img_path, resize=False)
else: # clean images, need resizing
image = load_image(img_path, resize=True)
resnet152 = models.resnet152(pretrained=True).cuda().eval()
# summary(resnet152, (3, 224, 224))
# print(resnet152)
smooth_list = [-1, 4] # diffusion layers
visual_list = list(range(8)) # visualization layers
start = time.time()
prob = forward_and_smooth(image=image, model=resnet152, smooth_list=smooth_list,
visual_list=visual_list, mode="modified-curvature-motion")
print("Forward and smooth time: ", time.time() - start)
prediction = np.argmax(prob)
print(prediction)