diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..e69de29 diff --git a/Readme.md b/Readme.md index c614482..6c40d22 100644 --- a/Readme.md +++ b/Readme.md @@ -20,14 +20,49 @@ Install packages with: $ pip install -r requirements.txt ``` +Or install with for Windows as per [PyTorch official site](https://pytorch.org/get-started/locally/): + +``` +$ pip install torch===1.6.0 torchvision===0.7.0 -f https://download.pytorch.org/whl/torch_s +table.html +$ pip install -r requirements.txt +``` + ## Configure and Run All configurations concerning data, model, training, visualization etc. can be made in _config.py_. The default configuration will run a training with paper-given parameters on the provided dummy dataset. This dataset contains images of 4 squares as normal examples and 4 circles as anomaly. -To start the training, just run _main.py_! If training on the dummy data does not lead to an AUROC of 1.0, something seems to be wrong. +If you encounter GPU Out of Memory issue, you can reduce the neuron numbers in _config.py_ +``` +fc_internal = 1536 # number of neurons in hidden layers of s-t-networks +``` + +To start the training, just run _main.py_ as follows! If training on the dummy data does not lead to an AUROC of 1.0, something seems to be wrong. Please report us if you have issues when using the code. +``` +$ python main.py +``` + ## Data +How to use Data extraction tool to extract data from video clips: + 1. Create folder structure like the example shows in the picture below. + + ![1](https://github.com/zerobox-ai/differnet/blob/zijian/dataset/data-generation/annotations/structure1.png) + + 2. Dump the videos and annotations (rename them use 1.xml, 1.avi as one pair annotation and video) into the folders under data-generation folder. + + ![2](https://github.com/zerobox-ai/differnet/blob/zijian/dataset/data-generation/annotations/structure2.png) + + 3. Modify the annotation files: Since the annotation uses label "defect" to indicate the defect area, while, both good and defective bottles are labeled as "bottle" which is confusing. To indicate which "bottle" is defective, we need to find the frames that labeled with defect, and then manully update the group's label from "bottle" to "defective" for the groups that falling in to those frames. + ![3](https://github.com/zerobox-ai/differnet/blob/zijian/dataset/data-generation/annotations/structure3.png) + + - For example: in the example image above, the frame 15 and 16 are labeled as "defect" which indicates those 2 frames has defect areas on the bottles. So we need to find the group that contains frame 15 and 16, and then manully update the label from "bottle" to "defective". and then delete the whole \ group that labeled as "defect" (since we don't care about the defect area in data extraction). + + 4. Modify the config.py, fill in appropriate value for num_videos, save_cropped_image_to and save_original_image_to + + 5. run the data extraction: python data_extraction.py + The given dummy dataset shows how the implementation expects the construction of a dataset. Coincidentally, the [MVTec AD dataset](https://www.mvtec.com/de/unternehmen/forschung/datasets/mvtec-ad/) is constructed in this way. diff --git a/apply_mask.py b/apply_mask.py new file mode 100644 index 0000000..8a0362d --- /dev/null +++ b/apply_mask.py @@ -0,0 +1,22 @@ +import cv2 +import os + + +def load_images_from_folder(folder): + images = [] + for filename in os.listdir(folder): + img = cv2.imread(os.path.join(folder, filename)) + if img is not None: + images.append(img) + return images + + +path = 'dataset/Experiment_4.1/validate/good' +mask = cv2.imread(os.path.join('dataset/Mask/', 'Mask_shrink.jpg')) +mask = mask / 255 # make the mask into 0/1 matrix for multiplication +imgs = load_images_from_folder(path) + +for i, img in enumerate(imgs): + img = cv2.resize(img, (400, 700), interpolation=cv2.INTER_AREA) + masked_img = img * mask + cv2.imwrite('dataset/Experiment_5.1/validate/good/good-Masked-' + str(i) + '.jpg', masked_img) diff --git a/config.py b/config.py index a650dba..d91230f 100644 --- a/config.py +++ b/config.py @@ -1,49 +1,65 @@ -'''This file configures the training procedure because handling arguments in every single function is so exhaustive for -research purposes. Don't try this code if you are a software engineer.''' - -# device settings -device = 'cuda' # or 'cpu' -import torch -torch.cuda.set_device(0) - -# data settings -dataset_path = "dummy_dataset" -class_name = "dummy_class" -modelname = "dummy_test" - -img_size = (448, 448) -img_dims = [3] + list(img_size) -add_img_noise = 0.01 - -# transformation settings -transf_rotations = True -transf_brightness = 0.0 -transf_contrast = 0.0 -transf_saturation = 0.0 -norm_mean, norm_std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225] - -# network hyperparameters -n_scales = 3 # number of scales at which features are extracted, img_size is the highest - others are //2, //4,... -clamp_alpha = 3 # see paper equation 2 for explanation -n_coupling_blocks = 8 -fc_internal = 2048 # number of neurons in hidden layers of s-t-networks -dropout = 0.0 # dropout in s-t-networks -lr_init = 2e-4 -n_feat = 256 * n_scales # do not change except you change the feature extractor - -# dataloader parameters -n_transforms = 4 # number of transformations per sample in training -n_transforms_test = 64 # number of transformations per sample in testing -batch_size = 24 # actual batch size is this value multiplied by n_transforms(_test) -batch_size_test = batch_size * n_transforms // n_transforms_test - -# total epochs = meta_epochs * sub_epochs -# evaluation after epochs -meta_epochs = 1 -sub_epochs = 8 - -# output settings -verbose = True -grad_map_viz = True -hide_tqdm_bar = True -save_model = True +'''This file configures the training procedure because handling arguments in every single function is so exhaustive for +research purposes. Don't try this code if you are a software engineer.''' + +# data extraction settings +num_videos = 21 +save_cropped_image_to = "dataset/zerobox-2010-1/" +save_original_image_to = "dataset/zerobox-2010-1-original/" + +# device settings +device = 'cuda' # 'cuda' or 'cpu' +import torch +torch.cuda.set_device(0) + +# data settings +dataset_path = "dataset" +class_name = "Experiment_6.1" +modelname = "Experiment_6.1_10epoch_239tainingdata_0.5BCS" + +img_size = (448, 448) +img_dims = [3] + list(img_size) +add_img_noise = 0.01 + +# transformation settings +transf_rotations = True +transf_brightness = 0.5 +transf_contrast = 0.5 +transf_saturation = 0.5 +norm_mean, norm_std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225] + +rotation_degree = 0 +crop_top = 0.10 +crop_left = 0.10 +crop_bottom = 0.10 +crop_right = 0.10 + +# network hyperparameters +n_scales = 3 # number of scales at which features are extracted, img_size is the highest - others are //2, //4,... +clamp_alpha = 3 # see paper equation 2 for explanation +n_coupling_blocks = 8 +# fc_internal = 2048 # number of neurons in hidden layers of s-t-networks +fc_internal = 1536 # number of neurons in hidden layers of s-t-networks +dropout = 0.0 # dropout in s-t-networks +lr_init = 2e-4 +n_feat = 256 * n_scales # do not change except you change the feature extractor + +# dataloader parameters +n_transforms = 4 # number of transformations per sample in training +n_transforms_test = 1 # number of transformations per sample in testing +batch_size = 4 # actual batch size is this value multiplied by n_transforms(_test) +batch_size_test = 1 + +# total epochs = meta_epochs * sub_epochs +# evaluation after epochs +meta_epochs = 5 +sub_epochs = 8 + +# output settings +verbose = True +grad_map_viz = True +hide_tqdm_bar = True +save_model = True +save_transformed_image = True +visualization = False +frame_name_is_given = False +target_tpr = 0.85 diff --git a/data_extraction.py b/data_extraction.py new file mode 100644 index 0000000..b0bb99f --- /dev/null +++ b/data_extraction.py @@ -0,0 +1,74 @@ +import cv2 +from xml.dom import minidom +import config as c + +# Load videos one by one +for i in range(c.num_videos): + print('Data generation on video-' + str(i+1)) + filename = str(i+1) + + # Opens the Video file + cap = cv2.VideoCapture('dataset/data-generation/videos/' + filename + '.avi') + + # Read annotations + annotation = minidom.parse('dataset/data-generation/annotations/' + filename + '.xml') + boxes = annotation.getElementsByTagName('box') + + frameList = [] + labelList = [] + boxesList = [] + + # Store the bounding box info along with frame number info into list + for i in range(boxes.length): + + # make sure not select the bounding box that outside the frame + if (boxes[i].attributes['outside'].value != '1'): + frame = int(boxes[i].attributes['frame'].value) + frameList.append(frame) + + labelList.append(boxes[i].parentNode.attributes['label'].value) + + ytl = int(float(boxes[i].attributes['ytl'].value)) + ybr = int(float(boxes[i].attributes['ybr'].value)) + xtl = int(float(boxes[i].attributes['xtl'].value)) + xbr = int(float(boxes[i].attributes['xbr'].value)) + boxesList.append([ytl, ybr, xtl, xbr]) + + # Set up shrink percentage + shrink_percentage = 0.02 + j = 0 + while(cap.isOpened()): + ret, frame = cap.read() + if(frame is not None and j in frameList): + ytl = boxesList[frameList.index(j)][0] + ybr = boxesList[frameList.index(j)][1] + xtl = boxesList[frameList.index(j)][2] + xbr = boxesList[frameList.index(j)][3] + label = 'good' if labelList[frameList.index(j)] == 'bottle' else 'defect' + + # draw bounding box on original frames + if label != 'defect': + cv2.rectangle(frame, (xtl, ytl), (xbr, ybr), (0, 255, 0), 5) + else: + cv2.rectangle(frame, (xtl, ytl), (xbr, ybr), (0, 0, 255), 5) + #cv2.imshow("Show", frame) + #cv2.waitKey() + #cv2.destroyAllWindows() + # Crop the frames with the bounding box position info + crop_frame = frame[int(ytl*(1+shrink_percentage)):int(ybr*(1-shrink_percentage)), + int(xtl*(1+shrink_percentage)):int(xbr*(1-shrink_percentage))] + + # output file formatting example "video1-frame4-defect.jpg" + print('Successfully generated: ' +c.save_cropped_image_to + label + '/video-' + filename + '-frame' + str(j) + + '-' + label + '.jpg') + cv2.imwrite(c.save_original_image_to + label + '/original-video-' + filename + '-frame' + str(j) + '-' + label + '.jpg', + frame) + cv2.imwrite(c.save_cropped_image_to + label + '/video-' + filename + '-frame' + str(j) + '-' + label + '.jpg', + crop_frame) + + if ret == False: + break + j += 1 + + cap.release() + cv2.destroyAllWindows() diff --git a/dataset/data-generation/annotations/1.xml b/dataset/data-generation/annotations/1.xml new file mode 100644 index 0000000..cee4ba3 --- /dev/null +++ b/dataset/data-generation/annotations/1.xml @@ -0,0 +1,438 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/dataset/data-generation/annotations/structure1.png b/dataset/data-generation/annotations/structure1.png new file mode 100644 index 0000000..18c72f6 Binary files /dev/null and b/dataset/data-generation/annotations/structure1.png differ diff --git a/dataset/data-generation/annotations/structure2.png b/dataset/data-generation/annotations/structure2.png new file mode 100644 index 0000000..ec5ae62 Binary files /dev/null and b/dataset/data-generation/annotations/structure2.png differ diff --git a/dataset/data-generation/annotations/structure3.png b/dataset/data-generation/annotations/structure3.png new file mode 100644 index 0000000..1ce82d5 Binary files /dev/null and b/dataset/data-generation/annotations/structure3.png differ diff --git a/drawMask.py b/drawMask.py new file mode 100644 index 0000000..fe1f253 --- /dev/null +++ b/drawMask.py @@ -0,0 +1,62 @@ +import cv2 +import os +import numpy as np + + +def load_images_from_folder(folder): + images = [] + for filename in os.listdir(folder): + image = cv2.imread(os.path.join(folder, filename)) + if image is not None: + images.append(image) + return images + + +imgs = load_images_from_folder('dataset/bgm/') +output_height = 700 +output_width = 400 + +for i, img in enumerate(imgs): + print('Original Dimensions : ', img.shape) + resized = cv2.resize(img, (output_width, output_height), interpolation=cv2.INTER_AREA) + masked = [] + + for r in resized: + new_c = [] + for c in r: + # check if the pixel value is green mask or original image pixels + if True in (abs([120, 255, 155] - c) > [30, 30, 30]): + new_c.append([0, 0, 0]) + else: + # keep the green mask pixels in the img + new_c.append([1, 1, 1]) + masked.append(new_c) + masked = np.array(masked, dtype=np.uint8) + if i == 0: + mask = masked + else: + mask = np.ceil((mask + masked) / 2) + + # output the results of each iteration after mask addition process + #cv2.imwrite('dataset/Mask/Mask-' + str(i) + '.jpg', 255 - (255 * mask)) + #print('Resized Dimensions : ', resized.shape) + +# output final mask addition result +cv2.imwrite('dataset/Mask/Mask.jpg', 255 - (255 * mask)) + +# shrink the mask area +shrink_percentage = 0.1 # shrink the mask by percentage from 4 orientations (top, bottom, left and right) +shrank_height = output_height * (1 - 2 * shrink_percentage) # shrink top and bottom +shrank_width = output_width * (1 - 2 * shrink_percentage) # shrink left and right +shrank_mask = cv2.resize(mask, (int(shrank_width), int(shrank_height)), interpolation=cv2.INTER_AREA) +# extend the border of the shrank_mask +shrank_mask = cv2.copyMakeBorder( + shrank_mask, + top=int(output_height * shrink_percentage), + bottom=int(output_height * shrink_percentage), + left=int(output_width * shrink_percentage), + right=int(output_width * shrink_percentage), + borderType=cv2.BORDER_CONSTANT, + value=[255, 255, 255] +) +cv2.imwrite('dataset/Mask/Mask_shrink.jpg', 255 - (255 * shrank_mask)) diff --git a/flows.py b/flows.py new file mode 100644 index 0000000..9af567c --- /dev/null +++ b/flows.py @@ -0,0 +1,528 @@ +import math +import types + +import numpy as np +import scipy as sp +import scipy.linalg +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def get_mask(in_features, out_features, in_flow_features, mask_type=None): + """ + mask_type: input | None | output + + See Figure 1 for a better illustration: + https://arxiv.org/pdf/1502.03509.pdf + """ + if mask_type == 'input': + in_degrees = torch.arange(in_features) % in_flow_features + else: + in_degrees = torch.arange(in_features) % (in_flow_features - 1) + + if mask_type == 'output': + out_degrees = torch.arange(out_features) % in_flow_features - 1 + else: + out_degrees = torch.arange(out_features) % (in_flow_features - 1) + + return (out_degrees.unsqueeze(-1) >= in_degrees.unsqueeze(0)).float() + + +class MaskedLinear(nn.Module): + def __init__(self, + in_features, + out_features, + mask, + cond_in_features=None, + bias=True): + super(MaskedLinear, self).__init__() + self.linear = nn.Linear(in_features, out_features) + if cond_in_features is not None: + self.cond_linear = nn.Linear( + cond_in_features, out_features, bias=False) + + self.register_buffer('mask', mask) + + def forward(self, inputs, cond_inputs=None): + output = F.linear(inputs, self.linear.weight * self.mask, + self.linear.bias) + if cond_inputs is not None: + output += self.cond_linear(cond_inputs) + return output + + +nn.MaskedLinear = MaskedLinear + + +class MADESplit(nn.Module): + """ An implementation of MADE + (https://arxiv.org/abs/1502.03509). + """ + + def __init__(self, + num_inputs, + num_hidden, + num_cond_inputs=None, + s_act='tanh', + t_act='relu', + pre_exp_tanh=False): + super(MADESplit, self).__init__() + + self.pre_exp_tanh = pre_exp_tanh + + activations = {'relu': nn.ReLU, 'sigmoid': nn.Sigmoid, 'tanh': nn.Tanh} + + input_mask = get_mask(num_inputs, num_hidden, num_inputs, + mask_type='input') + hidden_mask = get_mask(num_hidden, num_hidden, num_inputs) + output_mask = get_mask(num_hidden, num_inputs, num_inputs, + mask_type='output') + + act_func = activations[s_act] + self.s_joiner = nn.MaskedLinear(num_inputs, num_hidden, input_mask, + num_cond_inputs) + + self.s_trunk = nn.Sequential(act_func(), + nn.MaskedLinear(num_hidden, num_hidden, + hidden_mask), act_func(), + nn.MaskedLinear(num_hidden, num_inputs, + output_mask)) + + act_func = activations[t_act] + self.t_joiner = nn.MaskedLinear(num_inputs, num_hidden, input_mask, + num_cond_inputs) + + self.t_trunk = nn.Sequential(act_func(), + nn.MaskedLinear(num_hidden, num_hidden, + hidden_mask), act_func(), + nn.MaskedLinear(num_hidden, num_inputs, + output_mask)) + + def forward(self, inputs, cond_inputs=None, rev=False): + if rev == False: + h = self.s_joiner(inputs, cond_inputs) + m = self.s_trunk(h) + + h = self.t_joiner(inputs, cond_inputs) + a = self.t_trunk(h) + + if self.pre_exp_tanh: + a = torch.tanh(a) + + u = (inputs - m) * torch.exp(-a) + return u, -a.sum(-1, keepdim=True) + + else: + x = torch.zeros_like(inputs) + for i_col in range(inputs.shape[1]): + h = self.s_joiner(x, cond_inputs) + m = self.s_trunk(h) + + h = self.t_joiner(x, cond_inputs) + a = self.t_trunk(h) + + if self.pre_exp_tanh: + a = torch.tanh(a) + + x[:, i_col] = inputs[:, i_col] * torch.exp( + a[:, i_col]) + m[:, i_col] + return x, -a.sum(-1, keepdim=True) + +class MADE(nn.Module): + """ An implementation of MADE + (https://arxiv.org/abs/1502.03509). + """ + + def __init__(self, + num_inputs, + num_hidden, + num_cond_inputs=None, + act='relu', + pre_exp_tanh=False): + super(MADE, self).__init__() + + activations = {'relu': nn.ReLU, 'sigmoid': nn.Sigmoid, 'tanh': nn.Tanh} + act_func = activations[act] + + input_mask = get_mask( + num_inputs, num_hidden, num_inputs, mask_type='input') + hidden_mask = get_mask(num_hidden, num_hidden, num_inputs) + output_mask = get_mask( + num_hidden, num_inputs * 2, num_inputs, mask_type='output') + + self.joiner = nn.MaskedLinear(num_inputs, num_hidden, input_mask, + num_cond_inputs) + + self.trunk = nn.Sequential(act_func(), + nn.MaskedLinear(num_hidden, num_hidden, + hidden_mask), act_func(), + nn.MaskedLinear(num_hidden, num_inputs * 2, + output_mask)) + + def forward(self, inputs, cond_inputs=None, rev=False): + if rev == False: + h = self.joiner(inputs, cond_inputs) + m, a = self.trunk(h).chunk(2, 1) + u = (inputs - m) * torch.exp(-a) + return u, -a.sum(-1, keepdim=True) + + else: + x = torch.zeros_like(inputs) + for i_col in range(inputs.shape[1]): + h = self.joiner(x, cond_inputs) + m, a = self.trunk(h).chunk(2, 1) + x[:, i_col] = inputs[:, i_col] * torch.exp( + a[:, i_col]) + m[:, i_col] + return x, -a.sum(-1, keepdim=True) + + +class Sigmoid(nn.Module): + def __init__(self): + super(Sigmoid, self).__init__() + + def forward(self, inputs, cond_inputs=None, rev=False): + if rev == False: + s = torch.sigmoid + return s(inputs), torch.log(s(inputs) * (1 - s(inputs))).sum( + -1, keepdim=True) + else: + return torch.log(inputs / + (1 - inputs)), -torch.log(inputs - inputs**2).sum( + -1, keepdim=True) + + +class Logit(Sigmoid): + def __init__(self): + super(Logit, self).__init__() + + def forward(self, inputs, cond_inputs=None, rev=False): + if rev == False: + return super(Logit, self).forward(inputs, 'inverse') + else: + return super(Logit, self).forward(inputs, 'direct') + + + +class BatchNormFlow(nn.Module): + """ An implementation of a batch normalization layer from + Density estimation using Real NVP + (https://arxiv.org/abs/1605.08803). + """ + + def __init__(self, num_inputs, momentum=0.0, eps=1e-5): + super(BatchNormFlow, self).__init__() + + num_inputs = num_inputs[0][0] + self.log_gamma = nn.Parameter(torch.zeros(num_inputs)) + self.beta = nn.Parameter(torch.zeros(num_inputs)) + self.momentum = momentum + self.eps = eps + + self.register_buffer('running_mean', torch.zeros(num_inputs)) + self.register_buffer('running_var', torch.ones(num_inputs)) + + def forward(self, inputs, cond_inputs=None, rev=False): + if rev == False: + if self.training: + inputs = inputs[0] + #inputs = torch.Tensor(inputs) + #inputs = np.array(inputs) + self.batch_mean = inputs.mean(0) + self.batch_var = ( + inputs - self.batch_mean).pow(2).mean(0) + self.eps + + self.running_mean.mul_(self.momentum) + self.running_var.mul_(self.momentum) + + self.running_mean.add_(self.batch_mean.data * + (1 - self.momentum)) + self.running_var.add_(self.batch_var.data * + (1 - self.momentum)) + + mean = self.batch_mean + var = self.batch_var + else: + mean = self.running_mean + var = self.running_var + + x_hat = (inputs - mean) / var.sqrt() + y = torch.exp(self.log_gamma) * x_hat + self.beta + return y, (self.log_gamma - 0.5 * torch.log(var)).sum( + -1, keepdim=True) + else: + if self.training: + mean = self.batch_mean + var = self.batch_var + else: + mean = self.running_mean + var = self.running_var + + x_hat = (inputs - self.beta) / torch.exp(self.log_gamma) + + y = x_hat * var.sqrt() + mean + + return y, (-self.log_gamma + 0.5 * torch.log(var)).sum( + -1, keepdim=True) + + def output_dims(self, input_dims): + assert len(input_dims) == 1, "Can only use 1 input" + return input_dims + +class ActNorm(nn.Module): + """ An implementation of a activation normalization layer + from Glow: Generative Flow with Invertible 1x1 Convolutions + (https://arxiv.org/abs/1807.03039). + """ + + def __init__(self, num_inputs): + super(ActNorm, self).__init__() + num_inputs = num_inputs[0][0] + self.weight = nn.Parameter(torch.ones(num_inputs)) + self.bias = nn.Parameter(torch.zeros(num_inputs)) + self.initialized = False + + def forward(self, inputs, cond_inputs=None, rev=False): + inputs = inputs[0] + if self.initialized == False: + self.weight.data.copy_(torch.log(1.0 / (inputs.std(0) + 1e-12))) + self.bias.data.copy_(inputs.mean(0)) + self.initialized = True + + if rev == False: + return ( + inputs - self.bias) * torch.exp(self.weight), self.weight.sum( + -1, keepdim=True).unsqueeze(0).repeat(inputs.size(0), 1) + else: + return inputs * torch.exp( + -self.weight) + self.bias, -self.weight.sum( + -1, keepdim=True).unsqueeze(0).repeat(inputs.size(0), 1) + + def output_dims(self, input_dims): + assert len(input_dims) == 1, "Can only use 1 input" + return input_dims + + +class InvertibleMM(nn.Module): + """ An implementation of a invertible matrix multiplication + layer from Glow: Generative Flow with Invertible 1x1 Convolutions + (https://arxiv.org/abs/1807.03039). + """ + + def __init__(self, num_inputs): + super(InvertibleMM, self).__init__() + self.W = nn.Parameter(torch.Tensor(num_inputs, num_inputs)) + nn.init.orthogonal_(self.W) + + def forward(self, inputs, cond_inputs=None, rev=False): + if rev == False: + return inputs @ self.W, torch.slogdet( + self.W)[-1].unsqueeze(0).unsqueeze(0).repeat( + inputs.size(0), 1) + else: + return inputs @ torch.inverse(self.W), -torch.slogdet( + self.W)[-1].unsqueeze(0).unsqueeze(0).repeat( + inputs.size(0), 1) + + +class LUInvertibleMM(nn.Module): + """ An implementation of a invertible matrix multiplication + layer from Glow: Generative Flow with Invertible 1x1 Convolutions + (https://arxiv.org/abs/1807.03039). + """ + + def __init__(self, num_inputs): + super(LUInvertibleMM, self).__init__() + num_inputs = num_inputs[0][0] + self.W = torch.Tensor(num_inputs, num_inputs) + nn.init.orthogonal_(self.W) + self.L_mask = torch.tril(torch.ones(self.W.size()), -1) + self.U_mask = self.L_mask.t().clone() + + P, L, U = sp.linalg.lu(self.W.numpy()) + self.P = torch.from_numpy(P) + self.L = nn.Parameter(torch.from_numpy(L)) + self.U = nn.Parameter(torch.from_numpy(U)) + + S = np.diag(U) + sign_S = np.sign(S) + log_S = np.log(abs(S)) + self.sign_S = torch.from_numpy(sign_S) + self.log_S = nn.Parameter(torch.from_numpy(log_S)) + + self.I = torch.eye(self.L.size(0)) + + def forward(self, inputs, cond_inputs=None, rev=False): + if str(self.L_mask.device) != str(self.L.device): + self.L_mask = self.L_mask.to(self.L.device) + self.U_mask = self.U_mask.to(self.L.device) + self.I = self.I.to(self.L.device) + self.P = self.P.to(self.L.device) + self.sign_S = self.sign_S.to(self.L.device) + + L = self.L * self.L_mask + self.I + U = self.U * self.U_mask + torch.diag( + self.sign_S * torch.exp(self.log_S)) + W = self.P @ L @ U + + if rev == False: + return inputs[0] @ W, self.log_S.sum().unsqueeze(0).unsqueeze( + 0).repeat(inputs[0].size(0), 1) + else: + return inputs[0] @ torch.inverse( + W), -self.log_S.sum().unsqueeze(0).unsqueeze(0).repeat( + inputs[0].size(0), 1) + + def output_dims(self, input_dims): + assert len(input_dims) == 1, "Can only use 1 input" + return input_dims + + def jacobian(self, x, rev=False): + return 0. + +class Shuffle(nn.Module): + """ An implementation of a shuffling layer from + Density estimation using Real NVP + (https://arxiv.org/abs/1605.08803). + """ + + def __init__(self, num_inputs): + super(Shuffle, self).__init__() + self.perm = np.random.permutation(num_inputs) + self.inv_perm = np.argsort(self.perm) + + def forward(self, inputs, cond_inputs=None, rev=False): + if rev == False: + return inputs[:, self.perm], torch.zeros( + inputs.size(0), 1, device=inputs.device) + else: + return inputs[:, self.inv_perm], torch.zeros( + inputs.size(0), 1, device=inputs.device) + + +class Reverse(nn.Module): + """ An implementation of a reversing layer from + Density estimation using Real NVP + (https://arxiv.org/abs/1605.08803). + """ + + def __init__(self, num_inputs): + super(Reverse, self).__init__() + self.perm = np.array(np.arange(0, num_inputs)[::-1]) + self.inv_perm = np.argsort(self.perm) + + def forward(self, inputs, cond_inputs=None, rev=False): + if rev == False: + return inputs[:, self.perm], torch.zeros( + inputs.size(0), 1, device=inputs.device) + else: + return inputs[:, self.inv_perm], torch.zeros( + inputs.size(0), 1, device=inputs.device) + + +class CouplingLayer(nn.Module): + """ An implementation of a coupling layer + from RealNVP (https://arxiv.org/abs/1605.08803). + """ + + def __init__(self, + num_inputs, + num_hidden, + mask, + num_cond_inputs=None, + s_act='tanh', + t_act='relu'): + super(CouplingLayer, self).__init__() + self.num_inputs = num_inputs[0][0] + self.mask = mask + + activations = {'relu': nn.ReLU, 'sigmoid': nn.Sigmoid, 'tanh': nn.Tanh} + s_act_func = activations[s_act] + t_act_func = activations[t_act] + + if num_cond_inputs is not None: + total_inputs = num_inputs + num_cond_inputs + else: + total_inputs = num_inputs + + total_inputs = total_inputs[0][0] + self.scale_net = nn.Sequential( + nn.Linear(total_inputs, num_hidden), s_act_func(), + nn.Linear(num_hidden, num_hidden), s_act_func(), + nn.Linear(num_hidden, num_inputs[0][0])) + self.translate_net = nn.Sequential( + nn.Linear(total_inputs, num_hidden), t_act_func(), + nn.Linear(num_hidden, num_hidden), t_act_func(), + nn.Linear(num_hidden, num_inputs[0][0])) + + def init(m): + if isinstance(m, nn.Linear): + m.bias.data.fill_(0) + nn.init.orthogonal_(m.weight.data) + + def forward(self, inputs, cond_inputs=None, rev=False): + mask = self.mask + inputs = inputs[0] + masked_inputs = inputs * mask + if cond_inputs is not None: + masked_inputs = torch.cat([masked_inputs, cond_inputs], -1) + + if rev == False: + log_s = self.scale_net(masked_inputs) * (1 - mask) + t = self.translate_net(masked_inputs) * (1 - mask) + s = torch.exp(log_s) + return inputs * s + t, log_s.sum(-1, keepdim=True) + else: + log_s = self.scale_net(masked_inputs) * (1 - mask) + t = self.translate_net(masked_inputs) * (1 - mask) + s = torch.exp(-log_s) + return (inputs - t) * s, -log_s.sum(-1, keepdim=True) + + def output_dims(self, input_dims): + assert len(input_dims) == 1, "Can only use 1 input" + return input_dims + +class FlowSequential(nn.Sequential): + """ A sequential container for flows. + In addition to a forward pass it implements a backward pass and + computes log jacobians. + """ + + def forward(self, inputs, cond_inputs=None, rev=False, logdets=None): + """ Performs a forward or backward pass for flow modules. + Args: + inputs: a tuple of inputs and logdets + mode: to run direct computation or inverse + """ + self.num_inputs = inputs.size(-1) + + if logdets is None: + logdets = torch.zeros(inputs.size(0), 1, device=inputs.device) + + # assert mode in ['direct', 'inverse'] + if rev == False: + for module in self._modules.values(): + inputs, logdet = module(inputs, cond_inputs, rev) + logdets += logdet + else: + for module in reversed(self._modules.values()): + inputs, logdet = module(inputs, cond_inputs, rev) + logdets += logdet + + return inputs, logdets + + def log_probs(self, inputs, cond_inputs = None): + u, log_jacob = self(inputs, cond_inputs) + log_probs = (-0.5 * u.pow(2) - 0.5 * math.log(2 * math.pi)).sum( + -1, keepdim=True) + return (log_probs + log_jacob).sum(-1, keepdim=True) + + def sample(self, num_samples=None, noise=None, cond_inputs=None): + if noise is None: + noise = torch.Tensor(num_samples, self.num_inputs).normal_() + device = next(self.parameters()).device + noise = noise.to(device) + if cond_inputs is not None: + cond_inputs = cond_inputs.to(device) + samples = self.forward(noise, cond_inputs, mode='inverse')[0] + return samples diff --git a/localization.py b/localization.py index 17e5778..1e91dab 100644 --- a/localization.py +++ b/localization.py @@ -35,7 +35,7 @@ def save_imgs(inputs, grad, cnt): def export_gradient_maps(model, testloader, optimizer, n_batches=1): plt.figure(figsize=(10, 10)) - testloader.dataset.get_fixed = True + testloader.dataset.get_fixed = False cnt = 0 degrees = -1 * np.arange(c.n_transforms_test) * 360.0 / c.n_transforms_test @@ -50,13 +50,13 @@ def export_gradient_maps(model, testloader, optimizer, n_batches=1): loss.backward() grad = inputs.grad.view(-1, c.n_transforms_test, *inputs.shape[-3:]) - grad = grad[labels > 0] + grad = grad[labels >= 0] if grad.shape[0] == 0: continue grad = t2np(grad) inputs = inputs.view(-1, c.n_transforms_test, *inputs.shape[-3:])[:, 0] - inputs = np.transpose(t2np(inputs[labels > 0]), [0, 2, 3, 1]) + inputs = np.transpose(t2np(inputs[labels >= 0]), [0, 2, 3, 1]) inputs_unnormed = np.clip(inputs * c.norm_std + c.norm_mean, 0, 1) for i_item in range(c.n_transforms_test): diff --git a/logo_detection.py b/logo_detection.py new file mode 100644 index 0000000..e1a304a --- /dev/null +++ b/logo_detection.py @@ -0,0 +1,60 @@ +from skimage import io +import matplotlib.pyplot as plt +import numpy as np +import cv2 + +from skimage.color import rgb2gray +from skimage import feature + +for i in range(12): + break + +I1 = io.imread("bottle_logo_defective1.jpg") +cv2.imwrite("origin.jpg", I1) + +image = cv2.cvtColor(I1, cv2.COLOR_BGR2GRAY) +cv2.imwrite("Gray.jpg", image) + +image = cv2.GaussianBlur(image, (21, 21), 0) + +# seg_img = cv2.adaptiveThreshold(image, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 3, 1) +seg_img = cv2.threshold(image, 30, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)[1] +cv2.imwrite('binary.jpg', seg_img) + +h0, w0 = seg_img.shape + +if seg_img[0][0] == 255: + for i in range(w0): + for j in range(2): + if seg_img[j][i] == 255: + seg_img[j][i] = 0 + for j in range(h0): + for i in range(2): + if seg_img[j][i] == 255: + seg_img[j][i] = 0 + +cv2.imwrite("Threshold.jpg", seg_img) + +num, labels, stats, centroids = cv2.connectedComponentsWithStats(seg_img) + +h, w = seg_img.shape +print(h * w) + +first_stat = 0 +second_stat = 0 + +# Largest area should be the bottle +# Second largest area should be the logo +for istat in stats: + if istat[4] > 2000 and istat[4] > second_stat: + if istat[4] > first_stat: + first_stat = istat[4] + else: + second_stat = istat[4] + logo_stat = istat + +print(logo_stat) +cv2.rectangle(I1, (logo_stat[0], logo_stat[1]), (logo_stat[0] + logo_stat[2], logo_stat[1] + logo_stat[3]), + (255, 0, 255), 2) + +cv2.imwrite("segmented.jpg", I1) diff --git a/main.py b/main.py deleted file mode 100644 index 9fac6ca..0000000 --- a/main.py +++ /dev/null @@ -1,12 +0,0 @@ -'''This is the repo which contains the original code to the WACV 2021 paper -"Same Same But DifferNet: Semi-Supervised Defect Detection with Normalizing Flows" -by Marco Rudolph, Bastian Wandt and Bodo Rosenhahn. -For further information contact Marco Rudolph (rudolph@tnt.uni-hannover.de)''' - -import config as c -from train import train -from utils import load_datasets, make_dataloaders - -train_set, test_set = load_datasets(c.dataset_path, c.class_name) -train_loader, test_loader = make_dataloaders(train_set, test_set) -model = train(train_loader, test_loader) diff --git a/model.py b/model.py index 89f7594..dc06eaa 100644 --- a/model.py +++ b/model.py @@ -1,70 +1,99 @@ -import numpy as np -import os -import torch -import torch.nn.functional as F -from torch import nn -from torchvision.models import alexnet - -import config as c -from freia_funcs import permute_layer, glow_coupling_layer, F_fully_connected, ReversibleGraphNet, OutputNode, \ - InputNode, Node - -WEIGHT_DIR = './weights' -MODEL_DIR = './models' - - -def nf_head(input_dim=c.n_feat): - nodes = list() - nodes.append(InputNode(input_dim, name='input')) - for k in range(c.n_coupling_blocks): - nodes.append(Node([nodes[-1].out0], permute_layer, {'seed': k}, name=F'permute_{k}')) - nodes.append(Node([nodes[-1].out0], glow_coupling_layer, - {'clamp': c.clamp_alpha, 'F_class': F_fully_connected, - 'F_args': {'internal_size': c.fc_internal, 'dropout': c.dropout}}, - name=F'fc_{k}')) - nodes.append(OutputNode([nodes[-1].out0], name='output')) - coder = ReversibleGraphNet(nodes) - return coder - - -class DifferNet(nn.Module): - def __init__(self): - super(DifferNet, self).__init__() - self.feature_extractor = alexnet(pretrained=True) - self.nf = nf_head() - - def forward(self, x): - y_cat = list() - - for s in range(c.n_scales): - x_scaled = F.interpolate(x, size=c.img_size[0] // (2 ** s)) if s > 0 else x - feat_s = self.feature_extractor.features(x_scaled) - y_cat.append(torch.mean(feat_s, dim=(2, 3))) - - y = torch.cat(y_cat, dim=1) - z = self.nf(y) - return z - - -def save_model(model, filename): - if not os.path.exists(MODEL_DIR): - os.makedirs(MODEL_DIR) - torch.save(model, os.path.join(MODEL_DIR, filename)) - - -def load_model(filename): - path = os.path.join(MODEL_DIR, filename) - model = torch.load(path) - return model - - -def save_weights(model, filename): - if not os.path.exists(WEIGHT_DIR): - os.makedirs(WEIGHT_DIR) - torch.save(model.state_dict(), os.path.join(WEIGHT_DIR, filename)) - - -def load_weights(model, filename): - path = os.path.join(WEIGHT_DIR, filename) - model.load_state_dict(torch.load(path)) - return model +import os +import torch +import torch.nn.functional as F +from torch import nn +from torchvision.models import alexnet + +import config as c +from freia_funcs import permute_layer, glow_coupling_layer, F_fully_connected, ReversibleGraphNet, OutputNode, \ + InputNode, Node + +from datetime import datetime +import matplotlib.pyplot as plt +import json + +WEIGHT_DIR = './weights' +MODEL_DIR = './models' + + +def nf_head(input_dim=c.n_feat): + nodes = list() + nodes.append(InputNode(input_dim, name='input')) + for k in range(c.n_coupling_blocks): + nodes.append(Node([nodes[-1].out0], permute_layer, {'seed': k}, name=F'permute_{k}')) + nodes.append(Node([nodes[-1].out0], glow_coupling_layer, + {'clamp': c.clamp_alpha, 'F_class': F_fully_connected, + 'F_args': {'internal_size': c.fc_internal, 'dropout': c.dropout}}, + name=F'fc_{k}')) + nodes.append(OutputNode([nodes[-1].out0], name='output')) + coder = ReversibleGraphNet(nodes) + return coder + + +class DifferNet(nn.Module): + def __init__(self): + super(DifferNet, self).__init__() + self.feature_extractor = alexnet(pretrained=True) + self.nf = nf_head() + + def forward(self, x): + y_cat = list() + + for s in range(c.n_scales): + x_scaled = F.interpolate(x, size=c.img_size[0] // (2 ** s)) if s > 0 else x + feat_s = self.feature_extractor.features(x_scaled) + y_cat.append(torch.mean(feat_s, dim=(2, 3))) + + y = torch.cat(y_cat, dim=1) + z = self.nf(y) + return z + + +def save_model(model, filename): + if not os.path.exists(MODEL_DIR): + os.makedirs(MODEL_DIR) + torch.save(model, os.path.join(MODEL_DIR, filename)) + + +def load_model(filename): + path = os.path.join(MODEL_DIR, filename) + model = torch.load(path) + return model + + +def save_weights(model, filename): + if not os.path.exists(WEIGHT_DIR): + os.makedirs(WEIGHT_DIR) + torch.save(model.state_dict(), os.path.join(WEIGHT_DIR, filename)) + + +def load_weights(model, filename): + path = os.path.join(WEIGHT_DIR, filename) + model.load_state_dict(torch.load(path)) + return model + + +def save_parameters(model_parameters, filename): + if not os.path.exists(MODEL_DIR): + os.makedirs(MODEL_DIR) + + with open(MODEL_DIR + '/' + filename + '.json', 'w') as jsonfile: + jsonfile.write(json.dumps(model_parameters, indent=4)) + +def save_roc_plot(fpr, tpr, filename): + plt.figure() + lw = 2 + plt.figure(figsize=(10, 10)) + plt.plot(fpr.tolist(), tpr.tolist(), color='darkorange', + lw=lw, label='ROC curve') + plt.plot([0, 1], [0, 1], color='navy', lw=lw, linestyle='--') + plt.xlim([0.0, 1.0]) + plt.ylim([0.0, 1.0]) + plt.xlabel('False Positive Rate') + plt.ylabel('True Positive Rate') + plt.title('ROC Curve') + plt.legend(loc="lower right") + now = datetime.now() + dt_string = now.strftime("%Y%m%d%H%M%S") + # plt.savefig(MODEL_DIR + '/' +filename + '_ROC_' + dt_string + '.jpg') + plt.savefig(MODEL_DIR + '/' + filename + '_ROC.jpg') \ No newline at end of file diff --git a/multi_transform_loader.py b/multi_transform_loader.py index edfb5b8..44c3192 100644 --- a/multi_transform_loader.py +++ b/multi_transform_loader.py @@ -46,11 +46,15 @@ def __getitem__(self, index): samples = list() for i in range(self.n_transforms): if self.get_fixed: + # print(f"i={i}: calling fixed_rotation({sample}, {self.fixed_degrees[i]}))") samples.append(fixed_rotation(self, sample, self.fixed_degrees[i])) - else: - samples.append(self.transform(sample)) + else: + new = self.transform(sample) + samples.append(new) + # print(f"i={i}: calling transform({sample})") samples = torch.stack(samples, dim=0) if self.target_transform is not None: + # print(f"calling target_transform({target})") target = self.target_transform(target) return samples, target diff --git a/requirements.txt b/requirements.txt index eb100f1..508f397 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,3 +5,4 @@ torch>=1.00 torchvision>=0.2.2 matplotlib>=3.0.3 tqdm>=4.40.2 +opencv-python \ No newline at end of file diff --git a/run_test.py b/run_test.py new file mode 100644 index 0000000..f733db1 --- /dev/null +++ b/run_test.py @@ -0,0 +1,34 @@ +'''This is the repo which contains the original code to the WACV 2021 paper +"Same Same But DifferNet: Semi-Supervised Defect Detection with Normalizing Flows" +by Marco Rudolph, Bastian Wandt and Bodo Rosenhahn. +For further information contact Marco Rudolph (rudolph@tnt.uni-hannover.de)''' + +from test import * +from utils import load_datasets, make_dataloaders +import time +import gc +import json + +c.transf_brightness = 0.0 +c.transf_contrast = 0.0 +c.transf_saturation = 0.0 + +_, _, test_set = load_datasets(c.dataset_path, c.class_name, test=True) +_, _, test_loader = make_dataloaders(None, None, test_set, test=True) +model = torch.load('models/' + c.modelname + '.pth', map_location=torch.device('cpu')) + +with open('models/' + c.modelname + '.json') as jsonfile: + model_parameters = json.load(jsonfile) + +time_start = time.time() +test(model, model_parameters, test_loader) +time_end = time.time() +time_c = time_end - time_start +print("testing time cost: {:f} s".format(time_c)) + +# free memory +del test_set +del test_loader + +gc.collect() +torch.cuda.empty_cache() diff --git a/run_traning.py b/run_traning.py new file mode 100644 index 0000000..91e06a3 --- /dev/null +++ b/run_traning.py @@ -0,0 +1,28 @@ +'''This is the repo which contains the original code to the WACV 2021 paper +"Same Same But DifferNet: Semi-Supervised Defect Detection with Normalizing Flows" +by Marco Rudolph, Bastian Wandt and Bodo Rosenhahn. +For further information contact Marco Rudolph (rudolph@tnt.uni-hannover.de)''' + +from train import * +from utils import load_datasets, make_dataloaders +import time +import gc + +train_set, validate_set, _ = load_datasets(c.dataset_path, c.class_name) +train_loader, validate_loader, _ = make_dataloaders(train_set, validate_set, None) + +time_start = time.time() +model, model_parameters = train(train_loader, validate_loader) + +time_end = time.time() +time_c = time_end - time_start # 运行所花时间 +print("train time cost: {:f} s".format(time_c)) + +# free memory +del train_set +del validate_set +del train_loader +del validate_loader + +gc.collect() +torch.cuda.empty_cache() \ No newline at end of file diff --git a/test.py b/test.py new file mode 100644 index 0000000..0964ec9 --- /dev/null +++ b/test.py @@ -0,0 +1,131 @@ +from sklearn.metrics import roc_auc_score +from sklearn.metrics import roc_curve +from localization import export_gradient_maps +from model import save_roc_plot +from utils import * +from operator import itemgetter +import cv2 + +def test(model, model_parameters, test_loader): + print("Running test") + optimizer = torch.optim.Adam(model.nf.parameters(), lr=c.lr_init, betas=(0.8, 0.8), eps=1e-04, weight_decay=1e-5) + model.to(c.device) + model.eval() + if c.verbose: + print('\nCompute loss and scores on test set:') + test_z = list() + test_labels = list() + predictions = [] + with torch.no_grad(): + for i, data in enumerate(test_loader): + inputs, labels = preprocess_batch(data) + if c.frame_name_is_given: + frame = int(test_loader.dataset.imgs[i][0].split('frame',1)[1].split('-')[0]) + frame = i + #print(f"i={i}: frame#={frame}, labels={labels.cpu().numpy()[0]}, size of inputs={inputs.size()}") + predictions.append([frame, test_loader.dataset.imgs[i][0], labels.cpu().numpy()[0], 0, 0]) + z = model(inputs) + test_z.append(z) + test_labels.append(t2np(labels)) + + test_labels = np.concatenate(test_labels) + is_anomaly = np.array([0 if l == 0 else 1 for l in test_labels]) + + z_grouped = torch.cat(test_z, dim=0).view(-1, c.n_transforms_test, c.n_feat) + anomaly_score = t2np(torch.mean(z_grouped ** 2, dim=(-2, -1))) + AUROC = roc_auc_score(is_anomaly, anomaly_score) + fpr, tpr, thresholds = roc_curve(is_anomaly, anomaly_score) + save_roc_plot(fpr, tpr, c.modelname + "_{:.4f}_test".format(AUROC)) + + for i in range(len(model_parameters['tpr'])): + if model_parameters['tpr'][i] > c.target_tpr: + target_threshold = model_parameters['thresholds'][i] + break + + is_anomaly_detected = [] + i = 0 + for l in anomaly_score: + predictions[i][4] = l + if l < target_threshold: + is_anomaly_detected.append(0) + predictions[i][3] = 0 + else: + is_anomaly_detected.append(1) + predictions[i][3] = 1 + i += 1 + predictions = sorted(predictions, key=itemgetter(0)) + + # calculate test accuracy + error_count = 0 + for i in range(len(is_anomaly)): + if is_anomaly[i] != is_anomaly_detected[i]: + error_count += 1 + + test_accuracy = 1 - float(error_count) / len(is_anomaly) + #todo: tpr/fpr display. + + for i in range(len(predictions)): + msg = 'frame: ' + str(i) + '. ' + if (predictions[i][3] == 1): + msg += 'prediction: defective. ' + else: + msg += 'prediction: good. ' + + if (predictions[i][2] == 1): + msg += 'ground truth: defective. ' + else: + msg += 'ground truth: good. ' + + msg += 'anomaly score: ' + str(round(predictions[i][4], 4)) + '. ' + msg += 'threshold: ' + str(round(target_threshold, 4)) + '. ' + msg += 'accuracy: ' + str(round(test_accuracy * 100, 2)) + '%' + + print(msg) + + # print(f"test_labels={test_labels}, is_anomaly={is_anomaly},anomaly_score={anomaly_score},is_anomaly_detected={is_anomaly_detected}") + print(f"target_tpr={c.target_tpr}, target_threshold={target_threshold}, test_accuracy={test_accuracy}") + if c.grad_map_viz: + print("saving gradient maps...") + export_gradient_maps(model, test_loader, optimizer, -1) + + # visualize the prediction result + if c.visualization: + for i in range(len(predictions)): + # load file path + file_path = predictions[i][1] + idx = file_path.index('video') + file_path = file_path[:idx] + 'original-' + file_path[idx:] + file_path = file_path.replace("test\\test", "zerobox-2010-1-original") + + # rotate and resize image + img = cv2.imread(file_path) + img = cv2.rotate(img, cv2.cv2.ROTATE_90_COUNTERCLOCKWISE) + img = cv2.resize(img, (600, 900)) + + # display prediction on each frame + font = cv2.FONT_HERSHEY_DUPLEX + font_size = 0.65 + pos_x = 330 + if (predictions[i][3] == 1): + img = cv2.putText(img, 'prediction: defective', (pos_x, 810), font, + font_size, (0, 0, 255), 1, cv2.LINE_AA) + else: + img = cv2.putText(img, 'prediction: good', (pos_x, 810), font, + font_size, (0, 255, 0), 1, cv2.LINE_AA) + + if (predictions[i][2] == 1): + img = cv2.putText(img, 'ground truth: defective', (pos_x, 830), font, + font_size, (0, 0, 255), 1, cv2.LINE_AA) + else: + img = cv2.putText(img, 'ground truth: good', (pos_x, 830), font, + font_size, (0, 255, 0), 1, cv2.LINE_AA) + + img = cv2.putText(img, 'anomaly score: ' + str(round(predictions[i][4], 4)), (pos_x, 850), font, + font_size, (0, 255, 0), 1, cv2.LINE_AA) + img = cv2.putText(img, 'threshold: ' + str(round(target_threshold, 4)), (pos_x, 870), font, + font_size, (0, 255, 0), 1, cv2.LINE_AA) + img = cv2.putText(img, 'accuracy: ' + str(round(test_accuracy * 100, 2)) + '%', (pos_x, 890), font, + font_size, (0, 255, 0), 1, cv2.LINE_AA) + # show results + cv2.imshow('window', img) + cv2.waitKey(220) \ No newline at end of file diff --git a/train.py b/train.py index f3ebf73..689763a 100644 --- a/train.py +++ b/train.py @@ -1,14 +1,10 @@ -import numpy as np -import torch from sklearn.metrics import roc_auc_score +from sklearn.metrics import roc_curve from tqdm import tqdm - -import config as c from localization import export_gradient_maps -from model import DifferNet, save_model, save_weights +from model import DifferNet, save_model, save_weights, save_parameters, save_roc_plot from utils import * - class Score_Observer: '''Keeps an eye on the current and highest score so far''' @@ -30,12 +26,15 @@ def print_score(self): print('{:s}: \t last: {:.4f} \t max: {:.4f} \t epoch_max: {:d}'.format(self.name, self.last, self.max_score, self.max_epoch)) - -def train(train_loader, test_loader): +def train(train_loader, validate_loader): model = DifferNet() - optimizer = torch.optim.Adam(model.nf.parameters(), lr=c.lr_init, betas=(0.8, 0.8), eps=1e-04, weight_decay=1e-5) + optimizer = torch.optim.Adam([{'params': model.nf.parameters()}], lr=c.lr_init, betas=(0.8, 0.8), eps=1e-04, + weight_decay=1e-5) model.to(c.device) + save_name_pre = '{}_{}_{:.2f}_{:.2f}_{:.2f}_{:.2f}'.format(c.modelname, c.rotation_degree, + c.crop_top, c.crop_left, c.crop_bottom, c.crop_right) + # todo: learning rate score_obs = Score_Observer('AUROC') for epoch in range(c.meta_epochs): @@ -49,9 +48,6 @@ def train(train_loader, test_loader): for i, data in enumerate(tqdm(train_loader, disable=c.hide_tqdm_bar)): optimizer.zero_grad() inputs, labels = preprocess_batch(data) # move to device and reshape - # TODO inspect - # inputs += torch.randn(*inputs.shape).cuda() * c.add_img_noise - z = model(inputs) loss = get_loss(z, model.nf.jacobian(run_forward=False)) train_loss.append(t2np(loss)) @@ -62,39 +58,62 @@ def train(train_loader, test_loader): if c.verbose: print('Epoch: {:d}.{:d} \t train loss: {:.4f}'.format(epoch, sub_epoch, mean_train_loss)) - # evaluate - model.eval() - if c.verbose: - print('\nCompute loss and scores on test set:') - test_loss = list() - test_z = list() - test_labels = list() - with torch.no_grad(): - for i, data in enumerate(tqdm(test_loader, disable=c.hide_tqdm_bar)): - inputs, labels = preprocess_batch(data) - z = model(inputs) - loss = get_loss(z, model.nf.jacobian(run_forward=False)) - test_z.append(z) - test_loss.append(t2np(loss)) - test_labels.append(t2np(labels)) - - test_loss = np.mean(np.array(test_loss)) - if c.verbose: - print('Epoch: {:d} \t test_loss: {:.4f}'.format(epoch, test_loss)) + if not (validate_loader is None): + # evaluate + model.eval() + if c.verbose: + print('\nCompute loss and scores on validate set:') + test_loss = list() + test_z = list() + test_labels = list() + with torch.no_grad(): + for i, data in enumerate(tqdm(validate_loader, disable=c.hide_tqdm_bar)): + inputs, labels = preprocess_batch(data) + z = model(inputs) + loss = get_loss(z, model.nf.jacobian(run_forward=False)) + test_z.append(z) + test_loss.append(t2np(loss)) + test_labels.append(t2np(labels)) + + test_loss = np.mean(np.array(test_loss)) + + test_labels = np.concatenate(test_labels) + is_anomaly = np.array([0 if l == 0 else 1 for l in test_labels]) + + z_grouped = torch.cat(test_z, dim=0).view(-1, c.n_transforms_test, c.n_feat) + anomaly_score = t2np(torch.mean(z_grouped ** 2, dim=(-2, -1))) + AUROC = roc_auc_score(is_anomaly, anomaly_score) + score_obs.update(AUROC, epoch, + print_score=c.verbose or epoch == c.meta_epochs - 1) + + fpr, tpr, thresholds = roc_curve(is_anomaly, anomaly_score) + model_parameters = {} + model_parameters['fpr'] = fpr.tolist() + model_parameters['tpr'] = tpr.tolist() + model_parameters['thresholds'] = thresholds.tolist() + model_parameters['AUROC'] = AUROC + + if epoch == c.meta_epochs - 1: + save_parameters(model_parameters, c.modelname) + save_roc_plot(fpr, tpr, c.modelname + "_{:.4f}".format(AUROC)) - test_labels = np.concatenate(test_labels) - is_anomaly = np.array([0 if l == 0 else 1 for l in test_labels]) + if c.verbose: + print('Epoch: {:d} \t validate_loss: {:.4f}'.format(epoch, test_loss)) - z_grouped = torch.cat(test_z, dim=0).view(-1, c.n_transforms_test, c.n_feat) - anomaly_score = t2np(torch.mean(z_grouped ** 2, dim=(-2, -1))) - score_obs.update(roc_auc_score(is_anomaly, anomaly_score), epoch, - print_score=c.verbose or epoch == c.meta_epochs - 1) + # compare is_anomaly and anomaly_score + np.set_printoptions(precision=2, suppress=True) + print('is_anomaly: ', is_anomaly) + print('anomaly_score: ', anomaly_score) + print('fpr: ', fpr) + print('tpr: ', tpr) + print('thresholds: ', thresholds) - if c.grad_map_viz: - export_gradient_maps(model, test_loader, optimizer, -1) + if c.grad_map_viz and not (validate_loader is None): + export_gradient_maps(model, validate_loader, optimizer, 1) if c.save_model: model.to('cpu') - save_model(model, c.modelname) - save_weights(model, c.modelname) - return model + save_model(model, c.modelname + '.pth') + save_weights(model, c.modelname + '.weights.pth') + + return model, model_parameters diff --git a/utils.py b/utils.py index af3881c..d792fe8 100644 --- a/utils.py +++ b/utils.py @@ -1,105 +1,177 @@ -import os -import torch -from torch.utils.data import DataLoader -from torchvision import datasets, transforms - -import config as c -from multi_transform_loader import ImageFolderMultiTransform - - -def t2np(tensor): - '''pytorch tensor -> numpy array''' - return tensor.cpu().data.numpy() if tensor is not None else None - - -def get_loss(z, jac): - '''check equation 4 of the paper why this makes sense - oh and just ignore the scaling here''' - return torch.mean(0.5 * torch.sum(z ** 2, dim=(1,)) - jac) / z.shape[1] - - -def load_datasets(dataset_path, class_name): - ''' - Expected folder/file format to find anomalies of class from dataset location : - - train data: - - dataset_path/class_name/train/good/any_filename.png - dataset_path/class_name/train/good/another_filename.tif - dataset_path/class_name/train/good/xyz.png - [...] - - test data: - - 'normal data' = non-anomalies - - dataset_path/class_name/test/good/name_the_file_as_you_like_as_long_as_there_is_an_image_extension.webp - dataset_path/class_name/test/good/did_you_know_the_image_extension_webp?.png - dataset_path/class_name/test/good/did_you_know_that_filenames_may_contain_question_marks????.png - dataset_path/class_name/test/good/dont_know_how_it_is_with_windows.png - dataset_path/class_name/test/good/just_dont_use_windows_for_this.png - [...] - - anomalies - assume there are anomaly classes 'crack' and 'curved' - - dataset_path/class_name/test/crack/dat_crack_damn.png - dataset_path/class_name/test/crack/let_it_crack.png - dataset_path/class_name/test/crack/writing_docs_is_fun.png - [...] - - dataset_path/class_name/test/curved/wont_make_a_difference_if_you_put_all_anomalies_in_one_class.png - dataset_path/class_name/test/curved/but_this_code_is_practicable_for_the_mvtec_dataset.png - [...] - ''' - - def target_transform(target): - return class_perm[target] - - data_dir_train = os.path.join(dataset_path, class_name, 'train') - data_dir_test = os.path.join(dataset_path, class_name, 'test') - - classes = os.listdir(data_dir_test) - if 'good' not in classes: - print('There should exist a subdirectory "good". Read the doc of this function for further information.') - exit() - classes.sort() - class_perm = list() - class_idx = 1 - for cl in classes: - if cl == 'good': - class_perm.append(0) - else: - class_perm.append(class_idx) - class_idx += 1 - - augmentative_transforms = [] - if c.transf_rotations: - augmentative_transforms += [transforms.RandomRotation(180)] - if c.transf_brightness > 0.0 or c.transf_contrast > 0.0 or c.transf_saturation > 0.0: - augmentative_transforms += [transforms.ColorJitter(brightness=c.transf_brightness, contrast=c.transf_contrast, - saturation=c.transf_saturation)] - - tfs = [transforms.Resize(c.img_size)] + augmentative_transforms + [transforms.ToTensor(), - transforms.Normalize(c.norm_mean, c.norm_std)] - - transform_train = transforms.Compose(tfs) - - trainset = ImageFolderMultiTransform(data_dir_train, transform=transform_train, n_transforms=c.n_transforms) - testset = ImageFolderMultiTransform(data_dir_test, transform=transform_train, target_transform=target_transform, - n_transforms=c.n_transforms_test) - return trainset, testset - - -def make_dataloaders(trainset, testset): - trainloader = torch.utils.data.DataLoader(trainset, pin_memory=True, batch_size=c.batch_size, shuffle=True, - drop_last=False) - testloader = torch.utils.data.DataLoader(testset, pin_memory=True, batch_size=c.batch_size_test, shuffle=True, - drop_last=False) - return trainloader, testloader - - -def preprocess_batch(data): - '''move data to device and reshape image''' - inputs, labels = data - inputs, labels = inputs.to(c.device), labels.to(c.device) - inputs = inputs.view(-1, *inputs.shape[-3:]) - return inputs, labels +import os +import torch +from torchvision import datasets, transforms + +import config as c +from multi_transform_loader import ImageFolderMultiTransform + +import cv2 +import numpy as np +from datetime import datetime + +TRANSFORM_DIR = "./transform/" + +def TransformShow(name="img", wait=100): + def transform_show(img): + # path = "transform/" + # now = datetime.now() + # dt_string = now.strftime("%d%m%Y%H%M%S") + # cv2.imwrite(path + 'all_transform_' + dt_string + '.jpg', np.array(img)) + # cv2.imshow(name, np.array(img)) + # cv2.waitKey(wait) + return img + + return transform_show + +def cropImage(): + def crop_image(img): + x,y,w,h = shrinkEdges(img.size) + rs = transforms.functional.crop(img,y,x,h,w) + + if not os.path.exists(TRANSFORM_DIR): + os.makedirs(TRANSFORM_DIR) + + now = datetime.now() + dt_string = now.strftime("%d%m%Y%H%M%S") + if(c.save_transformed_image): + cv2.imwrite(TRANSFORM_DIR + 'transform_' + dt_string + '.jpg', np.array(rs)) + return rs + + return crop_image + +def shrinkEdges(img_size): + width, height = img_size + shrink_scale_top = c.crop_top + shrink_scale_bot = c.crop_bottom + shrink_scale_left = c.crop_left + shrink_scale_right = c.crop_right + left_reduction = shrink_scale_left * width + right_reduction = shrink_scale_right * width + top_reduction = shrink_scale_top * height + bot_reduction = shrink_scale_bot * height + new_height = int(height - top_reduction - bot_reduction) + new_width = int(width - left_reduction - right_reduction) + new_ul_x = int(left_reduction) + new_ul_y = int(top_reduction) + # print( + # f"shrinking {0, 0, width, height} to {new_ul_x, new_ul_y, new_width, new_height}" + # ) + return new_ul_x, new_ul_y, new_width, new_height + + +def t2np(tensor): + '''pytorch tensor -> numpy array''' + return tensor.cpu().data.numpy() if tensor is not None else None + + +def get_loss(z, jac): + '''check equation 4 of the paper why this makes sense - oh and just ignore the scaling here''' + return torch.mean(0.5 * torch.sum(z ** 2, dim=(1,)) - jac) / z.shape[1] + + +def load_datasets(dataset_path, class_name, test=False): + ''' + Expected folder/file format to find anomalies of class from dataset location : + + train data: + + dataset_path/class_name/train/good/any_filename.png + dataset_path/class_name/train/good/another_filename.tif + dataset_path/class_name/train/good/xyz.png + [...] + + test data: + + 'normal data' = non-anomalies + + dataset_path/class_name/test/good/name_the_file_as_you_like_as_long_as_there_is_an_image_extension.webp + dataset_path/class_name/test/good/did_you_know_the_image_extension_webp?.png + dataset_path/class_name/test/good/did_you_know_that_filenames_may_contain_question_marks????.png + dataset_path/class_name/test/good/dont_know_how_it_is_with_windows.png + dataset_path/class_name/test/good/just_dont_use_windows_for_this.png + [...] + + anomalies - assume there are anomaly classes 'crack' and 'curved' + + dataset_path/class_name/test/crack/dat_crack_damn.png + dataset_path/class_name/test/crack/let_it_crack.png + dataset_path/class_name/test/crack/writing_docs_is_fun.png + [...] + + dataset_path/class_name/test/curved/wont_make_a_difference_if_you_put_all_anomalies_in_one_class.png + dataset_path/class_name/test/curved/but_this_code_is_practicable_for_the_mvtec_dataset.png + [...] + ''' + + def target_transform(target): + return class_perm[target] + + data_dir_train = os.path.join(dataset_path, class_name, 'train') + data_dir_validate = os.path.join(dataset_path, class_name, 'validate') + data_dir_test = os.path.join(dataset_path, class_name, 'test') + + classes = os.listdir(data_dir_validate) + if 'good' not in classes: + print('There should exist a subdirectory "good". Read the doc of this function for further information.') + exit() + classes.sort() + class_perm = list() + class_idx = 1 + for cl in classes: + if cl == 'good': + class_perm.append(0) + else: + class_perm.append(class_idx) + class_idx += 1 + + augmentative_transforms = [] + if c.transf_rotations: + augmentative_transforms += [transforms.RandomRotation(c.rotation_degree)] + if c.transf_brightness > 0.0 or c.transf_contrast > 0.0 or c.transf_saturation > 0.0: + augmentative_transforms += [transforms.ColorJitter(brightness=c.transf_brightness, contrast=c.transf_contrast, + saturation=c.transf_saturation)] + + tfs = [cropImage(), transforms.Resize(c.img_size)] \ + + augmentative_transforms + [ TransformShow("Transformed Image", 10), transforms.ToTensor(), transforms.Normalize(c.norm_mean, c.norm_std)] + + transform_train = transforms.Compose(tfs) + + trainset = None + validateset = None + testset = None + if test == False: + trainset = ImageFolderMultiTransform(data_dir_train, transform=transform_train, n_transforms=c.n_transforms) + validateset = ImageFolderMultiTransform(data_dir_validate, transform=transform_train, target_transform=target_transform, + n_transforms=c.n_transforms_test) + else: + testset = ImageFolderMultiTransform(data_dir_test, transform=transform_train, target_transform=target_transform, + n_transforms=c.n_transforms_test) + + return trainset, validateset, testset + + +def make_dataloaders(trainset, validateset, testset, test=False): + trainloader = None + validateloader = None + testloader = None + if test == False: + trainloader = torch.utils.data.DataLoader(trainset, pin_memory=True, batch_size=c.batch_size, shuffle=True, + drop_last=False) + validateloader = torch.utils.data.DataLoader(validateset, pin_memory=True, batch_size=c.batch_size, shuffle=True, + drop_last=False) + else: + testloader = torch.utils.data.DataLoader(testset, pin_memory=True, batch_size=c.batch_size_test, shuffle=False, + drop_last=False) + + return trainloader, validateloader, testloader + + +def preprocess_batch(data): + '''move data to device and reshape image''' + inputs, labels = data + #print(f"begin: size of inputs={inputs.size()}") + inputs, labels = inputs.to(c.device), labels.to(c.device) + #print(f"to: size of inputs={inputs.size()}") + inputs = inputs.view(-1, *inputs.shape[-3:]) + #print(f"view: size of inputs={inputs.size()}") + return inputs, labels