-
Notifications
You must be signed in to change notification settings - Fork 0
/
anomaly_detection.py
355 lines (274 loc) · 10.4 KB
/
anomaly_detection.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
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
import argparse
import torch
from torch.nn import functional as F
from torch import nn
from torchvision import transforms, datasets
from torch.utils.data import Dataset, DataLoader
import cv2
import numpy as np
import os
import glob
import shutil
import time
from torchvision.models import resnet18
from PIL import Image
from sklearn.metrics import roc_auc_score
from torch import nn
import pytorch_lightning as pl
import string
import random
from sklearn.metrics import confusion_matrix
import ast
# dataset to load images
class MVTecDataset(Dataset):
"""Extract images from specified folder used for training model.
Args:
root: path to images
transform: transformation to apply to the images
wts: weights associated with each image
"""
def __init__(
self,
root,
transform,
wts,
):
# assign image paths to img_path
self.img_paths = glob.glob(root + "/*.png")
# transformation to apply to images
self.transform = transform
# weights to be assigned to images
# 9:1 when training not augmented samples and samples augmented not using template image
# 5:1:2 when training not augmented samples, samples augmented not using template image, samples trained using template images
if wts == "without-template":
self.wts = [9, 1]
else:
self.wts = [5, 1, 2]
def __len__(self):
return len(self.img_paths)
def __getitem__(self, idx):
# get image path and ground truth for particular index
img_path = self.img_paths[idx]
# using weighted
wt = self.wts[0]
if "stage1" in img_path:
# weights for training images generated using OK images
wt = self.wts[1]
elif "stage2" in img_path:
# weights for training images generated using OK and Template images
wt = self.wts[2]
# read the image and apply transformations
img = Image.open(img_path).convert("RGB")
img = self.transform(img)
# create ground truth labels
gt = torch.zeros([1, img.size()[-2], img.size()[-2]])
return (img, gt, wt)
class STPM(pl.LightningModule):
"""A teacher with pretrained weights and a student trained to extract
features similar to teacher(teacher is not optimized) on non-defect images, will produce different
features when images containing defect will be fed to them.
Args:
dataset_path: path to images
load_size: size of images
lr: learning rate of optimizer used
momentum: momentum of optimizer used
batch_size: number of images used to calculate gradients
weight_decay: learning rate decay
amap_mode: method of combining the similarity scores of different features maps default multiply each of them
"""
def __init__(
self,
dataset_path,
teacher_weights=None,
wts="without-template",
load_size=768,
lr=0.4,
momentum=0.9,
batch_size=32,
weight_decay=1e-4,
amap_mode="mul",
):
super(STPM, self).__init__()
self.save_hyperparameters()
self.amap_mode = amap_mode
self.teacher_weights = teacher_weights
self.load_size = load_size
self.weight_decay = weight_decay
self.lr = lr
self.wts = wts
self.momentum = momentum
self.dataset_path = dataset_path
self.batch_size = batch_size
# hook_t and hook_s append the extracted features to features_t and features_s respectively
def hook_t(module, input, output):
self.features_t.append(output)
def hook_s(module, input, output):
self.features_s.append(output)
self.model_t = resnet18(pretrained=True).eval()
if self.teacher_weights != None:
print("Teacher weights loaded")
self.model_t.load_state_dict(
torch.load(self.teacher_weights)["model_state_dict"], strict=False
)
for param in self.model_t.parameters():
param.requires_grad = False
# select layers of teacher and student models to extract features from
self.model_t.layer1[-1].register_forward_hook(hook_t)
self.model_t.layer2[-1].register_forward_hook(hook_t)
self.model_t.layer3[-1].register_forward_hook(hook_t)
self.model_s = resnet18(pretrained=False)
self.model_s.layer1[-1].register_forward_hook(hook_s)
self.model_s.layer2[-1].register_forward_hook(hook_s)
self.model_s.layer3[-1].register_forward_hook(hook_s)
# we are using the Mean Squared Error loss
self.criterion = torch.nn.MSELoss(reduce=False)
# statistics as calculated on imagenet dataset
self.data_transforms = transforms.Compose(
[
transforms.Resize((self.load_size, self.load_size), Image.ANTIALIAS),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
),
]
)
self.inv_normalize = transforms.Normalize(
mean=[-0.485 / 0.229, -0.456 / 0.224, -0.406 / 0.255],
std=[1 / 0.229, 1 / 0.224, 1 / 0.255],
)
# lists to store extracted features
def init_features(self):
self.features_t = []
self.features_s = []
# pass the image through the teacher and student models
def forward(self, x):
# reset the lists to store features
self.init_features()
x_t = self.model_t(x)
x_s = self.model_s(x)
return (self.features_t, self.features_s)
# weighted mean squared error for optimizing student model
def cal_loss(
self,
fs_list,
ft_list,
wts,
):
tot_loss = 0
for i in range(len(ft_list)):
fs = fs_list[i]
ft = ft_list[i]
(_, _, h, w) = fs.shape
# normalize the fetaures so both are within same scale
fs_norm = F.normalize(fs, p=2)
ft_norm = F.normalize(ft, p=2)
# calculate weighted loss as the difference between the features extracted by the teacher and student models
f_loss = (
0.5
/ (w * h)
* torch.sum(
float(wts) * torch.sum(self.criterion(fs_norm, ft_norm), [1, 2, 3]),
0,
)
)
tot_loss += f_loss
return tot_loss
# configure optimizer
def configure_optimizers(self):
return torch.optim.SGD(
self.model_s.parameters(),
lr=self.lr,
momentum=self.momentum,
weight_decay=self.weight_decay,
)
# create training dataloader
def train_dataloader(self):
image_datasets = MVTecDataset(
root=self.dataset_path, transform=self.data_transforms, wts=self.wts
)
train_loader = DataLoader(
image_datasets, batch_size=self.batch_size, shuffle=True, num_workers=0
)
return train_loader
# extract features from student and teacher model and optimize student model using weighted Mean Squared error loss
def training_step(self, batch, batch_idx):
# get image and weights used for current index
(x, _, wts) = batch
# pass the image through teacher and student model
(features_t, features_s) = self(x)
# calculate and return loss
loss = self.cal_loss(features_s, features_t, wts)
return loss
# normalize the extracted features and calculate their cosine similarity and combine either by adding them or multiplying
# each of the feature maps
def cal_anomaly_map(
self,
fs_list,
ft_list,
out_size=224,
):
if self.amap_mode == "mul":
anomaly_map = np.ones([out_size, out_size])
else:
anomaly_map = np.zeros([out_size, out_size])
a_map_list = []
for i in range(len(ft_list)):
fs = fs_list[i]
ft = ft_list[i]
fs_norm = F.normalize(fs, p=2)
ft_norm = F.normalize(ft, p=2)
a_map = 1 - F.cosine_similarity(fs_norm, ft_norm)
a_map = torch.unsqueeze(a_map, dim=1)
a_map = F.interpolate(a_map, size=out_size, mode="bilinear")
a_map = a_map[0, 0, :, :].to("cpu").detach().numpy()
a_map_list.append(a_map)
if self.amap_mode == "mul":
anomaly_map *= a_map
else:
anomaly_map += a_map
return (anomaly_map, a_map_list)
def get_args():
parser = argparse.ArgumentParser(description="ANOMALYDETECTION")
parser.add_argument("--teacher_weights", default=None)
parser.add_argument("--phase", choices=["train", "test"], default="train")
parser.add_argument("--dataset_path", default=r"D:\Dataset\mvtec_anomaly_detection")
parser.add_argument("--category", default="template")
parser.add_argument("--num_epochs", default=100)
parser.add_argument("--wts", default="without-template")
parser.add_argument("--lr", default=0.4)
parser.add_argument("--momentum", default=0.9)
parser.add_argument("--weight_decay", default=1e-4)
parser.add_argument("--batch_size", default=32)
parser.add_argument("--load_size", default=768)
parser.add_argument(
"--project_path",
default=r"D:\Project_Train_Results\mvtec_anomaly_detection\210624\test",
)
parser.add_argument("--amap_mode", choices=["mul", "sum"], default="mul")
args = parser.parse_args()
return args
if __name__ == "__main__":
# check available devices and assign cuda if available
device = torch.device(("cuda" if torch.cuda.is_available() else "cpu"))
# get assigned parameters
args = get_args()
# create trainer
trainer = pl.Trainer.from_argparse_args(
args,
default_root_dir=os.path.join(args.project_path, args.category),
max_epochs=int(args.num_epochs),
gpus=[0],
)
# create and fit the model
model = STPM(
dataset_path=args.dataset_path,
teacher_weights=args.teacher_weights,
load_size=args.load_size,
wts=args.wts,
lr=args.lr,
momentum=args.momentum,
batch_size=int(args.batch_size),
weight_decay=args.weight_decay,
amap_mode=args.amap_mode,
)
trainer.fit(model)