├── .gitignore ├── AE ├── autoencoder1.py ├── dataset.py ├── help_fucntions.py ├── pytorch_msssim │ ├── __init__.py │ └── ssim.py ├── pytorch_ssim │ └── __init__.py ├── train.py └── train.sh ├── AutoencodersForUnsupervisedAnomalyDetection └── Deep Autoeoncoding Models for Unsupervised Anomaly Detection in Brain MR Images.pdf ├── GANomaly-pytorch ├── README.md ├── ganomaly.py ├── loss.py ├── nets.py ├── reporter.py └── utils.py ├── LICENSE ├── README.md ├── Unet ├── Unet2d_pytorch.py ├── ganComponents.py ├── nnBuildUnits.py ├── runAnoDetectGAN_2D.py └── utils.py ├── anomaly-detection-image-completion ├── .gitignore ├── README.md ├── main.ipynb └── scripts │ ├── data_helpers.py │ ├── model_helpers.py │ └── train_helpers.py ├── anomalyDetection └── resource.md ├── cAAE ├── README.md ├── funcs │ ├── __init__.py │ ├── networks.py │ ├── preproc.py │ └── resblocks.py ├── import_dataset.py ├── main.py ├── model.py └── normalize_data.py ├── config ├── RED_Net_2skips-grid.json ├── RED_Net_2skips-mvtec.json ├── RED_Net_3skips-grid.json ├── RED_Net_3skips-mvtec.json ├── RED_Net_4skips-mvtec.json ├── SRGAN-mvtec.json ├── SSIM-mvtec.json ├── SSIM_lite-chip.json ├── VAE.json └── __init__.py ├── db ├── __init__.py ├── augment.py ├── chip.py ├── eval_func.py └── mvtec.py ├── doc └── architecture.jpg ├── factory.py ├── model ├── __init__.py ├── gan_trainer.py ├── loss │ ├── Mutil_SSIM_loss.py │ ├── SRGAN_loss.py │ ├── SSIM_loss.py │ ├── VAE_loss.py │ └── __init__.py ├── networks │ ├── RED_Net │ │ ├── red_net_2skips.py │ │ ├── red_net_3skips.py │ │ └── red_net_4skips.py │ ├── SRGAN │ │ ├── SRGAN.py │ │ └── vgg16.py │ ├── SSIM │ │ ├── ssim.py │ │ └── ssim_lite.py │ ├── VAE │ │ └── VAE.py │ └── __init__.py ├── rebuilder.py ├── segmentation │ ├── __init__.py │ ├── ssim_seg.py │ └── ssim_seg_cuda.py └── trainer.py ├── test.py ├── textured_surface_anomaly_detection ├── A_compact_convolutional_neura_network_for_textured_surface_anomaly_detection.pdf ├── README.md ├── network.py ├── provider.py ├── tf_utils.py ├── train.py ├── train_tb.py └── unpool.py ├── tools ├── __init__.py ├── log.py └── timer.py ├── train.py └── train.sh /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | *.egg-info/ 24 | .installed.cfg 25 | *.egg 26 | MANIFEST 27 | 28 | # PyInstaller 29 | # Usually these files are written by a python script from a template 30 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 31 | *.manifest 32 | *.spec 33 | 34 | # Installer logs 35 | pip-log.txt 36 | pip-delete-this-directory.txt 37 | 38 | # Unit test / coverage reports 39 | htmlcov/ 40 | .tox/ 41 | .coverage 42 | .coverage.* 43 | .cache 44 | nosetests.xml 45 | coverage.xml 46 | *.cover 47 | .hypothesis/ 48 | .pytest_cache/ 49 | 50 | # Translations 51 | *.mo 52 | *.pot 53 | 54 | # Django stuff: 55 | *.log 56 | local_settings.py 57 | db.sqlite3 58 | 59 | # Flask stuff: 60 | instance/ 61 | .webassets-cache 62 | 63 | # Scrapy stuff: 64 | .scrapy 65 | 66 | # Sphinx documentation 67 | docs/_build/ 68 | 69 | # PyBuilder 70 | target/ 71 | 72 | # Jupyter Notebook 73 | .ipynb_checkpoints 74 | 75 | # pyenv 76 | .python-version 77 | 78 | # celery beat schedule file 79 | celerybeat-schedule 80 | 81 | # SageMath parsed files 82 | *.sage.py 83 | 84 | # Environments 85 | .env 86 | .venv 87 | env/ 88 | venv/ 89 | ENV/ 90 | env.bak/ 91 | venv.bak/ 92 | 93 | # Spyder project settings 94 | .spyderproject 95 | .spyproject 96 | 97 | # Rope project settings 98 | .ropeproject 99 | 100 | # mkdocs documentation 101 | /site 102 | 103 | # mypy 104 | .mypy_cache/ 105 | -------------------------------------------------------------------------------- /AE/dataset.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | import os 4 | import random 5 | 6 | class Dataset(): 7 | def __init__(self, img_path, img_size=64): 8 | if os.path.isdir(img_path): 9 | files = [] 10 | for _root, dirs, _files in os.walk(img_pathv): 11 | for _file in _files: 12 | if _file.startswith('._'): 13 | continue 14 | files.append(os.path.join(_root, _file)) 15 | elif os.path.isfile(img_path): 16 | with open(img_path, 'r') as f: 17 | files = f.readlines() 18 | random.shuffle(files) 19 | self.files = [item.strip() for item in files] 20 | self.img_size = img_size 21 | 22 | def __getitem__(self, index): 23 | img_path = self.files[index] 24 | img = cv2.imread(img_path) 25 | img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 26 | img = cv2.resize(img, (self.img_size, self.img_size)) 27 | #img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) 28 | img = np.expand_dims(img, axis=-1) 29 | img = img.transpose(2, 0, 1).astype('float32') 30 | img = img / 255. 31 | return img_path, img 32 | 33 | def __len__(self): 34 | return len(self.files) 35 | 36 | 37 | import numpy as np 38 | from torchvision import datasets 39 | import torchvision.transforms as transforms 40 | import random 41 | from torch.utils.data.sampler import SubsetRandomSampler 42 | from torch.utils.data import DataLoader 43 | 44 | normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], 45 | std=[0.229, 0.224, 0.225]) 46 | 47 | def get_data_loader(img_size,img_path,batch_size,num_workers=8): 48 | my_transform = transforms.Compose([ 49 | transforms.Resize((img_size, img_size)), 50 | transforms.ToTensor() 51 | #normalize 52 | ]) 53 | train_dataset = datasets.ImageFolder(root = img_path, transform=my_transform) 54 | train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, 55 | shuffle=True, num_workers=num_workers,pin_memory=True) 56 | 57 | return train_loader -------------------------------------------------------------------------------- /AE/help_fucntions.py: -------------------------------------------------------------------------------- 1 | import torch as T 2 | import numpy as np 3 | 4 | def get_cuda(tensor): 5 | if T.cuda.is_available(): 6 | tensor = tensor.cuda() 7 | return tensor 8 | 9 | # def get_onehot_labels(labels, n_c): 10 | # batch_size = len(labels) 11 | # class_onehot = np.zeros((batch_size, n_c)) 12 | # class_onehot[np.arange(batch_size), labels] = 1 13 | # class_onehot = T.from_numpy(class_onehot) 14 | # return class_onehot 15 | 16 | def weights_init(m): 17 | classname = m.__class__.__name__ 18 | if classname.find('Conv') != -1: 19 | m.weight.data.normal_(0.0, 0.02) 20 | elif classname.find('BatchNorm') != -1: 21 | m.weight.data.normal_(1.0, 0.02) 22 | m.bias.data.fill_(0) -------------------------------------------------------------------------------- /AE/pytorch_msssim/__init__.py: -------------------------------------------------------------------------------- 1 | from .ssim import ssim, ms_ssim, SSIM, MS_SSIM -------------------------------------------------------------------------------- /AE/pytorch_msssim/ssim.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | 4 | 5 | def _fspecial_gauss_1d(size, sigma): 6 | r"""Create 1-D gauss kernel 7 | Args: 8 | size (int): the size of gauss kernel 9 | sigma (float): sigma of normal distribution 10 | 11 | Returns: 12 | torch.Tensor: 1D kernel 13 | """ 14 | coords = torch.arange(size).to(dtype=torch.float) 15 | coords -= size//2 16 | 17 | g = torch.exp(-(coords**2) / (2*sigma**2)) 18 | g /= g.sum() 19 | 20 | return g.unsqueeze(0).unsqueeze(0) 21 | 22 | 23 | def gaussian_filter(input, win): 24 | r""" Blur input with 1-D kernel 25 | Args: 26 | input (torch.Tensor): a batch of tensors to be blured 27 | window (torch.Tensor): 1-D gauss kernel 28 | 29 | Returns: 30 | torch.Tensor: blured tensors 31 | """ 32 | 33 | N, C, H, W = input.shape 34 | out = F.conv2d(input, win, stride=1, padding=0, groups=C) 35 | out = F.conv2d(out, win.transpose(2, 3), stride=1, padding=0, groups=C) 36 | return out 37 | 38 | 39 | def _ssim(X, Y, win, data_range=255, size_average=True, full=False): 40 | r""" Calculate ssim index for X and Y 41 | Args: 42 | X (torch.Tensor): images 43 | Y (torch.Tensor): images 44 | win (torch.Tensor): 1-D gauss kernel 45 | data_range (float or int, optional): value range of input images. (usually 1.0 or 255) 46 | size_average (bool, optional): if size_average=True, ssim of all images will be averaged as a scalar 47 | full (bool, optional): return sc or not 48 | 49 | Returns: 50 | torch.Tensor: ssim results 51 | """ 52 | 53 | K1 = 0.01 54 | K2 = 0.03 55 | batch, channel, height, width = X.shape 56 | compensation = 1.0 57 | 58 | C1 = (K1 * data_range)**2 59 | C2 = (K2 * data_range)**2 60 | 61 | win = win.to(X.device, dtype=X.dtype) 62 | 63 | mu1 = gaussian_filter(X, win) 64 | mu2 = gaussian_filter(Y, win) 65 | 66 | mu1_sq = mu1.pow(2) 67 | mu2_sq = mu2.pow(2) 68 | mu1_mu2 = mu1 * mu2 69 | 70 | sigma1_sq = compensation * ( gaussian_filter(X * X, win) - mu1_sq ) 71 | sigma2_sq = compensation * ( gaussian_filter(Y * Y, win) - mu2_sq ) 72 | sigma12 = compensation * ( gaussian_filter(X * Y, win) - mu1_mu2 ) 73 | 74 | cs_map = (2 * sigma12 + C2) / (sigma1_sq + sigma2_sq + C2) 75 | ssim_map = ((2 * mu1_mu2 + C1) / (mu1_sq + mu2_sq + C1)) * cs_map 76 | 77 | if size_average: 78 | ssim_val = ssim_map.mean() 79 | cs = cs_map.mean() 80 | else: 81 | ssim_val = ssim_map.mean(-1).mean(-1).mean(-1) # reduce along CHW 82 | cs = cs_map.mean(-1).mean(-1).mean(-1) 83 | 84 | if full: 85 | return ssim_val, cs 86 | else: 87 | return ssim_val 88 | 89 | 90 | def ssim(X, Y, win_size=11, win_sigma=1.5, win=None, data_range=255, size_average=True, full=False): 91 | r""" interface of ssim 92 | Args: 93 | X (torch.Tensor): a batch of images, (N,C,H,W) 94 | Y (torch.Tensor): a batch of images, (N,C,H,W) 95 | win_size: (int, optional): the size of gauss kernel 96 | win_sigma: (float, optional): sigma of normal distribution 97 | win (torch.Tensor, optional): 1-D gauss kernel. if None, a new kernel will be created according to win_size and win_sigma 98 | data_range (float or int, optional): value range of input images. (usually 1.0 or 255) 99 | size_average (bool, optional): if size_average=True, ssim of all images will be averaged as a scalar 100 | full (bool, optional): return sc or not 101 | 102 | Returns: 103 | torch.Tensor: ssim results 104 | """ 105 | 106 | if len(X.shape) != 4: 107 | raise ValueError('Input images must 4-d tensor.') 108 | 109 | if not X.type() == Y.type(): 110 | raise ValueError('Input images must have the same dtype.') 111 | 112 | if not X.shape == Y.shape: 113 | raise ValueError('Input images must have the same dimensions.') 114 | 115 | if not (win_size % 2 == 1): 116 | raise ValueError('Window size must be odd.') 117 | 118 | win_sigma = win_sigma 119 | if win is None: 120 | win = _fspecial_gauss_1d(win_size, win_sigma) 121 | win = win.repeat(X.shape[1], 1, 1, 1) 122 | else: 123 | win_size = win.shape[-1] 124 | 125 | ssim_val, cs = _ssim(X, Y, 126 | win=win, 127 | data_range=data_range, 128 | size_average=False, 129 | full=True) 130 | if size_average: 131 | ssim_val = ssim_val.mean() 132 | cs = cs.mean() 133 | 134 | if full: 135 | return ssim_val, cs 136 | else: 137 | return ssim_val 138 | 139 | 140 | def ms_ssim(X, Y, win_size=11, win_sigma=1.5, win=None, data_range=255, size_average=True, full=False, weights=None): 141 | r""" interface of ms-ssim 142 | Args: 143 | X (torch.Tensor): a batch of images, (N,C,H,W) 144 | Y (torch.Tensor): a batch of images, (N,C,H,W) 145 | win_size: (int, optional): the size of gauss kernel 146 | win_sigma: (float, optional): sigma of normal distribution 147 | win (torch.Tensor, optional): 1-D gauss kernel. if None, a new kernel will be created according to win_size and win_sigma 148 | data_range (float or int, optional): value range of input images. (usually 1.0 or 255) 149 | size_average (bool, optional): if size_average=True, ssim of all images will be averaged as a scalar 150 | full (bool, optional): return sc or not 151 | weights (list, optional): weights for different levels 152 | 153 | Returns: 154 | torch.Tensor: ms-ssim results 155 | """ 156 | if len(X.shape) != 4: 157 | raise ValueError('Input images must 4-d tensor.') 158 | 159 | if not X.type() == Y.type(): 160 | raise ValueError('Input images must have the same dtype.') 161 | 162 | if not X.shape == Y.shape: 163 | raise ValueError('Input images must have the same dimensions.') 164 | 165 | if not (win_size % 2 == 1): 166 | raise ValueError('Window size must be odd.') 167 | 168 | if weights is None: 169 | weights = torch.FloatTensor( 170 | [0.0448, 0.2856, 0.3001, 0.2363, 0.1333]).to(X.device, dtype=X.dtype) 171 | 172 | win_sigma = win_sigma 173 | if win is None: 174 | win = _fspecial_gauss_1d(win_size, win_sigma) 175 | win = win.repeat(X.shape[1], 1, 1, 1) 176 | else: 177 | win_size = win.shape[-1] 178 | 179 | levels = weights.shape[0] 180 | mcs = [] 181 | for _ in range(levels): 182 | ssim_val, cs = _ssim(X, Y, 183 | win=win, 184 | data_range=data_range, 185 | size_average=False, 186 | full=True) 187 | mcs.append(cs) 188 | 189 | padding = (X.shape[2] % 2, X.shape[3] % 2) 190 | X = F.avg_pool2d(X, kernel_size=2, padding=padding) 191 | Y = F.avg_pool2d(Y, kernel_size=2, padding=padding) 192 | 193 | mcs = torch.stack(mcs, dim=0) # mcs, (level, batch) 194 | # weights, (level) 195 | msssim_val = torch.prod((mcs[:-1] ** weights[:-1].unsqueeze(1)) 196 | * (ssim_val ** weights[-1]), dim=0) # (batch, ) 197 | 198 | if size_average: 199 | msssim_val = msssim_val.mean() 200 | return msssim_val 201 | 202 | 203 | # Classes to re-use window 204 | class SSIM(torch.nn.Module): 205 | def __init__(self, win_size=11, win_sigma=1.5, data_range=None, size_average=True, channel=3): 206 | r""" class for ssim 207 | Args: 208 | win_size: (int, optional): the size of gauss kernel 209 | win_sigma: (float, optional): sigma of normal distribution 210 | data_range (float or int, optional): value range of input images. (usually 1.0 or 255) 211 | size_average (bool, optional): if size_average=True, ssim of all images will be averaged as a scalar 212 | channel (int, optional): input channels (default: 3) 213 | """ 214 | 215 | super(SSIM, self).__init__() 216 | self.win = _fspecial_gauss_1d( 217 | win_size, win_sigma).repeat(channel, 1, 1, 1) 218 | self.size_average = size_average 219 | self.data_range = data_range 220 | 221 | def forward(self, X, Y): 222 | return ssim(X, Y, win=self.win, data_range=self.data_range, size_average=self.size_average) 223 | 224 | 225 | class MS_SSIM(torch.nn.Module): 226 | def __init__(self, win_size=11, win_sigma=1.5, data_range=None, size_average=True, channel=3, weights=None): 227 | r""" class for ms-ssim 228 | Args: 229 | win_size: (int, optional): the size of gauss kernel 230 | win_sigma: (float, optional): sigma of normal distribution 231 | data_range (float or int, optional): value range of input images. (usually 1.0 or 255) 232 | size_average (bool, optional): if size_average=True, ssim of all images will be averaged as a scalar 233 | channel (int, optional): input channels (default: 3) 234 | weights (list, optional): weights for different levels 235 | """ 236 | 237 | super(MS_SSIM, self).__init__() 238 | self.win = _fspecial_gauss_1d( 239 | win_size, win_sigma).repeat(channel, 1, 1, 1) 240 | self.size_average = size_average 241 | self.data_range = data_range 242 | self.weights = weights 243 | 244 | def forward(self, X, Y): 245 | return ms_ssim(X, Y, win=self.win, size_average=self.size_average, data_range=self.data_range, weights=self.weights) 246 | -------------------------------------------------------------------------------- /AE/pytorch_ssim/__init__.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | from torch.autograd import Variable 4 | import numpy as np 5 | from math import exp 6 | 7 | def gaussian(window_size, sigma): 8 | gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)]) 9 | return gauss/gauss.sum() 10 | 11 | def create_window(window_size, channel): 12 | _1D_window = gaussian(window_size, 1.5).unsqueeze(1) 13 | _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) 14 | window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) 15 | return window 16 | 17 | def _ssim(img1, img2, window, window_size, channel, size_average = True): 18 | mu1 = F.conv2d(img1, window, padding = window_size//2, groups = channel) 19 | mu2 = F.conv2d(img2, window, padding = window_size//2, groups = channel) 20 | 21 | mu1_sq = mu1.pow(2) 22 | mu2_sq = mu2.pow(2) 23 | mu1_mu2 = mu1*mu2 24 | 25 | sigma1_sq = F.conv2d(img1*img1, window, padding = window_size//2, groups = channel) - mu1_sq 26 | sigma2_sq = F.conv2d(img2*img2, window, padding = window_size//2, groups = channel) - mu2_sq 27 | sigma12 = F.conv2d(img1*img2, window, padding = window_size//2, groups = channel) - mu1_mu2 28 | 29 | C1 = 0.01**2 30 | C2 = 0.03**2 31 | 32 | ssim_map = ((2*mu1_mu2 + C1)*(2*sigma12 + C2))/((mu1_sq + mu2_sq + C1)*(sigma1_sq + sigma2_sq + C2)) 33 | 34 | if size_average: 35 | return ssim_map.mean() 36 | else: 37 | return ssim_map.mean(1).mean(1).mean(1) 38 | 39 | class SSIM(torch.nn.Module): 40 | def __init__(self, window_size = 11, size_average = True): 41 | super(SSIM, self).__init__() 42 | self.window_size = window_size 43 | self.size_average = size_average 44 | self.channel = 1 45 | self.window = create_window(window_size, self.channel) 46 | 47 | def forward(self, img1, img2): 48 | (_, channel, _, _) = img1.size() 49 | 50 | if channel == self.channel and self.window.data.type() == img1.data.type(): 51 | window = self.window 52 | else: 53 | window = create_window(self.window_size, channel) 54 | 55 | if img1.is_cuda: 56 | window = window.cuda(img1.get_device()) 57 | window = window.type_as(img1) 58 | 59 | self.window = window 60 | self.channel = channel 61 | 62 | 63 | return _ssim(img1, img2, window, self.window_size, channel, self.size_average) 64 | 65 | def ssim(img1, img2, window_size = 11, size_average = True): 66 | (_, channel, _, _) = img1.size() 67 | window = create_window(window_size, channel) 68 | 69 | if img1.is_cuda: 70 | window = window.cuda(img1.get_device()) 71 | window = window.type_as(img1) 72 | 73 | return _ssim(img1, img2, window, window_size, channel, size_average) 74 | -------------------------------------------------------------------------------- /AE/train.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import optim,nn 3 | from torchvision.utils import save_image 4 | from autoencoder1 import ResNet_autoencoder,Bottleneck,DeconvBottleneck,weights_init_kaiming 5 | from dataset import get_data_loader 6 | from help_fucntions import get_cuda 7 | import time 8 | import os 9 | import numpy as np 10 | import pickle 11 | import pytorch_ssim 12 | 13 | 14 | # parameters 15 | im_size = 256 16 | img_path = "/mnt/mfs/yiling/EL_surface" 17 | batch_size = 48 18 | save_path = "ae_model.pth" 19 | train_epoch = 80 20 | resume_training = False 21 | z_size = 512 22 | 23 | def load_model_from_checkpoint(): 24 | global ae_model 25 | checkpoint = torch.load(save_path) 26 | vae.load_state_dict(checkpoint['ae_model']) 27 | return checkpoint['epoch'] 28 | 29 | 30 | 31 | # load dataSet 32 | train_loader = get_data_loader(img_size=im_size,img_path=img_path,shuffle = True,batch_size=batch_size,num_workers=8) 33 | 34 | loss_function = pytorch_ssim.SSIM() 35 | 36 | 37 | 38 | def main(): 39 | #build Neural Network 40 | ae_model = ResNet_autoencoder(Bottleneck, DeconvBottleneck, [3, 4, 6, 3], 3) 41 | ae_model.apply(weights_init_kaiming) 42 | if torch.cuda.is_available(): 43 | torch.backends.cudnn.benchmark = True 44 | ae_model.cuda() 45 | ae_model.train() 46 | start_epoch = 0 47 | if resume_training: 48 | start_epoch = load_model_from_checkpoint() 49 | 50 | lr = 0.0005 51 | ae_model_optimizer = optim.Adam(ae_model.parameters(), lr=lr, betas=(0.5, 0.999), weight_decay=1e-5) 52 | device_ids = range(torch.cuda.device_count()) 53 | ae_model = nn.DataParallel(ae_model, device_ids) 54 | 55 | for epoch in range(start_epoch,train_epoch): 56 | ae_model.train() 57 | ssim_loss = 0 58 | epoch_start_time = time.time() 59 | if (epoch + 1) % 20 == 0: 60 | ae_model_optimizer.param_groups[0]['lr'] /= 4 61 | print("learning rate change!") 62 | i = 0 63 | for x, _ in train_loader: 64 | ae_model.train() 65 | ae_model.zero_grad() 66 | x = get_cuda(x) 67 | rec = ae_model(x) 68 | loss =1 - loss_function(x, rec[1]) 69 | loss.backward() 70 | ae_model_optimizer.step() 71 | ssim_loss += loss.item() 72 | os.makedirs('results_rec', exist_ok=True) 73 | epoch_end_time = time.time() 74 | per_epoch_ptime = epoch_end_time - epoch_start_time 75 | 76 | # report losses and save samples each 60 iterations 77 | m = 60 78 | i += 1 79 | if i % m == 0: 80 | ssim_loss /= m 81 | print('\n[%d/%d] - ptime: %.2f, ssim_loss: %.9f' % ((epoch + 1), train_epoch, per_epoch_ptime, ssim_loss)) 82 | ssim_loss = 0 83 | with torch.no_grad(): 84 | ae_model.eval() 85 | x_rec = ae_model(x) 86 | resultsample = torch.cat([x, x_rec[1]]) 87 | resultsample = resultsample.cpu() 88 | save_image(resultsample.view(-1, 3, im_size, im_size), 89 | 'results_rec/sample_' + str(epoch) + "_" + str(i) + '.png') 90 | torch.save({ 91 | 'epoch': epoch, 92 | "ae_model": model.state_dict(), 93 | }, save_path) 94 | 95 | print("Training finish!... save training results") 96 | torch.save(model.state_dict(), "ae_model.pkl") 97 | 98 | if __name__ == '__main__': 99 | main() -------------------------------------------------------------------------------- /AE/train.sh: -------------------------------------------------------------------------------- 1 | CUDA_VISIBLE_DEVICES=0,1,2,3 python train.py -------------------------------------------------------------------------------- /AutoencodersForUnsupervisedAnomalyDetection/Deep Autoeoncoding Models for Unsupervised Anomaly Detection in Brain MR Images.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ForrestPi/Unsupervised-Defect-Segmentation/e366ac7c757bb1b45f38ebbc502dfee7ccb72398/AutoencodersForUnsupervisedAnomalyDetection/Deep Autoeoncoding Models for Unsupervised Anomaly Detection in Brain MR Images.pdf -------------------------------------------------------------------------------- /GANomaly-pytorch/README.md: -------------------------------------------------------------------------------- 1 | # GANomaly-pytorch 2 | 本repo是论文[GANomaly: Semi-Supervised Anomaly Detection via Adversarial Training](https://arxiv.org/abs/1805.06725?context=cs) 的一个pytorch实现。 3 | 根据论文作者的[repo](https://github.com/samet-akcay/ganomaly)改写,原repo中脚本与相关数据集深度耦合,本repo没有复现其论文结果,而是使其可以灵活用于自定义的数据集,易于提取合适的threshold. 4 | 5 | ### from 6 | https://github.com/WellenWoo/GANomaly-pytorch 7 | 8 | -------------------------------------------------------------------------------- /GANomaly-pytorch/loss.py: -------------------------------------------------------------------------------- 1 | """ 2 | Losses 3 | """ 4 | import torch 5 | 6 | def l1_loss(input, target): 7 | """ L1 Loss without reduce flag. 8 | 9 | Args: 10 | input (FloatTensor): Input tensor 11 | target (FloatTensor): Output tensor 12 | 13 | Returns: 14 | [FloatTensor]: L1 distance between input and output 15 | """ 16 | 17 | return torch.mean(torch.abs(input - target)) 18 | 19 | def l2_loss(input, target, size_average=True): 20 | """ L2 Loss without reduce flag. 21 | 22 | Args: 23 | input (FloatTensor): Input tensor 24 | target (FloatTensor): Output tensor 25 | 26 | Returns: 27 | [FloatTensor]: L2 distance between input and output 28 | """ 29 | if size_average: 30 | return torch.mean(torch.pow((input-target), 2)) 31 | else: 32 | return torch.pow((input-target), 2) 33 | -------------------------------------------------------------------------------- /GANomaly-pytorch/nets.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | @author: Administrator 4 | """ 5 | import torch 6 | import torch.nn as nn 7 | 8 | def weights_init(mod): 9 | """ 10 | Custom weights initialization called on netG, netD and netE 11 | :param m: 12 | :return: 13 | """ 14 | classname = mod.__class__.__name__ 15 | if classname.find('Conv') != -1: 16 | mod.weight.data.normal_(0.0, 0.02) 17 | elif classname.find('BatchNorm') != -1: 18 | mod.weight.data.normal_(1.0, 0.02) 19 | mod.bias.data.fill_(0) 20 | # print(type(classname)) 21 | # print(classname) 22 | 23 | class Encoder(nn.Module): 24 | """ 25 | DCGAN ENCODER NETWORK 26 | 编码器; 27 | """ 28 | def __init__(self, isize, nz, nc, ndf, ngpu, n_extra_layers=0, add_final_conv=True): 29 | super(Encoder, self).__init__() 30 | self.ngpu = ngpu 31 | assert isize % 16 == 0, "isize has to be a multiple of 16" 32 | 33 | main = nn.Sequential() 34 | # input is nc x isize x isize 35 | main.add_module('initial-conv-{0}-{1}'.format(nc, ndf), 36 | nn.Conv2d(nc, ndf, 4, 2, 1, bias=False)) 37 | main.add_module('initial-relu-{0}'.format(ndf), 38 | nn.LeakyReLU(0.2, inplace=True)) 39 | csize, cndf = isize / 2, ndf 40 | 41 | # Extra layers 42 | for t in range(n_extra_layers): 43 | main.add_module('extra-layers-{0}-{1}-conv'.format(t, cndf), 44 | nn.Conv2d(cndf, cndf, 3, 1, 1, bias=False)) 45 | main.add_module('extra-layers-{0}-{1}-batchnorm'.format(t, cndf), 46 | nn.BatchNorm2d(cndf)) 47 | main.add_module('extra-layers-{0}-{1}-relu'.format(t, cndf), 48 | nn.LeakyReLU(0.2, inplace=True)) 49 | 50 | while csize > 4: 51 | in_feat = cndf 52 | out_feat = cndf * 2 53 | main.add_module('pyramid-{0}-{1}-conv'.format(in_feat, out_feat), 54 | nn.Conv2d(in_feat, out_feat, 4, 2, 1, bias=False)) 55 | main.add_module('pyramid-{0}-batchnorm'.format(out_feat), 56 | nn.BatchNorm2d(out_feat)) 57 | main.add_module('pyramid-{0}-relu'.format(out_feat), 58 | nn.LeakyReLU(0.2, inplace=True)) 59 | cndf = cndf * 2 60 | csize = csize / 2 61 | 62 | # state size. K x 4 x 4 63 | if add_final_conv: 64 | main.add_module('final-{0}-{1}-conv'.format(cndf, 1), 65 | nn.Conv2d(cndf, nz, 4, 1, 0, bias=False)) 66 | self.main = main 67 | 68 | def forward(self, data): 69 | if self.ngpu > 1: 70 | output = nn.parallel.data_parallel(self.main, data, range(self.ngpu)) 71 | else: 72 | output = self.main(data) 73 | 74 | return output 75 | 76 | class Decoder(nn.Module): 77 | """ 78 | DCGAN DECODER NETWORK 79 | 解码器; 80 | """ 81 | def __init__(self, isize, nz, nc, ngf, ngpu, n_extra_layers=0): 82 | super(Decoder, self).__init__() 83 | self.ngpu = ngpu 84 | assert isize % 16 == 0, "isize has to be a multiple of 16" 85 | 86 | cngf, tisize = ngf // 2, 4 87 | while tisize != isize: 88 | cngf = cngf * 2 89 | tisize = tisize * 2 90 | 91 | main = nn.Sequential() 92 | # input is Z, going into a convolution 93 | main.add_module('initial-{0}-{1}-convt'.format(nz, cngf), 94 | nn.ConvTranspose2d(nz, cngf, 4, 1, 0, bias=False)) 95 | main.add_module('initial-{0}-batchnorm'.format(cngf), 96 | nn.BatchNorm2d(cngf)) 97 | main.add_module('initial-{0}-relu'.format(cngf), 98 | nn.ReLU(True)) 99 | 100 | csize, _ = 4, cngf 101 | while csize < isize // 2: 102 | main.add_module('pyramid-{0}-{1}-convt'.format(cngf, cngf // 2), 103 | nn.ConvTranspose2d(cngf, cngf // 2, 4, 2, 1, bias=False)) 104 | main.add_module('pyramid-{0}-batchnorm'.format(cngf // 2), 105 | nn.BatchNorm2d(cngf // 2)) 106 | main.add_module('pyramid-{0}-relu'.format(cngf // 2), 107 | nn.ReLU(True)) 108 | cngf = cngf // 2 109 | csize = csize * 2 110 | 111 | # Extra layers 112 | for t in range(n_extra_layers): 113 | main.add_module('extra-layers-{0}-{1}-conv'.format(t, cngf), 114 | nn.Conv2d(cngf, cngf, 3, 1, 1, bias=False)) 115 | main.add_module('extra-layers-{0}-{1}-batchnorm'.format(t, cngf), 116 | nn.BatchNorm2d(cngf)) 117 | main.add_module('extra-layers-{0}-{1}-relu'.format(t, cngf), 118 | nn.ReLU(True)) 119 | 120 | main.add_module('final-{0}-{1}-convt'.format(cngf, nc), 121 | nn.ConvTranspose2d(cngf, nc, 4, 2, 1, bias=False)) 122 | main.add_module('final-{0}-tanh'.format(nc), 123 | nn.Tanh()) 124 | self.main = main 125 | 126 | def forward(self, data): 127 | if self.ngpu > 1: 128 | output = nn.parallel.data_parallel(self.main, data, range(self.ngpu)) 129 | else: 130 | output = self.main(data) 131 | return output 132 | 133 | class NetD(nn.Module): 134 | """ 135 | DISCRIMINATOR NETWORK 136 | 判别器网络; 137 | """ 138 | def __init__(self, isize, nc, ngf, ngpu, extralayers): 139 | super(NetD, self).__init__() 140 | model = Encoder(isize, 1, nc, ngf, ngpu, extralayers) 141 | layers = list(model.main.children()) 142 | 143 | self.features = nn.Sequential(*layers[:-1]) 144 | self.classifier = nn.Sequential(layers[-1]) 145 | self.classifier.add_module('Sigmoid', nn.Sigmoid()) 146 | 147 | def forward(self, x): 148 | features = self.features(x) 149 | features = features 150 | classifier = self.classifier(features) 151 | classifier = classifier.view(-1, 1).squeeze(1) 152 | 153 | return classifier, features 154 | 155 | class NetG(nn.Module): 156 | """ 157 | GENERATOR NETWORK 158 | 生成器网络; 159 | """ 160 | def __init__(self, isize, nz, nc, ngf, ngpu, extralayers): 161 | super(NetG, self).__init__() 162 | self.encoder1 = Encoder(isize, nz, nc, ngf, ngpu, extralayers) 163 | self.decoder = Decoder(isize, nz, nc, ngf, ngpu, extralayers) 164 | self.encoder2 = Encoder(isize, nz, nc, ngf, ngpu, extralayers) 165 | 166 | def forward(self, x): 167 | latent_i = self.encoder1(x) 168 | gen_imag = self.decoder(latent_i) 169 | latent_o = self.encoder2(gen_imag) 170 | return gen_imag, latent_i, latent_o -------------------------------------------------------------------------------- /GANomaly-pytorch/reporter.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | @author: WellenWoo 4 | """ 5 | import numpy as np 6 | import itertools 7 | import matplotlib.pyplot as plt 8 | from sklearn.metrics import classification_report, confusion_matrix, roc_curve, auc 9 | import torch 10 | import time 11 | 12 | class Clf_Reporter(object): 13 | """用于分类模型的预测结果评估器;""" 14 | def clf_report(self, y_test, y_pred): 15 | """分类报告; 16 | 只能用于分类任务, 17 | 不能用于回归任务;""" 18 | score = classification_report(y_test, y_pred) 19 | return score 20 | 21 | def matrix(self, y_test, y_pred): 22 | """混淆矩阵; 23 | 用于分类任务;""" 24 | score = confusion_matrix(y_test, y_pred) 25 | return score 26 | 27 | def roc(self, labels, scores): 28 | """Compute ROC curve and ROC area for each class""" 29 | fpr = dict() 30 | tpr = dict() 31 | roc_auc = dict() 32 | 33 | labels = labels.cpu() 34 | scores = scores.cpu() 35 | 36 | # True/False Positive Rates.\ 37 | fpr, tpr, _ = roc_curve(labels, scores) 38 | roc_auc = auc(fpr, tpr) 39 | 40 | return roc_auc 41 | 42 | def plot_cfm(self, cfm, classes = ["ng", "ok"], 43 | normalize=False, 44 | title='Confusion matrix', 45 | cmap=plt.cm.Blues): 46 | """ 47 | This function prints and plots the confusion matrix. 48 | Normalization can be applied by setting `normalize=True`. 49 | 绘制混淆矩阵图; 50 | args: 51 | cfm: 混淆矩阵, np.ndarray, 52 | classes: 类别列表, list, i.e. ["ng", "ok"], 53 | normalize: 是否归一化,bool,归一化则以百分比显示; 54 | """ 55 | if normalize: 56 | cfm = cfm.astype('float') / cfm.sum(axis=1)[:, np.newaxis] 57 | print("Normalized confusion matrix") 58 | else: 59 | print('Confusion matrix, without normalization') 60 | 61 | print(cfm) 62 | 63 | plt.imshow(cfm, interpolation='nearest', cmap=cmap) 64 | plt.title(title) 65 | plt.colorbar() 66 | tick_marks = np.arange(len(classes)) 67 | plt.xticks(tick_marks, classes, rotation=45) 68 | plt.yticks(tick_marks, classes) 69 | 70 | fmt = '.2f' if normalize else 'd' 71 | thresh = cfm.max() / 2. 72 | for i, j in itertools.product(range(cfm.shape[0]), range(cfm.shape[1])): 73 | plt.text(j, i, format(cfm[i, j], fmt), 74 | horizontalalignment="center", 75 | color="white" if cfm[i, j] > thresh else "black") 76 | 77 | plt.tight_layout() 78 | plt.ylabel('True label') 79 | plt.xlabel('Predicted label') 80 | plt.show() 81 | 82 | def get_threshold(self, y_test, pred, 83 | threshold_rang = [0.05,0.9], step = 0.01, 84 | plot = True): 85 | """获取最佳置信度,目前仅适用于二分类. 86 | 87 | args: 88 | ---------- 89 | y_test: 真实的分类标签, np.ndarray, 90 | 91 | pred: 预测的便签, np.ndarray, 92 | 93 | threshold_rang: 置信度范围, list, 94 | 95 | step: 步长, float; 96 | retrun: 97 | ---------- 98 | val: 最佳置信度对应的正确分类样本数, 99 | 100 | data[val]: 最佳置信度.""" 101 | pred_true_0 = [] 102 | pred_true_1 = [] 103 | pred_true_sum = [] 104 | 105 | threshold_rang = np.arange(*threshold_rang, step) 106 | data = dict() 107 | 108 | for th in threshold_rang: 109 | tmp = self.prob2int(pred, th) 110 | mat = self.matrix(y_test, tmp) 111 | 112 | pred_true_0.append(mat[0][0]) 113 | pred_true_1.append(mat[1][1]) 114 | 115 | sums = np.sum([mat[0][0], mat[1][1]]) 116 | pred_true_sum.append(sums) 117 | 118 | data[str(th)] = sums 119 | if plot: 120 | plot_true_pred(threshold_rang, pred_true_0, pred_true_1, pred_true_sum) 121 | 122 | th, val = get_dict_max_val(data) 123 | return th, val, data 124 | 125 | def plot_true_pred(rang, true0, true1, true_sum): 126 | fig,ax = plt.subplots() 127 | ax.plot(rang, true0, label = "true0") 128 | ax.plot(rang, true1, label = "true1") 129 | ax.plot(rang, true_sum, label = "true_sum") 130 | 131 | ax.set_xlabel("confidence") 132 | ax.set_ylabel("samples") 133 | 134 | plt.grid() 135 | plt.legend() 136 | plt.show() 137 | 138 | def get_dict_max_val(d): 139 | """获取字典中的最大值及其对应的键值; 140 | args: 141 | ---- 142 | d: 字典, dict, 143 | 144 | return: 145 | --- 146 | val: 最大值所对应的键, str, 147 | 148 | max_val: 最大值, int or float;""" 149 | max_val = max(d.values()) 150 | for i in d.keys(): 151 | if d[i] == max_val: 152 | break 153 | return i, max_val 154 | 155 | def print_current_performance(performance, best): 156 | """ Print current performance results. 157 | Args: 158 | performance ([OrderedDict]): Performance of the model 159 | best ([int]): Best performance. 160 | """ 161 | message = ' ' 162 | log_name = r"train_log.txt" 163 | dt = time.gmtime() 164 | message += str(dt.tm_year) +"-"+str(dt.tm_mon)+"-"+str(dt.tm_mday)+"-"+str(dt.tm_hour+8)+":"+str(dt.tm_min)+":"+str(dt.tm_sec) 165 | message += "roc: %.3f#######" % performance 166 | message += 'max AUC: %.3f' % best 167 | 168 | print(message) 169 | with open(log_name, "a") as log_file: 170 | log_file.write('%s\n' % message) 171 | -------------------------------------------------------------------------------- /GANomaly-pytorch/utils.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | @author: WellenWoo 4 | """ 5 | import torchvision as tv 6 | import torch 7 | import numpy as np 8 | 9 | def files2sample(fpath,img_size = 26,batch_size = 64, workers = 4, drop = False): 10 | """将文件夹中不同子文件夹的图片转为数据集; 11 | 传入的参数为r'sample'; 12 | 子文件夹分别为r'sample/ok_sam',r'sample/ng_sam'; 13 | args: 14 | fpath: str,图片存放的路径, 15 | img_size: seq or int, 图片缩放的尺寸, 16 | batch_size: int,每次批图片的数量, 17 | workers: int, 进程数量, 18 | return: 19 | torch.utils.data.dataloader.DataLoader; 20 | """ 21 | trans = tv.transforms.Compose([ 22 | tv.transforms.Resize(size = img_size),# 缩放图片(Image)，保持长宽比不变，最短边为img_size像素 23 | # tv.transforms.CenterCrop(img_size),# 从图片中间切出img_size*img_size的图片 24 | tv.transforms.ToTensor(),# 将图片(Image)转成Tensor，归一化至[0, 1] 25 | tv.transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))])# 标准化至[-1, 1]，规定均值和标准差 26 | 27 | dataset = tv.datasets.ImageFolder(fpath, transform=trans) 28 | 29 | dataloader = torch.utils.data.DataLoader(dataset, 30 | batch_size=batch_size, 31 | shuffle=False, 32 | num_workers= workers, 33 | drop_last = drop) #drap_last = True:如果样本总量除于batch_size后有余数,则丢弃余数部分 34 | return dataloader 35 | 36 | def loader2tensor(loader): 37 | """从dataloader中提取图像数据和便签数据, 38 | 再组装成torch.Tensor返回. 39 | args: 40 | ---- 41 | loader: 已经包含标签的数据,torch.utils.data.dataloader.DataLoader, 42 | 43 | return: 44 | ---- 45 | x: 图像数据,torch.Tensor,[batch,channel,h, w], 46 | y: 便签数据,torch.Tensor,[batch].""" 47 | xs, ys = [], [] 48 | for i, data in enumerate(loader): 49 | xs.append(data[0]) 50 | ys.append(data[1]) 51 | x = torch.cat(tuple(xs)) 52 | y = torch.cat(tuple(ys)) 53 | return x, y 54 | 55 | def data2np(data): 56 | """将cuda或cpu中的Tensor 数据转为cpu中的np数据; 57 | args: 58 | ---- 59 | data: torch.Tensor, 60 | 61 | return: 62 | ---- 63 | data: np.nadrray.""" 64 | device = torch.device("cpu") 65 | if isinstance(data, torch.Tensor): 66 | if not data.device == device: 67 | data = data.cpu() 68 | data = data.numpy() 69 | return data 70 | 71 | def prob2int(y_pred, threshold): 72 | """将预测的概率值转为类别; 73 | args: 74 | y_pred: 预测标签的概率值, np.ndarray, 75 | 76 | threshold: 阈值, float,[0, 1], 77 | return: 78 | val: 预测的标签, np.ndarray, dtype=np.int.""" 79 | val = np.empty_like(y_pred) 80 | val[y_pred >= threshold] = 1 81 | val[y_pred < threshold] = 0 82 | return val 83 | 84 | def get_primitive_th(th, pred): 85 | """根据缩放后的预测阈值来计算出原始预测阈值, 86 | args: 87 | ---- 88 | th: 缩放后的预测阈值,float, 89 | pred: 原始的预测概率序列, torch.Tensor. 90 | 91 | return: 92 | ---- 93 | th: 原始预测阈值,float.""" 94 | return (torch.max(pred) - torch.min(pred)) * th + torch.min(pred) 95 | 96 | def get_eval_th(th, pred): 97 | """根据测试集的原始预测阈值计算验证集的缩放预测阈值, 98 | args: 99 | ---- 100 | th: 测试集的原始预测阈值,float, 101 | pred: 验证集的原始预测概率序列,torch.Tensor, 102 | 103 | return: 104 | ---- 105 | th:验证集的缩放预测阈值,float.""" 106 | return (th - torch.min(pred))/(torch.max(pred) - torch.min(pred)) -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2019 Forrest-Zhu 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # DeeCamp-10Group-Unsupervised-Defect-Segmentation 2 | 3 | # Abstract 4 | Deep learning is well established in the mainstream of computer vision, but it relies on large amounts of data. 5 | especially in the industrial field, such as product surface defect detection task. Mainstream neural networks framework 6 | (detection, semantic segmentation) must have defect samples (all defect types to be detected) for training, and the number 7 | of defect samples should not be too small. Although defect learning is relatively easy, heavy defect collection and labeling 8 | work is introduced. Moreover, in some production processes in the industrial field, it is difficult to provide sufficient 9 | defect samples for training with a low defect rate, and it is difficult to meet the 'data-hungry deep'learning model.In contrast, 10 | there are plenty of positive samples in industrial production that have no defects at all. We use this characteristic 11 | of industrial production to gain the ability to detect bad samples by learning only good samples. This defect detection technology 12 | based on positive samples can be applied to AI at low cost in more and more complex industrial scenarios, so it has great practical significance. 13 | 14 | # Introduction 15 | In this work, we focus on unsupervised defect segmentation for visual inspection. 16 | The goal is to segment defective regions in images after having trained exclusively on 17 | non-defective samples. It has been shown that architectures based on convolutional neural networks (CNNs) such 18 | as autoencoders or generative adversarial networks can be used for this task. 19 | The input picture x into the autoencoder or GAN, get the restored image y. 20 | Then compare the features of each pixel of x and y, where the feature difference between x and y is large, that is 21 | the defect. 22 | 23 | # Architecture 24 | ![architecture](doc/architecture.jpg) 25 | 26 | # Directory 27 | - `/config`: This repo includes configuration file for training and evaluation. The parameters are saved in `.jason` file. 28 | - `/db`: This repo includes codes of data set tools 29 | - `/model`: Our networks codes are in `model/netwroks` and segmentation codes are in `mode/segmentation`. 30 | - `/tools`: This repo includes codes of auxiliary function just like timer, log writer and so on. 31 | -------------------------------------------------------------------------------- /anomaly-detection-image-completion/.gitignore: -------------------------------------------------------------------------------- 1 | **/.DS_Store 2 | **/__pycache__ 3 | **/.ipynb_checkpoints 4 | **/.pyc 5 | **/.idea 6 | -------------------------------------------------------------------------------- /anomaly-detection-image-completion/README.md: -------------------------------------------------------------------------------- 1 | # Welcome 2 | This repository implements the approach to detect surface anomalies in images presented in the paper [Anomaly Detection using Deep Learning based Image Completion](https://arxiv.org/pdf/1811.06861.pdf). 3 | ### from 4 | https://github.com/oezguensi/anomaly-detection-image-completion 5 | http://tongtianta.site/paper/51924 -------------------------------------------------------------------------------- /anomaly-detection-image-completion/main.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": null, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "import numpy as np\n", 10 | "import matplotlib.pyplot as plt\n", 11 | "import cv2\n", 12 | "import sys\n", 13 | "\n", 14 | "from collections import defaultdict\n", 15 | "from time import time\n", 16 | "from os import makedirs\n", 17 | "from os.path import join, isdir\n", 18 | "from glob import glob\n", 19 | "from keras.callbacks import TensorBoard\n", 20 | "\n", 21 | "sys.path.append('scripts')\n", 22 | "from model_helpers import *\n", 23 | "from data_helpers import *\n", 24 | "from train_helpers import *\n", 25 | "\n", 26 | "%load_ext autoreload\n", 27 | "%autoreload 2\n", 28 | "\n", 29 | "os.environ['CUDA_VISIBLE_DEVICES']='1'" 30 | ] 31 | }, 32 | { 33 | "cell_type": "code", 34 | "execution_count": null, 35 | "metadata": {}, 36 | "outputs": [], 37 | "source": [ 38 | "IMGS_PATH = '/data'\n", 39 | "OUT_PATH = 'saved/logs'\n", 40 | "if not isdir(OUT_PATH):\n", 41 | " makedirs(OUT_PATH)" 42 | ] 43 | }, 44 | { 45 | "cell_type": "markdown", 46 | "metadata": {}, 47 | "source": [ 48 | "# Hyperparameters" 49 | ] 50 | }, 51 | { 52 | "cell_type": "code", 53 | "execution_count": null, 54 | "metadata": {}, 55 | "outputs": [], 56 | "source": [ 57 | "patch_size = (128, 128)\n", 58 | "center_size = (32, 32)\n", 59 | "num_train_test = num_test_test = 10\n", 60 | "batch_size = 128\n", 61 | "num_epochs = 500\n", 62 | "model_width = 2" 63 | ] 64 | }, 65 | { 66 | "cell_type": "markdown", 67 | "metadata": {}, 68 | "source": [ 69 | "# Load data" 70 | ] 71 | }, 72 | { 73 | "cell_type": "code", 74 | "execution_count": null, 75 | "metadata": {}, 76 | "outputs": [], 77 | "source": [ 78 | "files = np.random.choice(glob(join(IMGS_PATH, '*jpg')), 10000, replace=False)\n", 79 | "train_files = files[:-(num_train_test + num_test_test)]\n", 80 | "train_test_files, test_test_files = files[-(num_train_test + num_test_test):-num_test_test], files[-num_test_test:]\n", 81 | "fake_files = glob('data/fake_files/*jpg')" 82 | ] 83 | }, 84 | { 85 | "cell_type": "code", 86 | "execution_count": null, 87 | "metadata": {}, 88 | "outputs": [], 89 | "source": [ 90 | "imgss = defaultdict(list)\n", 91 | "start = time()\n", 92 | "for dataset, files in zip(['train', 'train_test', 'test_test', 'fake'], \n", 93 | " [train_files, train_test_files, test_test_files, fake_files]):\n", 94 | " for i, f in enumerate(files):\n", 95 | " imgss[dataset].append(cv2.imread(f, 0))\n", 96 | " if i % 1000 == 0 and i != 0:\n", 97 | " end = time()\n", 98 | " print('Processing dataset {}: {:.0f} sec - ETA: {:.0f} sec'.format(\n", 99 | " dataset, end-start, ((end-start) / (i + 1)) * (len(files) - i)))" 100 | ] 101 | }, 102 | { 103 | "cell_type": "code", 104 | "execution_count": null, 105 | "metadata": {}, 106 | "outputs": [], 107 | "source": [ 108 | "fig, ax = plt.subplots(1, len(imgss.keys()), figsize=(20, 30))\n", 109 | "for i, (dataset, imgs) in enumerate(imgss.items()):\n", 110 | " ax[i].imshow(imgs[0], 'gray')\n", 111 | " ax[i].set_title('{} image'.format(dataset))" 112 | ] 113 | }, 114 | { 115 | "cell_type": "code", 116 | "execution_count": null, 117 | "metadata": {}, 118 | "outputs": [], 119 | "source": [ 120 | "plt.figure(figsize=(30, 20))\n", 121 | "plt.imshow(imgss['fake'][1], 'gray')" 122 | ] 123 | }, 124 | { 125 | "cell_type": "markdown", 126 | "metadata": {}, 127 | "source": [ 128 | "## Square images and resize to same size" 129 | ] 130 | }, 131 | { 132 | "cell_type": "code", 133 | "execution_count": null, 134 | "metadata": {}, 135 | "outputs": [], 136 | "source": [ 137 | "for dataset, imgs in imgss.items():\n", 138 | " res_imgs = [img[:min(img.shape), :min(img.shape)] for img in imgs] \n", 139 | " min_size = np.min([img.shape[0] for img in res_imgs])\n", 140 | "\n", 141 | " res_imgs = [cv2.resize(img, (min_size, min_size), interpolation=cv2.INTER_CUBIC) for img in res_imgs]\n", 142 | " imgss[dataset] = np.expand_dims(res_imgs, axis=3)" 143 | ] 144 | }, 145 | { 146 | "cell_type": "markdown", 147 | "metadata": {}, 148 | "source": [ 149 | "## Create training generator and evaluation images" 150 | ] 151 | }, 152 | { 153 | "cell_type": "code", 154 | "execution_count": null, 155 | "metadata": {}, 156 | "outputs": [], 157 | "source": [ 158 | "train_generator = DataGenerator(imgss['train'], patch_size, center_size, batch_size=batch_size, shuffle=True)" 159 | ] 160 | }, 161 | { 162 | "cell_type": "code", 163 | "execution_count": null, 164 | "metadata": {}, 165 | "outputs": [], 166 | "source": [ 167 | "x_test = []\n", 168 | "for patches in patchess.values():\n", 169 | " x_test.extend(patches)\n", 170 | "x_test = np.array(x_test)" 171 | ] 172 | }, 173 | { 174 | "cell_type": "code", 175 | "execution_count": null, 176 | "metadata": {}, 177 | "outputs": [], 178 | "source": [ 179 | "mask = create_center_mask(patch_size, center_size)\n", 180 | "patchess = defaultdict(list)\n", 181 | "for i, (dataset, imgs) in enumerate(imgss.items()):\n", 182 | " if dataset != 'train':\n", 183 | " for j, img in enumerate(imgs):\n", 184 | " patchess[dataset].append(prepare_patch(img, patch_size, mask, i + j))" 185 | ] 186 | }, 187 | { 188 | "cell_type": "code", 189 | "execution_count": null, 190 | "metadata": {}, 191 | "outputs": [], 192 | "source": [ 193 | "plt.imshow(np.squeeze(patchess['fake'][3]), 'gray')" 194 | ] 195 | }, 196 | { 197 | "cell_type": "markdown", 198 | "metadata": {}, 199 | "source": [ 200 | "# Create model" 201 | ] 202 | }, 203 | { 204 | "cell_type": "code", 205 | "execution_count": null, 206 | "metadata": {}, 207 | "outputs": [], 208 | "source": [ 209 | "model = create_anomaly_cnn(model_width=model_width)" 210 | ] 211 | }, 212 | { 213 | "cell_type": "code", 214 | "execution_count": null, 215 | "metadata": {}, 216 | "outputs": [], 217 | "source": [ 218 | "model.patch_size = patch_size\n", 219 | "model.center_size = center_size\n", 220 | "model.batch_size = batch_size\n", 221 | "model.num_epochs = num_epochs\n", 222 | "model.model_width = model_width" 223 | ] 224 | }, 225 | { 226 | "cell_type": "code", 227 | "execution_count": null, 228 | "metadata": { 229 | "scrolled": true 230 | }, 231 | "outputs": [], 232 | "source": [ 233 | "model.summary()" 234 | ] 235 | }, 236 | { 237 | "cell_type": "code", 238 | "execution_count": null, 239 | "metadata": {}, 240 | "outputs": [], 241 | "source": [ 242 | "model.compile(optimizer='adam', loss=reconstruction_loss(patch_size, mask=mask), metrics=['accuracy'])" 243 | ] 244 | }, 245 | { 246 | "cell_type": "code", 247 | "execution_count": null, 248 | "metadata": {}, 249 | "outputs": [], 250 | "source": [ 251 | "images_callback = TensorBoardImages(OUT_PATH, patchess, vis_every=1)\n", 252 | "checkpoint_callback = CustomModelCheckpoint(OUT_PATH, save_weights_only=False)\n", 253 | "losses_callback = TensorBoard(log_dir=OUT_PATH, batch_size=batch_size)" 254 | ] 255 | }, 256 | { 257 | "cell_type": "code", 258 | "execution_count": null, 259 | "metadata": {}, 260 | "outputs": [], 261 | "source": [ 262 | "history_dict = model.fit_generator(train_generator, \n", 263 | " validation_data=(x_test, x_test),\n", 264 | " epochs=2, \n", 265 | " verbose=1, \n", 266 | " callbacks=[images_callback, checkpoint_callback, losses_callback], \n", 267 | " workers=10, \n", 268 | " use_multiprocessing=True, \n", 269 | " shuffle=False, \n", 270 | " initial_epoch=0)" 271 | ] 272 | } 273 | ], 274 | "metadata": { 275 | "kernelspec": { 276 | "display_name": "Python (my-env)", 277 | "language": "python", 278 | "name": "myenv" 279 | }, 280 | "language_info": { 281 | "codemirror_mode": { 282 | "name": "ipython", 283 | "version": 3 284 | }, 285 | "file_extension": ".py", 286 | "mimetype": "text/x-python", 287 | "name": "python", 288 | "nbconvert_exporter": "python", 289 | "pygments_lexer": "ipython3", 290 | "version": "3.6.7" 291 | } 292 | }, 293 | "nbformat": 4, 294 | "nbformat_minor": 2 295 | } 296 | -------------------------------------------------------------------------------- /anomaly-detection-image-completion/scripts/data_helpers.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from keras.utils import Sequence 4 | 5 | 6 | def create_center_mask(patch_size, center_size): 7 | mask = np.zeros(patch_size[:2]) 8 | h, w = patch_size[:2] 9 | y_start = h // 2 - center_size[0] // 2 10 | x_start = w // 2 - center_size[1] // 2 11 | mask[y_start:y_start + center_size[0], x_start:x_start + center_size[1]] = 1 12 | 13 | return np.expand_dims(mask, axis=2) 14 | 15 | 16 | def create_rnd_patch(img, patch_size, seed): 17 | radius = img.shape[0] / 2 18 | 19 | np.random.seed(seed) 20 | angle = np.random.randint(0, 359) 21 | distance = np.random.randint(0, radius - np.sqrt(np.sum(np.power(patch_size, 2)))) 22 | x_start, y_start = (int(radius + distance * np.cos(angle)), int(radius - distance * np.sin(angle))) 23 | 24 | return img[y_start:y_start + patch_size[0], x_start:x_start + patch_size[1]] 25 | 26 | 27 | def prepare_patch(img, patch_size, mask, seed): 28 | patch = create_rnd_patch(img, patch_size, seed) 29 | patch = (1 - mask) * patch 30 | 31 | return patch.astype('float32') / 255.0 32 | 33 | 34 | class DataGenerator(Sequence): 35 | """Generates data for Keras""" 36 | def __init__(self, imgs, patch_size, center_size, batch_size=32, shuffle=True): 37 | 38 | self.imgs = imgs 39 | self.patch_size = patch_size 40 | self.center_size = center_size 41 | self.batch_size = batch_size 42 | self.shuffle = shuffle 43 | 44 | self.epoch = 0 45 | self.on_epoch_end() 46 | 47 | self.mask = create_center_mask(patch_size, center_size) 48 | 49 | def __len__(self): 50 | """Denotes the number of batches per epoch""" 51 | return int(np.floor(len(self.imgs) / self.batch_size)) 52 | 53 | def __getitem__(self, index): 54 | """Generate one batch of data""" 55 | # Generate indexes of the batch 56 | indexes = self.indexes[index * self.batch_size:(index + 1) * self.batch_size] 57 | 58 | # Generate data 59 | x = self.__data_generation(indexes, [self.imgs[k] for k in indexes]) 60 | 61 | return x, x 62 | 63 | def on_epoch_end(self): 64 | """Updates indexes after each epoch""" 65 | self.epoch += 1 66 | 67 | self.indexes = np.arange(len(self.imgs)) 68 | if self.shuffle == True: 69 | np.random.seed(self.epoch) 70 | np.random.shuffle(self.indexes) 71 | 72 | def __data_generation(self, indexes, batch): 73 | """Augments or/and pretransforms data""" 74 | patches = [prepare_patch(img, self.patch_size, self.mask, (self.epoch * idx) % (2**32 - 1)) 75 | for idx, img in zip(indexes, batch)] 76 | 77 | return np.array(patches) -------------------------------------------------------------------------------- /anomaly-detection-image-completion/scripts/model_helpers.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import keras.backend as K 3 | import numpy as np 4 | 5 | from keras.models import Model 6 | from keras.layers import Input, Conv2D, Lambda, UpSampling2D, ELU 7 | 8 | from data_helpers import * 9 | 10 | 11 | # Model parameters from the original paper: https://arxiv.org/pdf/1811.06861.pdf 12 | input_shape = (128, 128, 1) 13 | 14 | original_conv_layer_datas = [ 15 | {'kernel_size': 5, 'dilation_rate': 1, 'strides': 1, 'filters': 32}, 16 | {'kernel_size': 3, 'dilation_rate': 1, 'strides': 1, 'filters': 64}, 17 | {'kernel_size': 3, 'dilation_rate': 1, 'strides': 1, 'filters': 64}, 18 | {'kernel_size': 3, 'dilation_rate': 1, 'strides': 2, 'filters': 128}, 19 | {'kernel_size': 3, 'dilation_rate': 1, 'strides': 1, 'filters': 128}, 20 | {'kernel_size': 3, 'dilation_rate': 1, 'strides': 1, 'filters': 128}, 21 | {'kernel_size': 3, 'dilation_rate': 2, 'strides': 1, 'filters': 128}, 22 | {'kernel_size': 3, 'dilation_rate': 4, 'strides': 1, 'filters': 128}, 23 | {'kernel_size': 3, 'dilation_rate': 8, 'strides': 1, 'filters': 128}, 24 | {'kernel_size': 3, 'dilation_rate': 16, 'strides': 1, 'filters': 128}, 25 | {'kernel_size': 3, 'dilation_rate': 1, 'strides': 1, 'filters': 128}, 26 | {'kernel_size': 3, 'dilation_rate': 1, 'strides': 1, 'filters': 128}, 27 | {'kernel_size': 3, 'dilation_rate': 1, 'strides': 1, 'filters': 64}, 28 | {'kernel_size': 3, 'dilation_rate': 1, 'strides': 1, 'filters': 64}, 29 | {'kernel_size': 3, 'dilation_rate': 1, 'strides': 1, 'filters': 32}, 30 | {'kernel_size': 3, 'dilation_rate': 1, 'strides': 1, 'filters': 16}, 31 | {'kernel_size': 3, 'dilation_rate': 1, 'strides': 1, 'filters': 1} 32 | ] 33 | 34 | 35 | def conv_block(outputs, kernel, dilation, strides, filters, model_width): 36 | outputs = Lambda( 37 | lambda x: tf.pad(x, [[0, 0], [dilation, dilation], [dilation, dilation], [0, 0]], 'REFLECT'))(outputs) 38 | outputs = Conv2D(filters=filters * model_width, kernel_size=kernel, strides=strides, padding='valid', 39 | dilation_rate=dilation, activation='linear')(outputs) 40 | return ELU()(outputs) 41 | 42 | 43 | def create_anomaly_cnn(input_shape=input_shape, conv_layer_datas=None, model_width=1): 44 | """ 45 | Creates the CNN used for Anomaly Detection. 46 | :param input_shape: The shape of the inputed image 47 | :param conv_layer_datas: The layer parameters of the model 48 | :param model_width: Determines the rate to which factor the number of filter will be increased 49 | :return: Returns the model 50 | """ 51 | conv_layer_datas = conv_layer_datas if conv_layer_datas else original_conv_layer_datas 52 | 53 | assert len(input_shape) == 2 or input_shape[-1] == 1, 'Images must only have one channel (grayscale)!' 54 | inputs = Input(shape=(*input_shape[:2], 1)) 55 | outputs = inputs 56 | for i, data in enumerate(conv_layer_datas): 57 | outputs = conv_block(outputs, data['kernel_size'], data['dilation_rate'], data['strides'], data['filters'], 58 | model_width if i != len(conv_layer_datas) - 1 else 1) 59 | outputs = UpSampling2D(size=2)(outputs) if i == 11 else outputs 60 | outputs = Lambda(lambda x: tf.pad(x, [[0, 0], [1, 1], [1, 1], [0, 0]], 'REFLECT'))(outputs) 61 | outputs = Lambda(lambda x: K.clip(x, -1, 1), name='clip')(outputs) 62 | 63 | return Model(inputs=inputs, outputs=outputs) 64 | 65 | 66 | def l1_matrix_norm(M): 67 | return K.cast(K.max(K.sum(K.abs(M), axis=0)), 'float32') 68 | 69 | 70 | def reconstruction_loss(patch_size, mask=None, center_size=None, center_weight=0.9): 71 | assert mask is not None or center_size, 'You have to either specify the mask or the center_size' 72 | mask = mask if mask is not None else create_center_mask(patch_size, center_size[:2]) 73 | mask = mask.reshape(1, *mask.shape).astype('float32') 74 | mask_inv = 1 - mask 75 | 76 | def loss(y_true, y_pred): 77 | diff = y_true - y_pred 78 | center_part = mask * diff 79 | center_part_normed = l1_matrix_norm(center_part) 80 | 81 | surr_part = mask_inv * diff 82 | surr_part_normed = l1_matrix_norm(surr_part) 83 | 84 | num_pixels = np.prod(patch_size).astype('float32') 85 | 86 | numerator = center_weight * center_part_normed + (1 - center_weight) * surr_part_normed 87 | 88 | return numerator / num_pixels 89 | 90 | return loss -------------------------------------------------------------------------------- /anomaly-detection-image-completion/scripts/train_helpers.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import tensorflow as tf 3 | import io 4 | import warnings 5 | 6 | from glob import glob 7 | from os import remove 8 | from os.path import join 9 | from PIL import Image 10 | from keras.callbacks import Callback 11 | 12 | 13 | def transform_image(img): 14 | """ 15 | Convert a numpy representation image to Image protobuf 16 | """ 17 | height, width = img.shape[:2] 18 | num_channels = 1 if len(img.shape) < 3 else img.shape[-1] 19 | image = Image.fromarray(np.squeeze(img).astype(np.uint8)) 20 | 21 | output = io.BytesIO() 22 | image.save(output, format='PNG') 23 | image_string = output.getvalue() 24 | output.close() 25 | return tf.Summary.Image(height=height, 26 | width=width, 27 | colorspace=num_channels, 28 | encoded_image_string=image_string) 29 | 30 | 31 | class TensorBoardImages(Callback): 32 | """ 33 | Visualize images and their predictions at every `vis_every` epoch 34 | """ 35 | def __init__(self, logdir, imgss, vis_every=1): 36 | super(TensorBoardImages, self).__init__() 37 | self.logdir = logdir 38 | self.imgss = imgss 39 | self.vis_every = vis_every 40 | self.writer = tf.summary.FileWriter(logdir) 41 | 42 | def on_epoch_end(self, epoch, logs={}): 43 | if epoch % self.vis_every == 0: 44 | 45 | # Create tensorflow summaries for images 46 | for dataset, imgs in self.imgss.items(): 47 | for i, img in enumerate(imgs): 48 | orig = np.squeeze(img) 49 | pred = np.squeeze(self.model.predict(np.expand_dims(img, axis=0))) 50 | 51 | orig_summary = tf.Summary( 52 | value=[tf.Summary.Value(tag='Original_{}_{}'.format(dataset, i), image=transform_image(orig))]) 53 | pred_summary = tf.Summary( 54 | value=[tf.Summary.Value(tag='Predicted_{}_{}'.format(dataset, i), image=transform_image(pred))]) 55 | 56 | self.writer.add_summary(orig_summary, epoch) 57 | self.writer.add_summary(pred_summary, epoch) 58 | 59 | def on_train_end(self, logs={}): 60 | self.writer.close() 61 | 62 | 63 | class CustomModelCheckpoint(Callback): 64 | """ 65 | Save the last and best weights as well as the complete model according to the monitored metric 66 | """ 67 | def __init__(self, logdir, monitor='val_loss', verbose=0, save_weights_only=False, mode='auto', period=1): 68 | 69 | super(CustomModelCheckpoint, self).__init__() 70 | self.monitor = monitor 71 | self.verbose = verbose 72 | self.logdir = logdir 73 | self.weights_path = join( 74 | logdir, '{savetype}_{filetype}_epoch-{epoch:04d}_train_loss-{loss:.6f}_val_loss-{val_loss:.6f}.hdf5') 75 | self.save_weights_only = save_weights_only 76 | self.period = period 77 | self.epochs_since_last_save = 0 78 | 79 | if mode not in ['auto', 'min', 'max']: 80 | warnings.warn('ModelCheckpoint mode %s is unknown, ' 81 | 'fallback to auto mode.' % (mode), 82 | RuntimeWarning) 83 | mode = 'auto' 84 | 85 | if mode == 'min': 86 | self.monitor_op = np.less 87 | self.best = np.Inf 88 | elif mode == 'max': 89 | self.monitor_op = np.greater 90 | self.best = -np.Inf 91 | else: 92 | if 'acc' in self.monitor or self.monitor.startswith('fmeasure'): 93 | self.monitor_op = np.greater 94 | self.best = -np.Inf 95 | else: 96 | self.monitor_op = np.less 97 | self.best = np.Inf 98 | 99 | def on_epoch_end(self, epoch, logs=None): 100 | logs = logs or {} 101 | self.epochs_since_last_save += 1 102 | 103 | if self.epochs_since_last_save >= self.period: 104 | self.epochs_since_last_save = 0 105 | 106 | current = logs.get(self.monitor) 107 | if current is None: 108 | warnings.warn('Can save best model only with %s available, ' 109 | 'skipping.' % (self.monitor), RuntimeWarning) 110 | else: 111 | if self.monitor_op(current, self.best): 112 | if self.verbose > 0: 113 | print('\nEpoch %05d: %s improved from %0.5f to %0.5f' % ( 114 | epoch + 1, self.monitor, self.best, current)) 115 | self.best = current 116 | 117 | for ckpt_file in glob(join(self.logdir, 'best_weights*')): 118 | remove(ckpt_file) 119 | self.model.save_weights( 120 | self.weights_path.format(savetype='best', filetype='weights', epoch=epoch + 1, **logs)) 121 | 122 | if not self.save_weights_only: 123 | for ckpt_file in glob(join(self.logdir, 'best_model*')): 124 | remove(ckpt_file) 125 | self.model.save( 126 | self.weights_path.format(savetype='best', filetype='model', epoch=epoch + 1, **logs)) 127 | else: 128 | if self.verbose > 0: 129 | print('\nEpoch %05d: %s did not improve from %0.5f' % 130 | (epoch + 1, self.monitor, self.best)) 131 | 132 | 133 | for ckpt_file in glob(join(self.logdir, 'last_weights*')): 134 | remove(ckpt_file) 135 | self.model.save_weights( 136 | self.weights_path.format(savetype='last', filetype='weights', epoch=epoch + 1, **logs)) 137 | if not self.save_weights_only: 138 | for ckpt_file in glob(join(self.logdir, 'last_model*')): 139 | remove(ckpt_file) 140 | self.model.save(self.weights_path.format(savetype='last', filetype='model', epoch=epoch + 1, **logs)) -------------------------------------------------------------------------------- /anomalyDetection/resource.md: -------------------------------------------------------------------------------- 1 | ## resource 2 | https://github.com/yzhao062/anomaly-detection-resources 3 | 4 | https://github.com/yzhao062/pyod 5 | 6 | https://github.com/yzhao062/combo 7 | 8 | https://github.com/yzhao062/XGBOD 9 | 10 | https://github.com/yzhao062/SIML 11 | 12 | https://github.com/chen0040/keras-anomaly-detection 13 | 14 | https://github.com/vc1492a/PyNomaly 15 | 16 | https://github.com/Script-Boy/UnsupervisedAnomalyDetection 17 | 18 | 19 | ### detection 20 | https://github.com/chocolat-nya/Imbalanced_Image_Anomaly_Detection 21 | 22 | 23 | ### GAN 24 | https://github.com/houssamzenati/Efficient-GAN-Anomaly-Detection 25 | https://github.com/bianjiang1234567/CVAE-GAN-BASED-Anomaly-Detection 26 | https://github.com/donggong1/memae-anomaly-detection 27 | https://github.com/izikgo/AnomalyDetectionTransformations 28 | https://github.com/LeeDoYup/AnoGAN 29 | https://github.com/samet-akcay/ganomaly 30 | https://github.com/vantage-vision-vv/Anomaly-Detection-in-Surveillance-Videos 31 | https://github.com/otenim/AnomalyDetectionUsingAutoencoder 32 | https://github.com/trigrass2/wgan-gp-anomaly 33 | https://github.com/tkwoo/anogan-keras 34 | https://github.com/tSchlegl/f-AnoGAN 35 | https://github.com/WellenWoo/GANomaly-pytorch 36 | http://tongtianta.site/paper/33739 Skip-GANomaly: Skip Connected and Adversarially Trained Encoder-Decoder Anomaly Detection 37 | 38 | 39 | ### data mining algorithms 40 | https://github.com/Albertsr/Anomaly-Detection 41 | https://github.com/lishiyun19/AnomalyFilter 42 | https://github.com/a342058040/network_anomaly_detection 43 | https://github.com/kLabUM/rrcf 44 | https://github.com/tdda/pydatalondon2018ad 45 | https://github.com/JGuymont/vae-anomaly-detector 46 | https://github.com/caisr-hh/group-anomaly-detection 47 | https://github.com/shubhomoydas/ad_examples 48 | https://github.com/sapols/Satellite-Telemetry-Anomaly-Detection 49 | 50 | ### timeseries 51 | https://github.com/BBVA/timecop 52 | https://github.com/microsoft/anomalydetector 53 | https://github.com/LiDan456/MAD-GANs 54 | https://github.com/smallcowbaby/OmniAnomaly 55 | https://github.com/KDD-OpenSource/DeepADoTS 56 | https://github.com/Danyleb/Variational-Lstm-Autoencoder 57 | https://github.com/ZhouYuxuanYX/Unsupervised-Deep-Learning-Framework-for-Anomaly-Detection-in-Time-Series- 58 | https://github.com/LiDan456/GAN-AD 59 | https://github.com/Pengnengsong/AnomalyDetection 60 | https://github.com/arundo/adtk 61 | https://github.com/kentaroy47/AnomalyDetection.pytorch 62 | https://github.com/kentaroy47/easy-anomaly-detection-with-keras 63 | https://github.com/chickenbestlover/RNN-Time-series-Anomaly-Detection 64 | https://github.com/microsoft/TagAnomaly 65 | https://research.fb.com/wp-content/uploads/2018/11/AnoGen-Deep-Anomaly-Generator.pdf AnoGen: Deep Anomaly Generator 66 | 67 | ### GCN 68 | https://github.com/kaize0409/GCN_AnomalyDetection 69 | 70 | 71 | ### network 72 | https://github.com/Forcrush/Efficient-GAN-based-method-for-cyber-intrusion 73 | https://github.com/picsolab/DeepSphere 74 | 75 | 76 | ### VAE 77 | https://github.com/usersubsetscan/autoencoder_anomaly_subset 78 | https://github.com/Serhiy-Shekhovtsov/mars-anomalies-aae 79 | 80 | 81 | ### CV 82 | https://github.com/fjchange/object_centric_VAD 83 | https://github.com/ptirupat/AnomalyDetection_CVPR18 84 | https://github.com/ekosman/AnomalyDetectionCVPR2018-Pytorch 85 | https://github.com/WaqasSultani/AnomalyDetectionCVPR2018 86 | https://github.com/davenza/KDE-AMD 87 | https://github.com/microsoft/MLOps_VideoAnomalyDetection 88 | 89 | 90 | 91 | https://github.com/madara-tribe/images_anomaly_detection 92 | http://trafalbad.hatenadiary.jp/entry/2019/09/04/213627 93 | 94 | https://github.com/ivanbergonzani/adversarial-segmentation 95 | https://github.com/drsagitn/anomaly-detection-and-localization 96 | 97 | 98 | ### Unsupervised Anomaly Segmentation 99 | https://github.com/HuayueZhang/textured_surface_anomaly_detection 100 | https://github.com/msminhas93/CompactCNN 101 | https://github.com/oezguensi/anomaly-detection-image-completion 102 | https://github.com/bumuckl/AutoencodersForUnsupervisedAnomalyDetection Deep Autoeoncoding Models for Unsupervised Anomaly Detection in Brain MR Images 103 | http://tongtianta.site/paper/17042 Deep Autoencoding Models for Unsupervised Anomaly Segmentation in Brain MR Images 104 | 105 | 106 | https://github.com/xtarx/Unsupervised-Anomaly-Detection-with-Generative-Adversarial-Networks 107 | 108 | https://www.groundai.com/project/pnunet-anomaly-detection-using-positive-and-negative-noise-based-on-self-training-procedure 109 | 110 | 111 | https://github.com/h19920918/memae Memory-augmented Deep Autoencoder for Unsupervised Anomaly Detection 112 | https://github.com/bumuckl/AutoencodersForUnsupervisedAnomalyDetection 113 | https://github.com/khalilouardini/towards-practical-unsupervised-AD Towards Practical Unsupervised Anomaly Detection on Retinal images 114 | 115 | 116 | Generative Data Augmentation for Unsupervised Anomaly Detection with GAN 117 | 118 | 119 | AnoVAEGAN VAEGAN Context Encoders AE-GAN Spatial Variational Autoencoder Spatial Autoencoder 120 | 121 | 122 | https://github.com/xtarx/Unsupervised-Anomaly-Detection-with-Generative-Adversarial-Networks 123 | 124 | 125 | https://github.com/cheapthrillandwine/Improving_Unsupervised_Defect_Segmentation 126 | https://arxiv.org/abs/1807.02011 Improving Unsupervised Defect Segmentation by Applying Structural Similarity to Autoencoders 127 | 128 | 129 | 130 | https://github.com/yousuhang/Unsupervised-Lesion-Detection-via-Image-Restoration-with-a-Normative-Prior 131 | https://openreview.net/pdf?id=S1xg4W-leV 132 | 133 | 134 | 135 | https://github.com/aubreychen9012/cAAE Unsupervised Detection of Lesions in Brain MRI using constrained adversarial auto-encoders 136 | http://tongtianta.site/paper/12517 137 | https://github.com/aubreychen9012/LocalGaussianApproximation 138 | 139 | 140 | http://tongtianta.site/paper/59419 Iterative energy-based projection on a normal data manifold for anomaly localization 141 | http://tongtianta.site/paper/25792 Deep structured energy-based image inpainting 142 | 143 | 144 | ### AAE why convolutional AAE ???? 145 | https://github.com/Naresh1318/Adversarial_Autoencoder 146 | 147 | 148 | 149 | # for videos 150 | https://github.com/yiweilu3/CONV-VRNN-for-Anomaly-Detection Future Frame Prediction Using Convolutional VRNN for Anomaly Detection 151 | 152 | 153 | https://github.com/HeYuwei/GCN-Anomaly-Detection -------------------------------------------------------------------------------- /cAAE/README.md: -------------------------------------------------------------------------------- 1 | # cAAE 2 | 3 | code for Unsupervised Detection of Lesions in Brain MRI using constrained adversarial auto-encoders, [https://arxiv.org/abs/1806.04972](https://arxiv.org/abs/1806.04972) 4 | 5 | AAE model is implemented based on [this github repo](https://github.com/Naresh1318/Adversarial_Autoencoder) 6 | 7 | Required: python>=2.7, tensorlayer, tensorflow>=1.0, numpy 8 | 9 | train model: python main.py --model_name "cAAE" --z_dim "128" 10 | 11 | 12 | ### from 13 | https://github.com/aubreychen9012/cAAE Unsupervised Detection of Lesions in Brain MRI using constrained adversarial auto-encoders 14 | http://tongtianta.site/paper/12517 -------------------------------------------------------------------------------- /cAAE/funcs/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ForrestPi/Unsupervised-Defect-Segmentation/e366ac7c757bb1b45f38ebbc502dfee7ccb72398/cAAE/funcs/__init__.py -------------------------------------------------------------------------------- /cAAE/funcs/networks.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | def lrelu(x, alpha): 4 | return tf.nn.relu(x) - alpha * tf.nn.relu(-x) 5 | 6 | def dense(x, n1, n2, name): 7 | """ 8 | Used to create a dense layer. 9 | :param x: input tensor to the dense layer 10 | :param n1: no. of input neurons 11 | :param n2: no. of output neurons 12 | :param name: name of the entire dense layer.i.e, variable scope name. 13 | :return: tensor with shape [batch_size, n2] 14 | """ 15 | with tf.variable_scope(name, reuse=None): 16 | weights = tf.get_variable("weights", shape=[n1, n2], 17 | initializer=tf.random_normal_initializer(mean=0., stddev=0.01)) 18 | bias = tf.get_variable("bias", shape=[n2], initializer=tf.constant_initializer(0.0)) 19 | out = tf.add(tf.matmul(x, weights), bias, name='matmul') 20 | return out 21 | -------------------------------------------------------------------------------- /cAAE/funcs/preproc.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import skimage.measure 3 | import scipy 4 | from scipy.misc import imresize 5 | 6 | def generate_patch(img_3d, patch_size, step_size, threshold): 7 | img = np.asarray(img_3d) 8 | min_ = img.min() 9 | l_patch = [] 10 | _,x,y = img.shape 11 | for im in img: 12 | for xidx in np.arange(0, x-patch_size+step_size, step_size): 13 | for yidx in np.arange(0, y-patch_size+step_size, step_size): 14 | patch = im[xidx:xidx+patch_size, yidx:yidx+patch_size] 15 | if (patch==min_).sum() 1] 42 | img = img[idx] 43 | _, x, y = img.shape 44 | max_xy = max(x,y) 45 | 46 | a = (max_xy-x)/2 47 | b = (max_xy-y)/2 48 | 49 | if len(X_train_input)==0: 50 | print(img.shape) 51 | 52 | img = np.pad(img, ((0, 0),((a,a)), (b,b)), mode = 'edge') 53 | img = preproc.resize(img, 256 /233.0) 54 | 55 | if labels: 56 | z = np.genfromtxt(os.path.join(wd, str(i) + '/' + str(i) + "_label.txt")) 57 | assert len(z)==len(idx) 58 | X_train_target.extend(z) 59 | 60 | if ds_scale!=0: 61 | img = preproc.downsample_image(img[np.newaxis,:,:,:],ds_scale) 62 | X_train_input.extend(img) 63 | 64 | 65 | for j in subject_test: 66 | print(j) 67 | pathx=j 68 | img = nib.load(pathx).get_data() 69 | img = np.transpose(img, [2, 0, 1]) 70 | idx = [s for s in range(img.shape[0]) if len(set(img[s].flatten())) > 1] 71 | img = img[idx] 72 | img = np.pad(img, ((0, 0), (a,a), (b,b)), mode='edge') 73 | img = preproc.resize(img, 256 / 233.0) 74 | 75 | if ds_scale!=0: 76 | img = preproc.downsample_image(img[np.newaxis,:,:,:],ds_scale) 77 | X_dev_input.extend(img) 78 | if labels: 79 | z = np.genfromtxt(os.path.join(wd, str(j) + '/' + str(j) + "_label.txt")) 80 | print(len(z), img.shape) 81 | X_dev_target.extend(z) 82 | 83 | if not labels: 84 | return np.asarray(X_train_input), np.asarray(X_dev_input) 85 | else: 86 | return X_train_input, X_train_target, X_dev_input, X_dev_target -------------------------------------------------------------------------------- /cAAE/model.py: -------------------------------------------------------------------------------- 1 | from tensorlayer.layers import * 2 | from funcs.resblocks import ResBlock, ResBlockDown, ResBlockUp 3 | 4 | # The autoencoder network 5 | def encoder(x, z_dim, reuse=False, is_train=True): 6 | """ 7 | Encode part of the autoencoder. 8 | :param x: input to the autoencoder 9 | :param reuse: True -> Reuse the encoder variables, False -> Create or search of variables before creating 10 | :return: tensor which is the hidden latent variable of the autoencoder. 11 | """ 12 | input_dim = x.shape[-1] 13 | image_size = input_dim 14 | s2, s4, s8, s16 = int(image_size / 2), int(image_size / 4), int(image_size / 8), int(image_size / 16) 15 | gf_dim = 16 # Dimension of gen filters in first conv layer. [64] 16 | ft_size = 3 17 | with tf.variable_scope("Encoder", reuse=reuse): 18 | # x,y,z,_ = tf.shape(input_images) 19 | set_name_reuse(reuse) 20 | 21 | w_init = tf.truncated_normal_initializer(stddev=0.02) 22 | b_init = tf.constant_initializer(value=0.0) 23 | gamma_init = tf.random_normal_initializer(1., 0.01) 24 | 25 | inputs = InputLayer(x, name='e_inputs') 26 | conv1 = Conv2d(inputs, gf_dim, (ft_size, ft_size), act=tf.nn.leaky_relu(x, 0.2), 27 | padding='SAME',W_init=w_init, b_init=b_init, name="e_conv1") 28 | conv1 = BatchNormLayer(conv1, act=tf.nn.leaky_relu(x, 0.2), is_train=is_train, 29 | gamma_init=gamma_init, name='e_bn1') 30 | # image_size * image_size 31 | res1 = ResBlockDown(conv1.outputs, gf_dim, "res1", reuse, is_train) 32 | 33 | # s2*s2 34 | res2 = ResBlockDown(res1, gf_dim * 2, "res2", reuse, is_train) 35 | 36 | # s4*s4 37 | res3 = ResBlockDown(res2, gf_dim * 4, "res3", reuse, is_train) 38 | 39 | # s8*s8 40 | res4 = ResBlockDown(res3, gf_dim * 8, "res4", reuse, is_train) 41 | 42 | # s16*s16 43 | h_flat = tf.reshape(res4, shape=[-1, s16 * s16 * gf_dim * 16]) 44 | h_flat = InputLayer(h_flat, name='e_reshape') 45 | net_h = DenseLayer(h_flat, n_units=z_dim, act=tf.identity, name="e_dense_mean") 46 | return net_h.outputs 47 | 48 | 49 | def decoder(x, reuse=False, is_train=True): 50 | """ 51 | Decoder part of the autoencoder. 52 | :param x: input to the decoder 53 | :param reuse: True -> Reuse the decoder variables, False -> Create or search of variables before creating 54 | :return: tensor which should ideally be the input given to the encoder. 55 | """ 56 | input_dim = x.shape[-1] 57 | image_size = input_dim 58 | s2, s4, s8, s16 = int(image_size / 2), int(image_size / 4), int(image_size / 8), int(image_size / 16) 59 | gf_dim = 16 # Dimension of gen filters in first conv layer. [64] 60 | c_dim = 1 # n_color 3 61 | ft_size = 3 62 | batch_size = 16 # 64 63 | with tf.variable_scope("Decoder", reuse=reuse): 64 | set_name_reuse(reuse) 65 | w_init = tf.truncated_normal_initializer(stddev=0.02) 66 | b_init = tf.constant_initializer(value=0.0) 67 | inputs = InputLayer(x, name='g_inputs') 68 | 69 | # s16*s16 70 | z_develop = DenseLayer(inputs, s16 * s16 * gf_dim * 16, act=tf.nn.leaky_relu(x, 0.2), 71 | name='g_dense_z') 72 | z_develop = tf.reshape(z_develop.outputs, [-1, s16, s16, gf_dim * 16]) 73 | z_develop = InputLayer(z_develop, name='g_reshape') 74 | conv1 = Conv2d(z_develop, gf_dim * 8, (ft_size, ft_size), act=tf.nn.leaky_relu(x, 0.2), 75 | padding='SAME', W_init=w_init, b_init=b_init, name="g_conv1") 76 | 77 | # s16*s16 78 | res1 = ResBlockUp(conv1.outputs, s16, batch_size, gf_dim * 8, "gres1", reuse, is_train) 79 | 80 | # s8*s8 81 | res2 = ResBlockUp(res1, s8, batch_size, gf_dim * 4, "gres2", reuse, is_train) 82 | 83 | # s4*s4 84 | res3 = ResBlockUp(res2, s4, batch_size, gf_dim * 2, "gres3", reuse, is_train) 85 | 86 | # s2*s2 87 | res4 = ResBlockUp(res3, s2, batch_size, gf_dim, "gres4", reuse, is_train) 88 | 89 | # image_size*image_size 90 | res_inputs = InputLayer(res4, name='res_inputs') 91 | conv2 = Conv2d(res_inputs, c_dim, (ft_size, ft_size), act=None, padding='SAME', W_init=w_init, b_init=b_init, 92 | name="g_conv2") 93 | conv2_std = Conv2d(res_inputs, c_dim, (ft_size, ft_size), act=None, padding='SAME', W_init=w_init, 94 | b_init=b_init, 95 | name="g_conv2_std") 96 | return conv2.outputs, conv2_std.outputs 97 | 98 | 99 | def discriminator(x, reuse=False): 100 | """ 101 | Discriminator that is used to match the posterior distribution with a given prior distribution. 102 | :param x: tensor of shape [batch_size, z_dim] 103 | :param reuse: True -> Reuse the discriminator variables, 104 | False -> Create or search of variables before creating 105 | :return: tensor of shape [batch_size, 1] 106 | """ 107 | n_l1=200 108 | n_l2=400 109 | w_init = tf.random_normal_initializer(stddev=0.01) 110 | with tf.variable_scope("Discriminator", reuse=reuse): 111 | set_name_reuse(reuse) 112 | net_in = InputLayer(x, name='dc/in') 113 | net_h0 = DenseLayer(net_in, n_units=n_l1, 114 | W_init=w_init, 115 | act=tf.nn.leaky_relu(x, 0.2), name='dc/h0/lin') 116 | net_h1 = DenseLayer(net_h0, n_units=n_l2, 117 | W_init=w_init, 118 | act=tf.nn.leaky_relu(x, 0.2), name='dc/h1/lin') 119 | net_h2 = DenseLayer(net_h1, n_units=1, 120 | W_init=w_init, 121 | act=tf.identity, name='dc/h2/lin') 122 | logits = net_h2.outputs 123 | net_h2.outputs = tf.nn.sigmoid(net_h2.outputs) 124 | return net_h2.outputs, logits 125 | -------------------------------------------------------------------------------- /cAAE/normalize_data.py: -------------------------------------------------------------------------------- 1 | import nibabel as nib 2 | import glob 3 | import os 4 | import numpy as np 5 | import tensorlayer as tl 6 | 7 | ''' 8 | Before normalization, run N4 bias correction (https://www.ncbi.nlm.nih.gov/pubmed/20378467), 9 | then save the data under folder ./CamCAN_unbiased/CamCAN 10 | ''' 11 | 12 | modalities = ['T1w', 'T2w'] 13 | BraTS_modalities = ['T1w'] 14 | folders = ['HGG', 'LGG'] 15 | wd = './Data/CamCAN_unbiased/CamCAN' 16 | 17 | thumbnail_idx = [60, 70, 80, 90] 18 | 19 | for mod in modalities: 20 | wd_mod = os.path.join(wd, str(mod)) 21 | os.chdir(wd_mod) 22 | img_files = [i for i in glob.glob("*") if "_unbiased" in i] 23 | 24 | for img in img_files: 25 | print(img) 26 | img_data = nib.load(img) 27 | img_data = img_data.get_data() 28 | mask = img.split("_unbiased")[0] + "_brain_mask.nii.gz" 29 | mask_data = nib.load(mask).get_data() 30 | 31 | img_data = np.transpose(img_data, [2, 0, 1]) 32 | mask_data = np.transpose(mask_data, [2, 0, 1]) 33 | 34 | idx = [s for s in range(img_data.shape[0]) if mask_data[s].sum() > 1] 35 | img_data = img_data[idx, :, 17:215] 36 | mask_data = mask_data[idx, :, 17:215] 37 | 38 | img_data = np.pad(img_data, ((0, 0), (1, 2), (1, 1)), mode='edge') 39 | mask_data = np.pad(mask_data, ((0, 0), (1, 2), (1, 1)), mode='edge') 40 | img_data = np.rot90(img_data, 1, (2, 1)) 41 | mask_data = np.rot90(mask_data, 1, (2, 1)) 42 | 43 | ref_mean = np.mean(img_data[mask_data == 1]) 44 | ref_std = np.std(img_data[mask_data == 1]) 45 | 46 | normed_img = (img_data - ref_mean) / ref_std 47 | normed_img[normed_img == normed_img.min()] = -3.5 48 | 49 | x_nif = nib.Nifti1Image(normed_img, np.eye(4)) 50 | nib.save(x_nif, os.path.join(img.split("_unbiased")[0] + "_normalized_cropped_mask.nii.gz")) 51 | 52 | x_nif = nib.Nifti1Image(mask_data, np.eye(4)) 53 | nib.save(x_nif, os.path.join(img.split("_unbiased")[0] + "_mask_cropped_mask.nii.gz")) 54 | 55 | tl.visualize.save_images(normed_img[thumbnail_idx, :, :, np.newaxis], [2, 2], 56 | "/scratch_net/bmicdl01/Data/CamCAN_unbiased/preview/" + str(mod) 57 | + "/" + img.split("_unbiased")[0] + "_normed_img.png") 58 | print("---") -------------------------------------------------------------------------------- /config/RED_Net_2skips-grid.json: -------------------------------------------------------------------------------- 1 | { 2 | "model": { 3 | "name": "RED_Net_2skips", 4 | "type": "Encoder", 5 | "code_dim": 256, 6 | "img_channel": 1, 7 | "fp16": false 8 | }, 9 | "op": { 10 | "loss": "Multi_SSIM_loss", 11 | "window_size":[3, 7, 15], 12 | 13 | "start_epoch": 0, 14 | "max_epoch": 200, 15 | "snapshot": 10, 16 | "learning_rate": 1e-3, 17 | "decay_rate": 10, 18 | "epoch_steps": [90, 150, 200] 19 | }, 20 | "db": { 21 | "name": "chip", 22 | "data_dir": "D:/DataSet/chip_grid", 23 | "loader_threads": 4, 24 | "train_split": "train", 25 | "use_validation_set": true, 26 | "validation_split": "validation", 27 | "val_split": "test", 28 | "resize": [768, 768], 29 | "batch_size": 2 30 | }, 31 | "system": { 32 | "resume": false, 33 | "resume_path": "", 34 | "finetune": false, 35 | "finetune_path": "", 36 | "save_dir": "./weights/" 37 | } 38 | } 39 | -------------------------------------------------------------------------------- /config/RED_Net_2skips-mvtec.json: -------------------------------------------------------------------------------- 1 | { 2 | "model": { 3 | "name": "RED_Net_2skips", 4 | "type": "Encoder", 5 | "code_dim": 256, 6 | "img_channel": 3, 7 | "fp16": false 8 | }, 9 | "op": { 10 | "loss": "SSIM_loss", 11 | "window_size":11, 12 | 13 | "start_epoch": 0, 14 | "max_epoch": 200, 15 | "snapshot": 10, 16 | "learning_rate": 4e-3, 17 | "decay_rate": 10, 18 | "epoch_steps": [150, 200] 19 | }, 20 | "db": { 21 | "name": "mvtec", 22 | "data_dir": "D:/DataSet/mvtec_anomaly_detection", 23 | "loader_threads": 4, 24 | "train_split": "train", 25 | "use_validation_set": false, 26 | "validation_split": "validation", 27 | "val_split": "test", 28 | "resize": [256, 256], 29 | "batch_size": 16 30 | }, 31 | "system": { 32 | "resume": false, 33 | "resume_path": "", 34 | "finetune": false, 35 | "finetune_path": "", 36 | "save_dir": "./weights/" 37 | } 38 | } 39 | -------------------------------------------------------------------------------- /config/RED_Net_3skips-grid.json: -------------------------------------------------------------------------------- 1 | { 2 | "model": { 3 | "name": "RED_Net_3skips", 4 | "type": "Encoder", 5 | "code_dim": 256, 6 | "img_channel": 1, 7 | "fp16": false 8 | }, 9 | "op": { 10 | "loss": "Multi_SSIM_loss", 11 | "window_size":[3, 7, 15], 12 | 13 | "start_epoch": 0, 14 | "max_epoch": 100, 15 | "snapshot": 10, 16 | "learning_rate": 1e-3, 17 | "decay_rate": 10, 18 | "epoch_steps": [60, 100] 19 | }, 20 | "db": { 21 | "name": "chip", 22 | "data_dir": "D:/DataSet/chip_grid", 23 | "loader_threads": 4, 24 | "train_split": "train", 25 | "use_validation_set": true, 26 | "validation_split": "validation", 27 | "val_split": "test", 28 | "resize": [768, 768], 29 | "batch_size": 32 30 | }, 31 | "system": { 32 | "resume": false, 33 | "resume_path": "", 34 | "finetune": false, 35 | "finetune_path": "", 36 | "save_dir": "./weights/" 37 | } 38 | } 39 | -------------------------------------------------------------------------------- /config/RED_Net_3skips-mvtec.json: -------------------------------------------------------------------------------- 1 | { 2 | "model": { 3 | "name": "RED_Net_3skips", 4 | "type": "Encoder", 5 | "code_dim": 256, 6 | "img_channel": 3, 7 | "fp16": false 8 | }, 9 | "op": { 10 | "loss": "SSIM_loss", 11 | "window_size":11, 12 | 13 | "start_epoch": 0, 14 | "max_epoch": 200, 15 | "snapshot": 10, 16 | "learning_rate": 4e-3, 17 | "decay_rate": 10, 18 | "epoch_steps": [150, 200] 19 | }, 20 | "db": { 21 | "name": "mvtec", 22 | "data_dir": "D:/DataSet/mvtec_anomaly_detection", 23 | "loader_threads": 4, 24 | "train_split": "train", 25 | "use_validation_set": true, 26 | "validation_split": "validation", 27 | "val_split": "test", 28 | "resize": [256, 256], 29 | "batch_size": 16 30 | }, 31 | "system": { 32 | "resume": false, 33 | "resume_path": "", 34 | "finetune": false, 35 | "finetune_path": "", 36 | "save_dir": "./weights/" 37 | } 38 | } 39 | -------------------------------------------------------------------------------- /config/RED_Net_4skips-mvtec.json: -------------------------------------------------------------------------------- 1 | { 2 | "model": { 3 | "name": "RED_Net_4skips", 4 | "type": "Encoder", 5 | "code_dim": 256, 6 | "img_channel": 3, 7 | "fp16": false 8 | }, 9 | "op": { 10 | "loss": "SSIM_loss", 11 | "window_size":11, 12 | 13 | "start_epoch": 0, 14 | "max_epoch": 200, 15 | "snapshot": 10, 16 | "learning_rate": 4e-3, 17 | "decay_rate": 10, 18 | "epoch_steps": [150, 200] 19 | }, 20 | "db": { 21 | "name": "mvtec", 22 | "data_dir": "./mvtec_anomaly_detection", 23 | "loader_threads": 4, 24 | "train_split": "train", 25 | "use_validation_set": true, 26 | "validation_split": "validation", 27 | "val_split": "test", 28 | "resize": [256, 256], 29 | "batch_size": 16 30 | }, 31 | "system": { 32 | "resume": false, 33 | "resume_path": "", 34 | "finetune": false, 35 | "finetune_path": "", 36 | "save_dir": "./weights/" 37 | } 38 | } 39 | -------------------------------------------------------------------------------- /config/SRGAN-mvtec.json: -------------------------------------------------------------------------------- 1 | { 2 | "model": { 3 | "name": "SRGAN", 4 | "type": "GAN", 5 | "code_dim": 256, 6 | "img_channel": 3, 7 | "upscale_factor":4, 8 | "fp16": false 9 | }, 10 | "op": { 11 | "loss": "SRGAN_loss", 12 | "vgg16_weight_path": "./weights/vgg16_reducedfc.pth", 13 | "start_epoch": 0, 14 | "max_epoch": 120, 15 | "snapshot": 10, 16 | "learning_rate": 5e-4, 17 | "decay_rate": 10, 18 | "epoch_steps": [60, 120] 19 | }, 20 | "db": { 21 | "name": "mvtec", 22 | "data_dir": "D:/DataSet/mvtec_anomaly_detection", 23 | "loader_threads": 4, 24 | "train_split": "train", 25 | "use_validation_set": true, 26 | "validation_split": "validation", 27 | "val_split": "test", 28 | "resize": [256, 256], 29 | "batch_size": 2 30 | }, 31 | "system": { 32 | "resume": false, 33 | "resume_path": "", 34 | "finetune": false, 35 | "finetune_path": "", 36 | "save_dir": "./weights/" 37 | } 38 | } 39 | -------------------------------------------------------------------------------- /config/SSIM-mvtec.json: -------------------------------------------------------------------------------- 1 | { 2 | "model": { 3 | "name": "SSIM_Net", 4 | "type": "Encoder", 5 | "code_dim": 256, 6 | "img_channel": 3, 7 | "fp16": false 8 | }, 9 | "op": { 10 | "loss": "Multi_SSIM_loss", 11 | "window_size":[3, 7, 15], 12 | 13 | "start_epoch": 0, 14 | "max_epoch": 120, 15 | "snapshot": 10, 16 | "learning_rate": 1e-3, 17 | "decay_rate": 10, 18 | "epoch_steps": [50, 120] 19 | }, 20 | "db": { 21 | "name": "mvtec", 22 | "data_dir": "D:/DataSet/mvtec_anomaly_detection", 23 | "loader_threads": 4, 24 | "train_split": "train", 25 | "use_validation_set": true, 26 | "validation_split": "validation", 27 | "val_split": "test", 28 | "resize": [256, 256], 29 | "batch_size": 2 30 | }, 31 | "system": { 32 | "resume": false, 33 | "resume_path": "", 34 | "finetune": false, 35 | "finetune_path": "", 36 | "save_dir": "./weights/" 37 | } 38 | } 39 | -------------------------------------------------------------------------------- /config/SSIM_lite-chip.json: -------------------------------------------------------------------------------- 1 | { 2 | "model": { 3 | "name": "SSIM_Net_lite", 4 | "type": "Encoder", 5 | "code_dim": 128, 6 | "img_channel": 1, 7 | "fp16": false 8 | }, 9 | "op": { 10 | "loss": "SSIM_loss", 11 | "window_size":11, 12 | 13 | "start_epoch": 0, 14 | "max_epoch": 100, 15 | "snapshot": 10, 16 | "learning_rate": 5e-4, 17 | "decay_rate": 10, 18 | "epoch_steps": [95, 100] 19 | }, 20 | "db": { 21 | "name": "chip", 22 | "data_dir": "D:/DataSet/chip_cell", 23 | "loader_threads": 4, 24 | "train_split": "train", 25 | "val_split": "test", 26 | "resize": [256, 256], 27 | "batch_size": 16 28 | }, 29 | "system": { 30 | "resume": false, 31 | "resume_path": "", 32 | "finetune": false, 33 | "finetune_path": "", 34 | "save_dir": "./weights/" 35 | } 36 | } 37 | -------------------------------------------------------------------------------- /config/VAE.json: -------------------------------------------------------------------------------- 1 | { 2 | "model": { 3 | "name": "VAE_Net0", 4 | "type": "Encoder", 5 | "code_dim": 256, 6 | "fp16": false 7 | }, 8 | "op": { 9 | "loss": "VAE_loss", 10 | "window_size":11, 11 | 12 | "start_epoch": 0, 13 | "max_epoch": 200, 14 | "snapshot": 4, 15 | "learning_rate": 1e-4, 16 | "decay_rate": 10, 17 | "epoch_steps": [150, 200] 18 | }, 19 | "db": { 20 | "name": "mvtec", 21 | "data_dir": "/home/yangxusheng/Desktop/baimingliang/torch-minist/mvtec_anomaly_detection", 22 | "loader_threads": 4, 23 | "train_split": "train", 24 | "use_validation_set": true, 25 | "validation_split": "validation", 26 | "val_split": "test", 27 | "resize": [256, 256], 28 | "crop_size": [128, 128], 29 | "batch_size": 16 30 | }, 31 | "system": { 32 | "resume": false, 33 | "resume_path": "", 34 | "finetune": false, 35 | "finetune_path": "", 36 | "save_dir": "./weights/" 37 | } 38 | } -------------------------------------------------------------------------------- /config/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ForrestPi/Unsupervised-Defect-Segmentation/e366ac7c757bb1b45f38ebbc502dfee7ccb72398/config/__init__.py -------------------------------------------------------------------------------- /db/__init__.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import numpy as np 3 | import torch 4 | from .mvtec import MVTEC, MVTEC_with_val 5 | from .mvtec import Preproc as MVTEC_pre 6 | from .chip import CHIP 7 | from .chip import Preproc as CHIP_pre 8 | 9 | 10 | def training_collate(batch): 11 | """Custom collate fn for dealing with batches of images. 12 | 13 | Arguments: 14 | batch: (tuple) A tuple of tensor images 15 | 16 | Return: 17 | (tensor) batch of images stacked on their 0 dim 18 | """ 19 | # imgs = list() 20 | # for img in batch: 21 | # _c, _h, _w = img.shape 22 | # imgs.append(img.view(1, _c, _h, _w)) 23 | 24 | return torch.stack(batch, 0) 25 | 26 | 27 | class Transform(object): 28 | def __init__(self, resize): 29 | self.resize = resize 30 | 31 | def __call__(self, image): 32 | image = cv2.resize(image, self.resize) 33 | ori_img = image.copy() 34 | image = image.astype(np.float32) / 255. 35 | if len(image.shape) == 3: 36 | image = image.transpose((2, 0, 1)) 37 | image = torch.from_numpy(image) 38 | else: 39 | image = torch.from_numpy(image) 40 | image = image.unsqueeze(0) 41 | 42 | 43 | return ori_img, image.unsqueeze(0) 44 | 45 | -------------------------------------------------------------------------------- /db/augment.py: -------------------------------------------------------------------------------- 1 | """Augmentation function 2 | 3 | author: Haixin wang 4 | e-mail: haixinwa@gmail.com 5 | """ 6 | 7 | import cv2 8 | import random 9 | import numpy as np 10 | 11 | 12 | def mirror(image): 13 | image_m = image[:, ::-1] 14 | 15 | return image_m 16 | 17 | 18 | def flip(image): 19 | image_f = image[::-1, :] 20 | 21 | return image_f 22 | 23 | 24 | def rotation(image, range): 25 | _h, _w = image.shape[0: 2] 26 | center = (_w // 2, _h // 2) 27 | rot = random.uniform(range[0], range[1]) 28 | M = cv2.getRotationMatrix2D(center, rot, 1) 29 | image_r = cv2.warpAffine(image, M, (_w, _h), borderMode=cv2.BORDER_REPLICATE) 30 | 31 | return image_r 32 | 33 | 34 | def shift(image, dis): 35 | _h, _w = image.shape[0:2] 36 | y_s = random.uniform(dis[0], dis[1]) 37 | x_s = random.uniform(dis[0], dis[1]) 38 | M = np.float32([[1, 0, x_s], [0, 1, y_s]]) 39 | image_s = cv2.warpAffine(image, M, (_w, _h), borderMode=cv2.BORDER_REPLICATE) 40 | 41 | return image_s 42 | 43 | 44 | def lighting_adjust(image, k, b): 45 | slope = random.uniform(k[0], k[1]) 46 | bias = random.uniform(b[0], b[1]) 47 | image = image * slope + bias 48 | image = np.clip(image, 0, 255) 49 | 50 | return image.astype(np.uint8) 51 | 52 | 53 | def normalize_(image, mean, std): 54 | image -= mean 55 | image /= std 56 | 57 | 58 | def crop(image, crop_size): 59 | height, width, _ = image.shape 60 | x_offset = random.randint(0, width - crop_size[0]) 61 | y_offset = random.randint(0, height - crop_size[1]) 62 | 63 | return image[y_offset: y_offset+crop_size[1], x_offset: x_offset+crop_size[0]] 64 | -------------------------------------------------------------------------------- /db/chip.py: -------------------------------------------------------------------------------- 1 | """Data set tool of MVTEC 2 | 3 | author: Haixin wang 4 | e-mail: haixinwa@gmail.com 5 | """ 6 | import os 7 | import re 8 | import cv2 9 | import random 10 | import torch 11 | import torch.utils.data as data 12 | from collections import OrderedDict 13 | from .augment import * 14 | from .eval_func import * 15 | 16 | 17 | class Preproc(object): 18 | """Pre-procession of input image includes resize, crop & data augmentation 19 | 20 | Arguments: 21 | resize: tup(int width, int height): resize shape 22 | crop: tup(int width, int height): crop shape 23 | """ 24 | def __init__(self, resize): 25 | self.resize = resize 26 | 27 | def __call__(self, image): 28 | image = cv2.resize(image, self.resize) 29 | # random transformation 30 | p = random.uniform(0, 1) 31 | if (p > 0.2) and (p <= 0.4): 32 | image = mirror(image) 33 | elif (p > 0.4) and (p <= 0.6): 34 | image = flip(image) 35 | # elif (p > 0.6) and (p <= 0.8): 36 | # image = shift(image, (-12, 12)) 37 | # else: 38 | # image = rotation(image, (-10, 10)) 39 | 40 | # light adjustment 41 | p = random.uniform(0, 1) 42 | if p > 0.5: 43 | image = lighting_adjust(image, k=(0.8, 0.95), b=(-10, 10)) 44 | 45 | # image normal 46 | image = image.astype(np.float32) / 255. 47 | # normalize_(tile, self.mean, self.std) 48 | image = torch.from_numpy(image) 49 | 50 | return image.unsqueeze(0) 51 | 52 | 53 | class CHIP(data.Dataset): 54 | """A tiny data set for chip cell 55 | 56 | Arguments: 57 | root (string): root directory to root folder. 58 | set (string): image set to use ('train', or 'test') 59 | preproc(callable, optional): pre-procession on the input image 60 | """ 61 | 62 | def __init__(self, root, set, preproc=None): 63 | self.root = root 64 | self.preproc = preproc 65 | self.set = set 66 | 67 | if set == 'train': 68 | self.ids = list() 69 | set_path = os.path.join(self.root, set) 70 | for img in os.listdir(set_path): 71 | item_path = os.path.join(set_path, img) 72 | self.ids.append(item_path) 73 | elif set == 'test': 74 | self.test_len = 0 75 | self.test_dict = OrderedDict() 76 | set_path = os.path.join(self.root, set) 77 | for type in os.listdir(set_path): 78 | type_dir = os.path.join(set_path, type) 79 | if os.path.isfile(type_dir): 80 | continue 81 | ids = list() 82 | for img in os.listdir(type_dir): 83 | if re.search('.png', img) is None: 84 | continue 85 | ids.append(os.path.join(type_dir, img)) 86 | self.test_len += 1 87 | self.test_dict[type] = ids 88 | else: 89 | raise Exception("Invalid set name") 90 | 91 | def __getitem__(self, index): 92 | """Returns training image 93 | """ 94 | img_path = self.ids[index] 95 | img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) 96 | img = self.preproc(img) 97 | 98 | return img 99 | 100 | def __len__(self): 101 | if self.set == 'train': 102 | return len(self.ids) 103 | else: 104 | return self.test_len -------------------------------------------------------------------------------- /db/eval_func.py: -------------------------------------------------------------------------------- 1 | """Evaluation function 2 | 3 | author: Haixin wang 4 | e-mail: haixinwa@gmail.com 5 | """ 6 | import cv2 7 | import numpy as np 8 | import matplotlib.pyplot as plt 9 | 10 | 11 | def cal_iou(mask, gt): 12 | mask_defect = (mask == 255) 13 | gt_defect = (gt == 255) 14 | overlap = (mask_defect == gt_defect) & (mask == 255) & (gt == 255) 15 | mask_defect_sum = float(mask_defect.sum()) 16 | gt_defect_sum = float(gt_defect.sum()) 17 | overlap_sum = float(overlap.sum()) 18 | iou = overlap_sum / (mask_defect_sum + gt_defect_sum - overlap_sum) 19 | 20 | return iou 21 | 22 | 23 | def cal_pixel_accuracy(mask, gt): 24 | _h, _w = gt.shape 25 | positive = (mask == gt) 26 | positive_sum = float(positive.sum()) 27 | pixel_acc = positive_sum / (_h * _w) 28 | 29 | return pixel_acc 30 | 31 | 32 | def cal_TPR(s_map,gt,threshold): 33 | # True Postive Rate 34 | gt_defect = (gt == 255) 35 | overlap = (s_map= threshold] = 0 30 | return mask 31 | 32 | -------------------------------------------------------------------------------- /model/segmentation/ssim_seg_cuda.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding:utf-8 -*- 3 | # author: forrestpi 4 | # data: 2019.09.04 5 | 6 | import torch 7 | import torch.nn.functional as F 8 | from model.loss.SSIM_loss import create_window 9 | 10 | def ssim_seg(img1, img2, window_size=11,threshold=128): 11 | (_, channel_1, _, _) = img1.size() 12 | (_, channel_2, _, _) = img1.size() 13 | BGR = torch.Tensor([[[[0.114,0.587,0.299]]]]).reshape((1,-1,1,1)) 14 | if channel_1 == 3: 15 | img1 = torch.sum(img1 * BGR, dim=1,keepdim=True) 16 | 17 | if channel_2 == 3: 18 | img2 = img2 * BGR 19 | img2 = torch.sum(img2 * BGR, dim=1,keepdim=True) 20 | 21 | 22 | channel = 1 23 | 24 | window = create_window(window_size, channel) 25 | 26 | if img1.is_cuda: 27 | window = window.cuda(img1.get_device()) 28 | # window = window.type_as(img1) 29 | mu1 = F.conv2d(img1.type_as(window), window, padding=window_size // 2, groups=channel) 30 | mu2 = F.conv2d(img2.type_as(window), window, padding=window_size // 2, groups=channel) 31 | 32 | mu1_sq = mu1.pow(2) 33 | mu2_sq = mu2.pow(2) 34 | mu1_mu2 = mu1 * mu2 35 | 36 | sigma1_sq = F.conv2d((img1 * img1).type_as(window), window, padding=window_size // 2, groups=channel) - mu1_sq 37 | sigma2_sq = F.conv2d((img2 * img2).type_as(window), window, padding=window_size // 2, groups=channel) - mu2_sq 38 | sigma12 = F.conv2d((img1 * img2).type_as(window), window, padding=window_size // 2, groups=channel) - mu1_mu2 39 | 40 | C1 = 0.01 ** 2 41 | C2 = 0.03 ** 2 42 | 43 | ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) 44 | mask = torch.randn_like(ssim_map, dtype=torch.float) 45 | 46 | #mask 47 | mask[ssim_map >= threshold/255] = 0 48 | mask[ssim_map < threshold/255] = 255 49 | 50 | return mask 51 | 52 | 53 | -------------------------------------------------------------------------------- /model/trainer.py: -------------------------------------------------------------------------------- 1 | """training container 2 | author: forrestpi 3 | e-mail: forrrest_zhu@foxmail.com 4 | """ 5 | 6 | 7 | import torch 8 | import torch.nn as nn 9 | 10 | 11 | class Network(nn.Module): 12 | def __init__(self, model, loss): 13 | super(Network, self).__init__() 14 | self.model = model 15 | self.loss = loss 16 | 17 | def forward(self, images): 18 | preds = self.model(images) 19 | loss = self.loss(preds, images) 20 | 21 | return loss 22 | 23 | 24 | class Trainer(): 25 | def __init__(self, net, loss, loss_name, optimizer, ngpu): 26 | self.net = net 27 | self.loss = loss 28 | self.loss_name = loss_name 29 | self.loss_value = None 30 | self.optimizer = optimizer 31 | self.network = torch.nn.DataParallel(Network(self.net, self.loss), device_ids=list(range(ngpu))) 32 | self.network.train() 33 | self.network.cuda() 34 | torch.backends.cudnn.benchmark = True 35 | 36 | def save_params(self, save_path): 37 | print("saving model to {}".format(save_path)) 38 | with open(save_path, "wb") as f: 39 | params = self.net.state_dict() 40 | torch.save(params, f) 41 | 42 | def load_params(self, path): 43 | from collections import OrderedDict 44 | new_state_dict = OrderedDict() 45 | w_dict = torch.load(path) 46 | for k, v in w_dict.items(): 47 | head = k[:7] 48 | if head == 'module.': 49 | name = k[7:] # remove `module.` 50 | else: 51 | name = k 52 | new_state_dict[name] = v 53 | self.net.load_state_dict(new_state_dict) 54 | 55 | def set_lr(self, lr): 56 | # print("setting learning rate to: {}".format(lr)) 57 | for param_group in self.optimizer.param_groups: 58 | param_group["lr"] = lr 59 | 60 | def train(self, input_tensor): 61 | if self.loss_name == 'SSIM_loss' or self.loss_name == 'VAE_loss': 62 | self.optimizer.zero_grad() 63 | loss = self.network(input_tensor) 64 | loss = loss.mean() 65 | loss.backward() 66 | self.optimizer.step() 67 | self.loss_value = loss.item() 68 | 69 | elif self.loss_name == 'Multi_SSIM_loss': 70 | self.loss_value = list() 71 | total_loss = list() 72 | self.optimizer.zero_grad() 73 | loss_multi = self.network(input_tensor) 74 | for loss in loss_multi: 75 | loss = loss.mean() 76 | total_loss.append(loss) 77 | self.loss_value.append(loss.item()) 78 | total_loss = torch.stack(total_loss, 0).sum() 79 | total_loss.backward() 80 | self.optimizer.step() 81 | 82 | else: 83 | raise Exception('Wrong loss name') 84 | 85 | def get_loss_message(self): 86 | if self.loss_name == 'SSIM_loss' or self.loss_name == 'VAE_loss': 87 | mes = 'ssim loss:{:.4f};'.format(self.loss_value) 88 | 89 | elif self.loss_name == 'Multi_SSIM_loss': 90 | mes = '' 91 | for k, loss in enumerate(self.loss_value): 92 | mes += 'size{:d} ssim loss:{:.4f}; '.format(k, loss) 93 | else: 94 | raise Exception('Wrong loss name') 95 | 96 | return mes -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | import os 2 | import cv2 3 | import json 4 | import argparse 5 | import numpy as np 6 | from db import Transform 7 | from model.rebuilder import Rebuilder 8 | from model.segmentation import ssim_seg, seg_mask 9 | from tools import Timer 10 | from factory import * 11 | from db.eval_func import cal_good_index 12 | 13 | 14 | def parse_args(): 15 | parser = argparse.ArgumentParser(description='Object detection base on anchor.') 16 | parser.add_argument('--cfg', help="Path of config file", type=str, required=True) 17 | parser.add_argument('--model_path', help="Path of model", type=str,required=True) 18 | parser.add_argument('--gpu_id', help="ID of GPU", type=int, default=0) 19 | parser.add_argument('--res_dir', help="Directory path of result", type=str, default='./eval_result') 20 | parser.add_argument('--retest', default=False, type=bool) 21 | 22 | return parser.parse_args() 23 | 24 | def val_mvtec(val_set, rebuilder, transform): 25 | threshold_seg_dict = dict() 26 | for item in val_set.val_dict: 27 | item_list = list() 28 | item_list = val_set.val_dict[item] 29 | good_count = 0 30 | for threshold_temp in range(0, 256): 31 | for path in item_list: 32 | image = cv2.imread(path, cv2.IMREAD_COLOR) 33 | ori_h, ori_w, _ = image.shape 34 | ori_img, input_tensor = transform(image) 35 | out = rebuilder.inference(input_tensor) 36 | re_img = out.transpose((1, 2, 0)) 37 | s_map = ssim_seg(ori_img, re_img) 38 | s_map = cv2.resize(s_map, (ori_w, ori_h)) 39 | mask = seg_mask(s_map, threshold_temp) 40 | good_count += cal_good_index(mask, 400) 41 | if good_count >= int(0.99*len(item_list)): 42 | threshold_seg_dict[item] = threshold_temp 43 | break 44 | print('validation: Item:{} finishes'.format(item)) 45 | return threshold_seg_dict 46 | 47 | 48 | def test_mvtec(test_set, rebuilder, transform, save_dir, threshold_seg_dict, val_index): 49 | _t = Timer() 50 | cost_time = list() 51 | threshold_dict = dict() 52 | if not os.path.exists(os.path.join(save_dir, 'ROC_curve')): 53 | os.mkdir(os.path.join(save_dir, 'ROC_curve')) 54 | for item in test_set.test_dict: 55 | threshold_list = list() 56 | item_dict = test_set.test_dict[item] 57 | 58 | if not os.path.exists(os.path.join(save_dir, item)): 59 | os.mkdir(os.path.join(save_dir, item)) 60 | os.mkdir(os.path.join(save_dir, item, 'ori')) 61 | os.mkdir(os.path.join(save_dir, item, 'gen')) 62 | os.mkdir(os.path.join(save_dir, item, 'mask')) 63 | for type in item_dict: 64 | if not os.path.exists(os.path.join(save_dir, item, 'ori', type)): 65 | os.mkdir(os.path.join(save_dir, item, 'ori', type)) 66 | if not os.path.exists(os.path.join(save_dir, item, 'gen', type)): 67 | os.mkdir(os.path.join(save_dir, item, 'gen', type)) 68 | if not os.path.exists(os.path.join(save_dir, item, 'mask', type)): 69 | os.mkdir(os.path.join(save_dir, item, 'mask', type)) 70 | _time = list() 71 | img_list = item_dict[type] 72 | for path in img_list: 73 | image = cv2.imread(path, cv2.IMREAD_COLOR) 74 | ori_h, ori_w, _ = image.shape 75 | _t.tic() 76 | ori_img, input_tensor = transform(image) 77 | out = rebuilder.inference(input_tensor) 78 | re_img = out.transpose((1, 2, 0)) 79 | s_map = ssim_seg(ori_img, re_img, win_size=11, gaussian_weights=True) 80 | s_map = cv2.resize(s_map, (ori_w, ori_h)) 81 | if val_index == 1: 82 | mask = seg_mask(s_map, threshold=threshold_seg_dict[item]) 83 | elif val_index == 0: 84 | mask = seg_mask(s_map, threshold=threshold_seg_dict) 85 | else: 86 | raise Exception("Invalid val_index") 87 | 88 | inference_time = _t.toc() 89 | img_id = path.split('.')[0][-3:] 90 | cv2.imwrite(os.path.join(save_dir, item, 'ori', type, '{}.png'.format(img_id)), ori_img) 91 | cv2.imwrite(os.path.join(save_dir, item, 'gen', type, '{}.png'.format(img_id)), re_img) 92 | cv2.imwrite(os.path.join(save_dir, item, 'mask', type, '{}.png'.format(img_id)), mask) 93 | _time.append(inference_time) 94 | 95 | if type != 'good': 96 | threshold_list.append(s_map) 97 | else: 98 | pass 99 | 100 | cost_time += _time 101 | mean_time = np.array(_time).mean() 102 | print('Evaluate: Item:{}; Type:{}; Mean time:{:.1f}ms'.format(item, type, mean_time*1000)) 103 | _t.clear() 104 | threshold_dict[item] = threshold_list 105 | # calculate mean time 106 | cost_time = np.array(cost_time) 107 | cost_time = np.sort(cost_time) 108 | num = cost_time.shape[0] 109 | num90 = int(num*0.9) 110 | cost_time = cost_time[0:num90] 111 | mean_time = np.mean(cost_time) 112 | print('Mean_time: {:.1f}ms'.format(mean_time*1000)) 113 | 114 | # evaluate results 115 | print('Evaluating...') 116 | test_set.eval(save_dir,threshold_dict) 117 | 118 | 119 | def test_chip(test_set, rebuilder, transform, save_dir): 120 | _t = Timer() 121 | cost_time = list() 122 | for type in test_set.test_dict: 123 | img_list = test_set.test_dict[type] 124 | if not os.path.exists(os.path.join(save_dir, type)): 125 | os.mkdir(os.path.join(save_dir, type)) 126 | for k, path in enumerate(img_list): 127 | image = cv2.imread(path, cv2.IMREAD_GRAYSCALE) 128 | _t.tic() 129 | ori_img, input_tensor = transform(image) 130 | out = rebuilder.inference(input_tensor) 131 | re_img = out[0] 132 | s_map = ssim_seg(ori_img, re_img, win_size=11, gaussian_weights=True) 133 | mask = seg_mask(s_map, threshold=32) 134 | inference_time = _t.toc() 135 | cat_img = np.concatenate((ori_img, re_img, mask), axis=1) 136 | cv2.imwrite(os.path.join(save_dir, type, '{:d}.png'.format(k)), cat_img) 137 | cost_time.append(inference_time) 138 | if (k+1) % 20 == 0: 139 | print('{}th image, cost time: {:.1f}'.format(k+1, inference_time*1000)) 140 | _t.clear() 141 | # calculate mean time 142 | cost_time = np.array(cost_time) 143 | cost_time = np.sort(cost_time) 144 | num = cost_time.shape[0] 145 | num90 = int(num*0.9) 146 | cost_time = cost_time[0:num90] 147 | mean_time = np.mean(cost_time) 148 | print('Mean_time: {:.1f}ms'.format(mean_time*1000)) 149 | 150 | 151 | if __name__ == '__main__': 152 | args = parse_args() 153 | 154 | # load config file 155 | cfg_file = os.path.join('./config', args.cfg + '.json') 156 | with open(cfg_file, "r") as f: 157 | configs = json.load(f) 158 | 159 | if not os.path.exists(args.res_dir): 160 | os.mkdir(args.res_dir) 161 | 162 | 163 | # load data set 164 | test_set = load_data_set_from_factory(configs, 'test') 165 | print('Data set: {} has been loaded'.format(configs['db']['name'])) 166 | 167 | # retest 168 | if args.retest is True: 169 | print('Evaluating...') 170 | test_set.eval(args.res_dir) 171 | exit(0) 172 | 173 | # init and load Rebuilder 174 | # load model 175 | transform = Transform(tuple(configs['db']['resize'])) 176 | net = load_test_model_from_factory(configs) 177 | rebuilder = Rebuilder(net, gpu_id=args.gpu_id) 178 | rebuilder.load_params(args.model_path) 179 | print('Model: {} has been loaded'.format(configs['model']['name'])) 180 | 181 | threshold_seg_dict = {} 182 | val_index = 0 183 | if configs['db']['name'] == 'mvtec': 184 | if configs['db']['use_validation_set'] is True: 185 | # load validation set 186 | val_index = 1 187 | val_set = load_data_set_from_factory(configs, 'validation') 188 | print('Data set: {} has been loaded'.format(configs['db']['name'])) 189 | # validation for threshold selection 190 | print('Start Validation... ') 191 | threshold_seg_dict = val_mvtec(val_set, rebuilder, transform) 192 | elif configs['db']['use_validation_set'] is False: 193 | val_index = 0 194 | else: 195 | raise Exception("Invalid input") 196 | elif configs['db']['name'] == 'chip': 197 | pass 198 | else: 199 | raise Exception("Invalid set name") 200 | 201 | # test each image 202 | print('Start Testing... ') 203 | if configs['db']['name'] == 'mvtec': 204 | if configs['db']['use_validation_set'] is True: 205 | test_mvtec(test_set, rebuilder, transform, args.res_dir, threshold_seg_dict, val_index) 206 | elif configs['db']['use_validation_set'] is False: 207 | test_mvtec(test_set, rebuilder, transform, args.res_dir, 64, val_index) 208 | else: 209 | raise Exception("Invalid input") 210 | elif configs['db']['name'] == 'chip': 211 | test_chip(test_set, rebuilder, transform, args.res_dir) 212 | else: 213 | raise Exception("Invalid set name") -------------------------------------------------------------------------------- /textured_surface_anomaly_detection/A_compact_convolutional_neura_network_for_textured_surface_anomaly_detection.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ForrestPi/Unsupervised-Defect-Segmentation/e366ac7c757bb1b45f38ebbc502dfee7ccb72398/textured_surface_anomaly_detection/A_compact_convolutional_neura_network_for_textured_surface_anomaly_detection.pdf -------------------------------------------------------------------------------- /textured_surface_anomaly_detection/README.md: -------------------------------------------------------------------------------- 1 | # textured_surface_anomaly_detection 2 | Textured surface anomaly detection 3 | Implemented by tensorflow and python2.7 4 | Refer "A compact convolutional neural network for textured surface anomaly detection" 5 | 6 | ### from 7 | https://github.com/HuayueZhang/textured_surface_anomaly_detection -------------------------------------------------------------------------------- /textured_surface_anomaly_detection/network.py: -------------------------------------------------------------------------------- 1 | # !/usr/bin/python 2 | # -*- coding: utf-8 -*- 3 | 4 | import tensorflow as tf 5 | import tf_utils 6 | 7 | 8 | def placeholder_inputs(batch_size, img_w, img_h): 9 | input_img_pl = tf.placeholder(tf.float32, shape=(batch_size, img_w, img_h)) 10 | label_seg_pl = tf.placeholder(tf.int32, shape=(batch_size, img_w, img_h)) 11 | label_cls_pl = tf.placeholder(tf.int32, shape=(batch_size)) 12 | return input_img_pl, label_seg_pl, label_cls_pl 13 | 14 | 15 | 16 | def cls_model(seg_out, seg_out_former, is_training, bn_decay): 17 | out_1 = tf_utils.conv2d(seg_out_former, kernel_shape=[1,1], strides=1, channel=32, 18 | activation_fn=tf.nn.relu, scope='conv11') 19 | 20 | out_2 = tf_utils.maxpool2d(seg_out, kernel_shape=[128,128], strides=1, padding='VALID') 21 | out_3 = tf_utils.avgpool2d(seg_out, kernel_shape=[128,128], strides=1, padding='VALID') 22 | out_4 = tf_utils.maxpool2d(out_1, kernel_shape=[128, 128], strides=1, padding='VALID') 23 | out_5 = tf_utils.avgpool2d(out_1, kernel_shape=[128, 128], strides=1, padding='VALID') 24 | 25 | out_6 = tf.concat([out_2, out_3, out_4, out_5], -1) 26 | 27 | out_7 = tf_utils.conv2d(out_6, kernel_shape=[1,1], strides=1, channel=1, 28 | bn=True, bn_decay=bn_decay, is_training=is_training, 29 | activation_fn=tf.nn.sigmoid, scope='conv12') 30 | 31 | return out_7 32 | 33 | 34 | def seg_model(input, is_training, bn_decay): 35 | extend_image = tf.expand_dims(input, axis=-1) 36 | 37 | out_1 = tf_utils.conv2d(extend_image, kernel_shape=[11,11], strides=2, channel=32, 38 | bn=True, bn_decay=bn_decay, is_training=is_training, 39 | activation_fn=tf.nn.relu, scope='conv1') 40 | out_2 = tf_utils.conv2d(out_1, kernel_shape=[11,11], strides=1, channel=32, 41 | bn=True, bn_decay=bn_decay, is_training=is_training, 42 | activation_fn=tf.nn.relu, scope='conv2') 43 | out_3 = tf_utils.conv2d(out_2, kernel_shape=[11,11], strides=1, channel=32, 44 | bn=True, bn_decay=bn_decay, is_training=is_training, 45 | activation_fn=tf.nn.relu, scope='conv3') 46 | 47 | out_4 = tf_utils.conv2d(out_3, kernel_shape=[7, 7], strides=2, channel=64, 48 | bn=True, bn_decay=bn_decay, is_training=is_training, 49 | activation_fn=tf.nn.relu, scope='conv4') 50 | out_5 = tf_utils.conv2d(out_4, kernel_shape=[7, 7], strides=1, channel=64, 51 | bn=True, bn_decay=bn_decay, is_training=is_training, 52 | activation_fn=tf.nn.relu, scope='conv5') 53 | out_6 = tf_utils.conv2d(out_5, kernel_shape=[7, 7], strides=1, channel=64, 54 | bn=True, bn_decay=bn_decay, is_training=is_training, 55 | activation_fn=tf.nn.relu, scope='conv6') 56 | 57 | out_7 = tf_utils.conv2d(out_6, kernel_shape=[3, 3], strides=1, channel=128, 58 | bn=True, bn_decay=bn_decay, is_training=is_training, 59 | activation_fn=tf.nn.relu, scope='conv7') 60 | out_8 = tf_utils.conv2d(out_7, kernel_shape=[3, 3], strides=1, channel=128, 61 | bn=True, bn_decay=bn_decay, is_training=is_training, 62 | activation_fn=tf.nn.relu, scope='conv8') 63 | out_9 = tf_utils.conv2d(out_8, kernel_shape=[3, 3], strides=1, channel=128, 64 | bn=True, bn_decay=bn_decay, is_training=is_training, 65 | activation_fn=tf.nn.relu, scope='conv9') 66 | 67 | seg_layer = tf_utils.conv2d(out_9, kernel_shape=[1, 1], strides=1, channel=1, 68 | bn=True, bn_decay=bn_decay, is_training=is_training, 69 | activation_fn=tf.nn.relu, scope='conv10') 70 | 71 | return seg_layer, out_9 72 | 73 | 74 | def get_loss(label, model_out, stage): 75 | label = tf.to_float(label) 76 | if stage == 'seg': 77 | # label or loss: shape = [batch_size, pixel_w, pixel_h] 78 | 79 | # shape = model_out.get_shape() 80 | # # tf.get_variable函数，变量名称name是一个必填的参数，它会根据变量名称去创建或者获取变量 81 | # linear_W = tf.get_variable(name='W', shape=shape) 82 | # linear_b = tf.get_variable(name='b', shape=shape) 83 | # pred = tf.multiply(model_out, linear_W) + linear_b 84 | # # tf.multiply对应位置点乘，tf.matmul矩阵数学乘法 85 | 86 | pred = tf_utils.unpool(model_out) 87 | pred = tf_utils.unpool(pred) 88 | pred = tf.squeeze(pred, -1) 89 | 90 | loss = tf.reduce_mean(tf.square(label - pred)) 91 | else: 92 | # stage == cls 93 | # label or loss: shape = [batch_size], 0 or 1 94 | pred = model_out 95 | loss = - tf.reduce_mean(label*tf.log(pred) + (1-label)*tf.log(1-pred)) 96 | return pred, loss -------------------------------------------------------------------------------- /textured_surface_anomaly_detection/provider.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | from scipy import misc 4 | import re 5 | import numpy as np 6 | 7 | 8 | def LOAD_DATA(data_path): 9 | label_path = data_path + 'Label/' 10 | cls_label = [] 11 | with open(label_path + 'Labels.txt') as f: 12 | for line in f.readlines(): 13 | line = re.findall('\d+', line) 14 | if len(line)>1: 15 | cls_label.append(int(line[1])) 16 | cls_label = np.array(cls_label) 17 | 18 | path_dir = os.listdir(data_path) 19 | path_dir.sort() 20 | img = [] 21 | for line in path_dir: 22 | if len(line) == 8: 23 | img.append(misc.imread(data_path + line)) 24 | img = np.array(img) 25 | 26 | label_dir = os.listdir(label_path) 27 | label_dir.sort() 28 | labeldir = [] 29 | for line in label_dir: 30 | if len(line) == 14: 31 | labeldir.append(line) 32 | seg_label = [] 33 | shape = img[0].shape 34 | i = 0 35 | for obj in cls_label: 36 | if obj == 0: 37 | seg_label.append(-np.ones(shape=shape)) 38 | else: 39 | temp = misc.imread(label_path + labeldir[i]) 40 | temp2 = 2 * (temp/255 - 0.5) 41 | seg_label.append(temp2) 42 | i += 1 43 | seg_label = np.array(seg_label) 44 | 45 | return img, seg_label, cls_label 46 | 47 | 48 | def shuffle_data(data, seg_label, cls_label): 49 | idx = np.arange(len(cls_label)) 50 | np.random.shuffle(idx) 51 | return data[idx], seg_label[idx], cls_label[idx] 52 | 53 | # img, sl, cl = LOAD_DATA(os.path.join(os.path.dirname(__file__), '../data/Class1/Train/')) 54 | # shuffle_data(img, sl, cl) -------------------------------------------------------------------------------- /textured_surface_anomaly_detection/tf_utils.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python 2 | # -*- coding: utf-8 -*- 3 | 4 | import tensorflow as tf 5 | 6 | def _variable_on_cpu(name, shape, initializer, use_fp16=False): 7 | """Helper to create a Variable stored on CPU memory. 8 | Args: 9 | name: name of the variable 10 | shape: list of ints 11 | initializer: initializer for Variable 12 | Returns: 13 | Variable Tensor 14 | """ 15 | with tf.device('/cpu:0'): 16 | dtype = tf.float16 if use_fp16 else tf.float32 17 | var = tf.get_variable(name, shape, initializer=initializer, dtype=dtype) 18 | return var 19 | 20 | def _variable_with_weight_decay(name, shape, stddev, wd, use_xavier=True): 21 | """Helper to create an initialized Variable with weight decay. 22 | 23 | Note that the Variable is initialized with a truncated normal distribution. 24 | A weight decay is added only if one is specified. 25 | 26 | Args: 27 | name: name of the variable 28 | shape: list of ints 29 | stddev: standard deviation of a truncated Gaussian 30 | wd: add L2Loss weight decay multiplied by this float. If None, weight 31 | decay is not added for this Variable. 32 | use_xavier: bool, whether to use xavier initializer 33 | 34 | Returns: 35 | Variable Tensor 36 | """ 37 | if use_xavier: 38 | initializer = tf.contrib.layers.xavier_initializer() 39 | else: 40 | initializer = tf.truncated_normal_initializer(stddev=stddev) 41 | var = _variable_on_cpu(name, shape, initializer) 42 | if wd is not None: 43 | weight_decay = tf.multiply(tf.nn.l2_loss(var), wd, name='weight_loss') 44 | tf.add_to_collection('losses', weight_decay) 45 | return var 46 | 47 | 48 | def conv2d(data, 49 | kernel_shape, 50 | strides, 51 | channel, 52 | activation_fn, 53 | scope, 54 | bn = False, 55 | bn_decay = None, 56 | is_training=None): # default参数必须位于非default参数之后 57 | 58 | 59 | with tf.variable_scope(scope) as sc: 60 | in_channel = data.get_shape()[-1] 61 | filter_shape = [kernel_shape[0], kernel_shape[1], in_channel, channel] 62 | # filter = tf.get_variable(name='filter', shape=filter_shape) 63 | 64 | filter = _variable_with_weight_decay('weights', 65 | shape=filter_shape, 66 | use_xavier=True, 67 | stddev=1e-3, 68 | wd=0.0) 69 | 70 | # filter接收的是一个实实在在的tensor，有内容，不仅仅是一个shape!! 71 | # 要搞清楚filter/kernel和它们shape的区别!! 72 | outputs = tf.nn.conv2d(input=data, 73 | filter=filter, 74 | strides=[1, strides, strides, 1], 75 | padding="SAME", 76 | use_cudnn_on_gpu=None, 77 | data_format=None, 78 | name=None) 79 | # tf.nn.conv2d只有卷积操作（乘法运算）， 没有加bias！！手动加上 80 | biases = _variable_on_cpu('biases', [channel], tf.constant_initializer(0.0)) 81 | outputs = tf.nn.bias_add(outputs, biases) 82 | 83 | if bn: 84 | outputs = batch_norm_for_conv2d(outputs, is_training, 85 | bn_decay=bn_decay, scope='bn') 86 | 87 | if activation_fn is not None: 88 | outputs = activation_fn(outputs) 89 | return outputs 90 | 91 | 92 | def maxpool2d(data, kernel_shape, strides, padding): 93 | return tf.nn.max_pool(value=data, 94 | ksize=[1, kernel_shape[0], kernel_shape[1], 1], 95 | strides=[1, strides, strides, 1], 96 | padding=padding, 97 | name=None) 98 | 99 | 100 | def avgpool2d(data, kernel_shape, strides, padding): 101 | return tf.nn.avg_pool(value=data, 102 | ksize=[1, kernel_shape[0], kernel_shape[1], 1], 103 | strides=[1, strides, strides, 1], 104 | padding=padding, 105 | name=None) 106 | 107 | 108 | def unpool(value, name='unpool'): 109 | """N-dimensional version of the unpooling operation from 110 | https://www.robots.ox.ac.uk/~vgg/rg/papers/Dosovitskiy_Learning_to_Generate_2015_CVPR_paper.pdf 111 | 112 | :param value: A Tensor of shape [b, d0, d1, ..., dn, ch] 113 | :return: A Tensor of shape [b, 2*d0, 2*d1, ..., 2*dn, ch] 114 | """ 115 | with tf.name_scope(name) as scope: 116 | sh = value.get_shape().as_list() 117 | dim = len(sh[1:-1]) 118 | out = (tf.reshape(value, [-1] + sh[-dim:])) 119 | # [-1] + sh[-dim:] 列表list的扩充 120 | 121 | for i in range(dim, 0, -1): 122 | # out = tf.concat([out, tf.zeros_like(out)], i) 123 | out = tf.concat([out, out], i) 124 | out_size = [-1] + [s * 2 for s in sh[1:-1]] + [sh[-1]] 125 | out = tf.reshape(out, out_size, name=scope) 126 | return out 127 | 128 | # value = tf.get_variable(name='value', shape=[3,2,2,4,5], dtype=tf.int32) 129 | # unpool(value=value) 130 | 131 | 132 | def batch_norm_template(inputs, is_training, scope, moments_dims, bn_decay): 133 | """ Batch normalization on convolutional maps and beyond... 134 | Ref.: http://stackoverflow.com/questions/33949786/how-could-i-use-batch-normalization-in-tensorflow 135 | 136 | Args: 137 | inputs: Tensor, k-D input ... x C could be BC or BHWC or BDHWC 138 | is_training: boolean tf.Variable, true indicates training phase 139 | scope: string, variable scope 140 | ********************************************************************************* 141 | * moments_dims: a list of ints, indicating dimensions for moments calculation * 142 | * 如果是图像数据, 可以传入 [0, 1, 2], 相当于求[batch, height, width] * 143 | * 的均值/方差, 注意不要加入 channel 维度 * 144 | ********************************************************************************* 145 | bn_decay: float or float tensor variable, controlling moving average weight 146 | Return: 147 | normed: batch-normalized maps 148 | """ 149 | with tf.variable_scope(scope) as sc: 150 | num_channels = inputs.get_shape()[-1].value 151 | beta = tf.Variable(tf.constant(0.0, shape=[num_channels]), 152 | name='beta', trainable=True) 153 | gamma = tf.Variable(tf.constant(1.0, shape=[num_channels]), 154 | name='gamma', trainable=True) 155 | batch_mean, batch_var = tf.nn.moments(inputs, moments_dims, name='moments') 156 | 157 | # 我们使用batch进行每次的更新, 那每个batch的mean/var都会不同, 所以我们可以使用moving average 158 | # 的方法记录并慢慢改进mean/var的值，尽可能的让mean/var建立在所有数据的基础上；然后将修改提升后的 159 | # mean/var放入tf.nn.batch_normalization().而且在test阶段, 我们就可以直接调用最后一次修改的 160 | # mean/var值进行测试, 而不是采用test时的mean/var. 161 | # 采用加权滑动平均，权重系数以decay衰减。 162 | decay = bn_decay if bn_decay is not None else 0.9 163 | ema = tf.train.ExponentialMovingAverage(decay=decay) 164 | 165 | # Operator that maintains moving average variables 166 | # 那如何确定我们是在train阶段还是在test阶段呢, 想办法传入is_training参数 167 | # 条件函数tf.cond（true_or_false, true_op, false_op） 168 | ema_apply_op = tf.cond(is_training, 169 | lambda: ema.apply([batch_mean, batch_var]), # is_training=true 170 | lambda: tf.no_op()) # is_training=false 171 | 172 | # Update moving average and return current batch's avg and var 173 | def mean_var_with_update(): 174 | with tf.control_dependencies([ema_apply_op]): 175 | return tf.identity(batch_mean), tf.identity(batch_var) 176 | 177 | # ema.average returns the Variable holding the average of var 178 | # 那如何确定我们是在train阶段还是在test阶段呢, 想办法传入is_training参数 179 | mean, var = tf.cond(is_training, # is_training值为True/False。 180 | mean_var_with_update, # 若True, 更新mean/var; 181 | lambda: (ema.average(batch_mean), ema.average(batch_var))) 182 | # 若False，返回之前mean/var的滑动平均值 183 | 184 | normed = tf.nn.batch_normalization(inputs, mean, var, beta, gamma, 1e-3) 185 | # 最后一步tf.nn.batch_normalization, 在做如下事情, inputs = Wx_plus_b: 186 | # Wx_plus_b = (Wx_plus_b - mean) / tf.sqrt(var + 1e-3) 187 | # Wx_plus_b = Wx_plus_b * gamma + beta 188 | return normed 189 | 190 | 191 | def batch_norm_for_fc(inputs, is_training, bn_decay, scope): 192 | return batch_norm_template(inputs, is_training, scope, [0,], bn_decay) 193 | 194 | 195 | def batch_norm_for_conv2d(inputs, is_training, bn_decay, scope): 196 | return batch_norm_template(inputs, is_training, scope, [0,1,2], bn_decay) 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 | -------------------------------------------------------------------------------- /textured_surface_anomaly_detection/train.py: -------------------------------------------------------------------------------- 1 | # !/usr/bin/python 2 | # -*- coding: utf-8 -*- 3 | 4 | import tensorflow as tf 5 | import numpy as np 6 | import os 7 | import sys 8 | import argparse 9 | 10 | import network 11 | import provider 12 | import tf_utils 13 | import math 14 | 15 | BASE_DIR = os.path.dirname(__file__) 16 | DATA_ROOT = os.path.join(BASE_DIR, '../data/Class1/') 17 | 18 | parser = argparse.ArgumentParser() 19 | parser.add_argument('--gpu', type=int, default=0) 20 | parser.add_argument('--batch_size', type=int, default=32) 21 | parser.add_argument('--base_learning_rate', type=float, default=0.001) 22 | parser.add_argument('--max_epoch', type=int, default=10) 23 | parser.add_argument('--max_epoch_seg', type=int, default=25) 24 | parser.add_argument('--max_epoch_cls', type=int, default=10) 25 | parser.add_argument('--optimizer', default='adam') 26 | parser.add_argument('--image_size', type=int, default=512) 27 | parser.add_argument('--decay_step', type=int, default=1000) 28 | parser.add_argument('--decay_rate', type=float, default=0.7) 29 | FLAGS = parser.parse_args() 30 | 31 | GPU_IDX = FLAGS.gpu 32 | BATCH_SIZE = FLAGS.batch_size 33 | BASE_LEARNING_RATE = FLAGS.base_learning_rate 34 | MAX_EPOCH = FLAGS.max_epoch 35 | MAX_EPOCH_seg = FLAGS.max_epoch_seg 36 | MAX_EPOCH_cls = FLAGS.max_epoch_cls 37 | OPTIMIZER = FLAGS.optimizer 38 | IMAGE_SIZE = FLAGS.image_size 39 | DECAY_STEP = FLAGS.decay_step 40 | DECAY_RATE = FLAGS.decay_rate 41 | 42 | LOG_DIR = os.path.join(BASE_DIR, '../log/') 43 | if not os.path.exists(LOG_DIR): os.mkdir(LOG_DIR) 44 | LOG_FOUT = open(os.path.join(LOG_DIR, 'log_train.txt'), 'w') 45 | LOG_FOUT.write(str(FLAGS)+'\n') 46 | 47 | BN_INIT_DECAY = 0.5 48 | BN_DECAY_DECAY_RATE = 0.5 49 | BN_DECAY_DECAY_STEP = float(DECAY_STEP) 50 | BN_DECAY_CLIP = 0.99 51 | 52 | def log_string(out_str): # 把每次print的东西都写入log 53 | LOG_FOUT.write(out_str+'\n') 54 | LOG_FOUT.flush() 55 | print(out_str) 56 | 57 | def get_learning_rate(global_step): 58 | learning_rate = tf.train.exponential_decay( 59 | BASE_LEARNING_RATE, global_step*BATCH_SIZE, DECAY_STEP, 60 | DECAY_RATE, staircase=True) 61 | learning_rate = tf.maximum(learning_rate, 0.00001) 62 | return learning_rate 63 | 64 | 65 | def get_bn_decay(global_step): 66 | bn_momentum = tf.train.exponential_decay( 67 | BN_INIT_DECAY, global_step*BATCH_SIZE, BN_DECAY_DECAY_STEP, 68 | BN_DECAY_DECAY_STEP, staircase=True) 69 | bn_decay = tf.minimum(BN_DECAY_CLIP, 1-bn_momentum) 70 | return bn_decay 71 | 72 | 73 | def train(): 74 | with tf.device('/gpu:'+str(GPU_IDX)): 75 | # placeholder的作用是什么，我们现在需要用到一些参数构建graph，但这些参数要到后面才能给 76 | # 而且不同的阶段这些参数不同，因此不能直接将它们给定，先用placeholder占个位 77 | # 在后面向graph喂数据的时候再把这些数据具体输进去 78 | img_pl, label_seg_pl, label_cls_pl = \ 79 | network.placeholder_inputs(BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE) 80 | is_training_pl = tf.placeholder(tf.bool, shape=()) 81 | print(is_training_pl) 82 | 83 | global_step = tf.Variable(0) 84 | bn_decay = get_bn_decay(global_step) 85 | 86 | seg_out, seg_out_former = network.seg_model(img_pl, is_training_pl, bn_decay) 87 | seg_pred, loss_seg = network.get_loss(label_seg_pl, seg_out, 'seg') 88 | 89 | cls_out = network.cls_model(seg_out, seg_out_former, is_training_pl, bn_decay) 90 | cls_pred, loss_cls = network.get_loss(label_cls_pl, cls_out, 'cls') 91 | 92 | learning_rate = get_learning_rate(global_step) 93 | tf.summary.scalar('learning_rate', learning_rate) 94 | 95 | if OPTIMIZER == 'momentum': 96 | optimizer = tf.train.MomentumOptimizer(learning_rate, momentum='MOMENTUM') 97 | elif OPTIMIZER == 'adam': 98 | optimizer = tf.train.AdamOptimizer(learning_rate) 99 | 100 | train_op_seg = optimizer.minimize(loss_seg, global_step=global_step) 101 | train_op_cls = optimizer.minimize(loss_cls, global_step=global_step) 102 | # global_step: Optional 'Variable' to increment by one after the 103 | # variables have been updated every time 104 | # 所有数据集训练完一次，称为一个epoch。 105 | # 在一个epoch内，每训练一个batchsize，参数更新一次，称为一个iteration/step， 106 | # 累计在所有epoch的训练过程中，iteration的数目成为global_step 107 | 108 | saver = tf.train.Saver() 109 | 110 | # Create a session 111 | config = tf.ConfigProto() 112 | config.gpu_options.allow_growth = True 113 | config.allow_soft_placement = True 114 | config.log_device_placement = True 115 | sess = tf.Session(config=config) 116 | 117 | init = tf.global_variables_initializer() 118 | 119 | sess.run(init) 120 | 121 | ops_seg = {'stage': 'seg', 122 | 'img_pl': img_pl, 123 | 'label_pl': label_seg_pl, 124 | 'is_training_pl': is_training_pl, 125 | 'pred': seg_pred, 126 | 'loss': loss_seg, 127 | 'train_op': train_op_seg, 128 | 'step': global_step} 129 | 130 | ops_cls = {'stage': 'cls', 131 | 'img_pl': img_pl, 132 | 'label_pl': label_cls_pl, 133 | 'is_training_pl': is_training_pl, 134 | 'pred': cls_pred, 135 | 'loss': loss_cls, 136 | 'train_op': train_op_cls, 137 | 'step': global_step} 138 | 139 | for e in range(MAX_EPOCH): 140 | log_string('----- EPOCH %03d -----'% e) 141 | for epoch in range(MAX_EPOCH_seg): 142 | log_string('----- SEG EPOCH %03d -----'% epoch) 143 | train_one_epoch(sess, ops_seg) 144 | eval_one_epoch(sess, ops_seg) 145 | 146 | for epoch in range(MAX_EPOCH_cls): 147 | log_string('----- CLS EPOCH %03d -----' % epoch) 148 | train_one_epoch(sess, ops_cls) 149 | eval_one_epoch(sess, ops_cls) 150 | 151 | if e % 1 == 0: 152 | save_path = saver.save(sess, os.path.join(LOG_DIR, "model.ckpt")) 153 | log_string("Model saved in file: %s" % save_path) 154 | 155 | 156 | def train_one_epoch(sess, ops): 157 | is_training = True 158 | img, seg_label, cls_label = provider.LOAD_DATA(DATA_ROOT + 'Train/') 159 | img, seg_label, cls_label = provider.shuffle_data(img, seg_label, cls_label) 160 | label = {'seg': seg_label, 'cls': cls_label} 161 | img_num = img.shape[0] 162 | batch_num = img_num / BATCH_SIZE 163 | 164 | loss_sum = 0 165 | 166 | for batch_idx in range(batch_num): 167 | start_idx = batch_idx * BATCH_SIZE 168 | end_idx = (batch_idx+1) * BATCH_SIZE 169 | 170 | current_imgs = img[start_idx:end_idx] 171 | current_label = label[ops['stage']][start_idx:end_idx] 172 | 173 | feed_dict = {ops['img_pl']: current_imgs, 174 | ops['label_pl']: current_label, 175 | ops['is_training_pl']: is_training} 176 | 177 | global_step, pred, loss, _ = sess.run([ops['step'], ops['pred'], ops['loss'], ops['train_op']], 178 | feed_dict=feed_dict) 179 | 180 | log_string('global_step: %d; iter: %d; loss: %f' % (global_step, batch_idx+1, loss)) 181 | loss_sum += loss 182 | 183 | log_string('train mean loss: %f' % (loss_sum / float(batch_num))) 184 | # print('accuracy: %f' % (total_correct / float(total_seen))) 185 | 186 | 187 | def eval_one_epoch(sess, ops): 188 | is_training = False 189 | img, seg_label, cls_label = provider.LOAD_DATA(DATA_ROOT + 'Test/') 190 | label = {'seg': seg_label, 'cls': cls_label} 191 | img_num = len(img) 192 | batch_num = img_num / BATCH_SIZE 193 | 194 | loss_sum = 0 195 | 196 | for batch_idx in range(batch_num): 197 | start_idx = batch_idx * BATCH_SIZE 198 | end_idx = (batch_idx + 1) * BATCH_SIZE 199 | 200 | current_imgs = img[start_idx:end_idx] 201 | current_label = label[ops['stage']][start_idx:end_idx] 202 | 203 | feed_dict = {ops['img_pl']: current_imgs, 204 | ops['label_pl']: current_label, 205 | ops['is_training_pl']: is_training} 206 | 207 | global_step, pred, loss = sess.run([ops['step'], ops['pred'], ops['loss']], 208 | feed_dict=feed_dict) 209 | 210 | loss_sum += loss 211 | 212 | log_string('eval mean loss: %f' % (loss_sum / float(batch_num))) 213 | # print('accuracy: %f' % (total_correct / float(total_seen))) 214 | 215 | if __name__ == "__main__": 216 | train() 217 | LOG_FOUT.close() -------------------------------------------------------------------------------- /textured_surface_anomaly_detection/unpool.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | 3 | def unpool(): 4 | 5 | value = tf.constant([[[[1],[2]],[[3],[4]]],[[[1],[2]],[[3],[4]]],[[[1],[2]],[[3],[4]]]]) 6 | 7 | sh = value.get_shape().as_list() 8 | dim = len(sh[1:-1]) 9 | out = (tf.reshape(value, [-1] + sh[-dim:])) 10 | 11 | for i in range(dim, 0, -1): 12 | out = tf.concat([out, out], i) 13 | out_size = [-1] + [s * 2 for s in sh[1:-1]] + [sh[-1]] 14 | out = tf.reshape(out, out_size) 15 | 16 | init = tf.global_variables_initializer() 17 | sess = tf.Session() 18 | sess.run(init) 19 | 20 | print(sess.run(value)) 21 | print(sess.run(out)) 22 | return out 23 | # unpool() -------------------------------------------------------------------------------- /tools/__init__.py: -------------------------------------------------------------------------------- 1 | from .log import Log 2 | from .timer import Timer -------------------------------------------------------------------------------- /tools/log.py: -------------------------------------------------------------------------------- 1 | import os 2 | import datetime 3 | 4 | 5 | def get_sys_date(): 6 | return datetime.datetime.now().strftime('%Y-%m-%d-h%H') 7 | 8 | 9 | class Log: 10 | def __init__(self, log_dir, cfg_name): 11 | if os.path.exists(log_dir) is not True: 12 | os.mkdir(log_dir) 13 | date = get_sys_date() 14 | f_name = '{}_{}.txt'.format(date, cfg_name) 15 | f_path = os.path.join(log_dir, f_name) 16 | self.log = open(f_path, 'w') 17 | 18 | def wr_cfg(self, configs): 19 | for _dict in configs: 20 | self.log.write(str(_dict) + ':\n') 21 | cfg_dict = configs[_dict] 22 | for item in cfg_dict: 23 | mes = ' {}: {}'.format(str(item), str(cfg_dict[item])) 24 | self.log.write(mes + '\n') 25 | self.log.write('loss:' + '\n') 26 | 27 | def wr_mes(self, mes): 28 | self.log.write(mes + '\n') 29 | 30 | def close(self): 31 | self.log.close() 32 | -------------------------------------------------------------------------------- /tools/timer.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | 4 | class Timer(object): 5 | """A simple timer.""" 6 | def __init__(self): 7 | self.total_time = 0. 8 | self.calls = 0 9 | self.start_time = 0. 10 | self.diff = 0. 11 | self.average_time = 0. 12 | 13 | def tic(self): 14 | # using time.time instead of time.clock because time time.clock 15 | # does not normalize for multithreading 16 | self.start_time = time.time() 17 | 18 | def toc(self, average=True): 19 | self.diff = time.time() - self.start_time 20 | self.total_time += self.diff 21 | self.calls += 1 22 | self.average_time = self.total_time / self.calls 23 | if average: 24 | return self.average_time 25 | else: 26 | return self.diff 27 | 28 | def clear(self): 29 | self.total_time = 0. 30 | self.calls = 0 31 | self.start_time = 0. 32 | self.diff = 0. 33 | self.average_time = 0. -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | import argparse 4 | from db import training_collate 5 | from tools import Timer, Log 6 | from factory import * 7 | 8 | 9 | def parse_args(): 10 | parser = argparse.ArgumentParser() 11 | parser.add_argument('--cfg', help='Name of .json file', type=str, required=True) 12 | parser.add_argument('--ngpu', help='Numbers of GPU', type=int, default=1) 13 | parser.add_argument('--log_dir', help="Directory of training log", type=str, default='./log') 14 | args = parser.parse_args() 15 | 16 | return args 17 | 18 | 19 | def adjust_learning_rate(trainer, init_lr, decay_rate, epoch, step_index, iteration, epoch_size): 20 | if epoch < 6: 21 | lr = 1e-6 + (init_lr-1e-6) * iteration / (epoch_size * 5) 22 | else: 23 | lr = init_lr / (decay_rate ** (step_index)) 24 | 25 | trainer.set_lr(lr) 26 | 27 | return lr 28 | 29 | 30 | if __name__ == '__main__': 31 | args = parse_args() 32 | 33 | # load config 34 | cfg_file = os.path.join('./config', args.cfg + '.json') 35 | with open(cfg_file, "r") as f: 36 | configs = json.load(f) 37 | start_epoch = configs['op']['start_epoch'] 38 | max_epoch = configs['op']['max_epoch'] 39 | learning_rate = configs['op']['learning_rate'] 40 | decay_rate = configs['op']['decay_rate'] 41 | epoch_steps = configs['op']['epoch_steps'] 42 | snapshot = configs['op']['snapshot'] 43 | batch_size = configs['db']['batch_size'] 44 | loader_threads = configs['db']['loader_threads'] 45 | save_dir = configs['system']['save_dir'] 46 | 47 | # init Timer 48 | if not os.path.exists(save_dir): 49 | os.mkdir(save_dir) 50 | _t = Timer() 51 | 52 | # create log file 53 | log = Log(args.log_dir, args.cfg) 54 | log.wr_cfg(configs) 55 | 56 | # load data set 57 | training_set = load_data_set_from_factory(configs, 'train') 58 | print('Data set: {} has been loaded'.format(configs['db']['name'])) 59 | 60 | # load model 61 | trainer = load_training_model_from_factory(configs, ngpu=args.ngpu) 62 | if configs['system']['resume']: 63 | trainer.load_params(configs['system']['resume_path']) 64 | print('Model: {} has been loaded'.format(configs['model']['name'])) 65 | 66 | # start training 67 | epoch_size = len(training_set) // batch_size # init learning rate & iters 68 | start_iter = start_epoch * epoch_size 69 | max_iter = max_epoch * epoch_size 70 | print('Start training...') 71 | epoch = 0 72 | iters_steps = [epoch_step*epoch_size for epoch_step in epoch_steps] 73 | for iteration in range(start_iter, max_iter): 74 | # reset batch iterator 75 | if iteration % epoch_size == 0: 76 | batch_iterator = iter(torch.utils.data.DataLoader(training_set, batch_size, shuffle=True, 77 | num_workers=loader_threads, collate_fn=training_collate)) 78 | # save parameters 79 | if epoch % snapshot == 0 and iteration > start_iter: 80 | save_name = '{}-{:d}.pth'.format(args.cfg, epoch) 81 | save_path = os.path.join(save_dir, save_name) 82 | trainer.save_params(save_path) 83 | epoch += 1 84 | 85 | # adjust learning rate 86 | step_index = len(iters_steps) 87 | for k, step in enumerate(iters_steps): 88 | if iteration < step: 89 | step_index = k 90 | break 91 | lr = adjust_learning_rate(trainer, learning_rate, decay_rate, epoch, step_index, iteration, epoch_size) 92 | 93 | # load data 94 | _t.tic() 95 | images = next(batch_iterator) 96 | if configs['model']['type'] == 'Encoder': 97 | trainer.train(images) 98 | elif configs['model']['type'] == 'GAN': 99 | # Update Discriminator 100 | trainer.train(images, phase='discriminate') 101 | # Update Generator 102 | trainer.train(images, phase='generate') 103 | else: 104 | raise Exception("Wrong model type!") 105 | batch_time = _t.toc() 106 | 107 | # print message 108 | if iteration % 10 == 0: 109 | _t.clear() 110 | mes = 'Epoch:' + repr(epoch) + '||epochiter: ' + repr(iteration % epoch_size) + '/' + repr(epoch_size) 111 | mes += '||Totel iter: ' + repr(iteration) 112 | mes += '||{}'.format(trainer.get_loss_message()) 113 | mes += '||LR: %.8f' % (lr) 114 | mes += '||Batch time: %.4f sec.' % batch_time 115 | log.wr_mes(mes) 116 | print(mes) 117 | save_name = '{}-{:d}.pth'.format(args.cfg, epoch) 118 | save_path = os.path.join(save_dir, save_name) 119 | trainer.save_params(save_path) 120 | log.close() 121 | exit(0) 122 | 123 | -------------------------------------------------------------------------------- /train.sh: -------------------------------------------------------------------------------- 1 | CUDA_VISIBLE_DEVICES=0 python train.py --cfg="RED_Net_4skips-mvtec" --ngpu=1 --log_dir="./log" --------------------------------------------------------------------------------