first commit

Evaluate.py

@@ -0,0 +1,155 @@
+import numpy as np
+import os
+import sys
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import torchvision
+import torch.nn.init as init
+import torch.utils.data as data
+import torch.utils.data.dataset as dataset
+import torchvision.datasets as dset
+import torchvision.transforms as transforms
+from torch.autograd import Variable
+import torchvision.utils as v_utils
+import matplotlib.pyplot as plt
+import cv2
+import math
+from collections import OrderedDict
+import copy
+import time
+from model.utils import DataLoader
+from model.Reconstruction import *
+from sklearn.metrics import roc_auc_score
+from utils import *
+import random
+
+import argparse
+
+
+parser = argparse.ArgumentParser(description="MNAD")
+parser.add_argument('--gpus', nargs='+', type=str, help='gpus')
+parser.add_argument('--batch_size', type=int, default=4, help='batch size for training')
+parser.add_argument('--test_batch_size', type=int, default=1, help='batch size for test')
+parser.add_argument('--epochs', type=int, default=60, help='number of epochs for training')
+parser.add_argument('--loss_compact', type=float, default=0.01, help='weight of the feature compactness loss')
+parser.add_argument('--loss_separate', type=float, default=0.01, help='weight of the feature separateness loss')
+parser.add_argument('--h', type=int, default=256, help='height of input images')
+parser.add_argument('--w', type=int, default=256, help='width of input images')
+parser.add_argument('--c', type=int, default=3, help='channel of input images')
+parser.add_argument('--lr', type=float, default=2e-4, help='initial learning rate')
+parser.add_argument('--t_length', type=int, default=5, help='length of the frame sequences')
+parser.add_argument('--fdim', type=int, default=512, help='channel dimension of the features')
+parser.add_argument('--mdim', type=int, default=512, help='channel dimension of the memory items')
+parser.add_argument('--msize', type=int, default=10, help='number of the memory items')
+parser.add_argument('--alpha', type=float, default=0.7, help='weight for the anomality score')
+parser.add_argument('--num_workers', type=int, default=2, help='number of workers for the train loader')
+parser.add_argument('--num_workers_test', type=int, default=1, help='number of workers for the test loader')
+parser.add_argument('--dataset_type', type=str, default='ped2', help='type of dataset: ped2, avenue, shanghai')
+parser.add_argument('--dataset_path', type=str, default='./dataset/', help='directory of data')
+parser.add_argument('--model_dir', type=str, help='directory of model')
+parser.add_argument('--m_items_dir', type=str, help='directory of model')
+
+args = parser.parse_args()
+
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+if args.gpus is None:
+    gpus = "0"
+    os.environ["CUDA_VISIBLE_DEVICES"]= gpus
+else:
+    gpus = ""
+    for i in range(len(args.gpus)):
+        gpus = gpus + args.gpus[i] + ","
+    os.environ["CUDA_VISIBLE_DEVICES"]= gpus[:-1]
+
+torch.backends.cudnn.enabled = True # make sure to use cudnn for computational performance
+
+test_folder = args.dataset_path+args.dataset_type+"/testing/frames"
+
+# Loading dataset
+test_dataset = DataLoader(test_folder, transforms.Compose([
+             transforms.ToTensor(),            
+             ]), resize_height=args.h, resize_width=args.w, time_step=args.t_length-1)
+
+test_size = len(test_dataset)
+
+test_batch = data.DataLoader(test_dataset, batch_size = args.test_batch_size, 
+                             shuffle=False, num_workers=args.num_workers_test, drop_last=False)
+
+
+# Loading the trained model
+model = torch.load(args.model_dir)
+model.cuda()
+m_items = torch.load(args.m_itmes_dir)
+labels = np.load('./data/frame_labels_'+args.dataset_type+'.npy')
+
+videos = OrderedDict()
+videos_list = sorted(glob.glob(os.path.join(test_folder, '*')))
+for video in videos_list:
+    video_name = video.split('/')[-1]
+    videos[video_name] = {}
+    videos[video_name]['path'] = video
+    videos[video_name]['frame'] = glob.glob(os.path.join(video, '*.jpg'))
+    videos[video_name]['frame'].sort()
+    videos[video_name]['length'] = len(videos[video_name]['frame'])
+
+labels_list = []
+label_length = 0
+psnr_list = {}
+feature_distance_list = {}
+
+print('Evaluation of', args.dataset_type)
+
+# Setting for video anomaly detection
+for video in sorted(videos_new):
+    video_name = video.split('/')[-1]
+    labels_list = np.append(labels_list, labels[0][4+label_length:videos[video_name]['length']+label_length])
+    label_length += videos[video_name]['length']
+    psnr_list[video_name] = []
+    feature_distance_list[video_name] = []
+
+label_length = 0
+video_num = 0
+label_length += videos[videos_list[video_num].split('/')[-1]]['length']
+m_items_test = m_items.clone()
+
+model.eval()
+
+for k,(imgs) in enumerate(test_batch):
+
+    if k == label_length-4*(video_num+1):
+        video_num += 1
+        label_length += videos[videos_list[video_num].split('/')[-1]]['length']
+
+    imgs = Variable(imgs).cuda()
+
+    outputs, feas, updated_feas, m_items_test, softmax_score_query, softmax_score_memory, compactness_loss = model.forward(imgs[:,0:3*4], m_items_test, False)
+    mse_imgs = torch.mean(loss_func_mse((outputs[0]+1)/2, (imgs[0,3*4:]+1)/2)).item()
+    mse_feas = compactness_loss.item()
+
+    # Calculating the threshold for updating at the test time
+    point_sc = point_score(outputs, imgs[:,3*4:])
+
+    if  point_sc < threshold:
+        query = F.normalize(feas, dim=1)
+        query = query.permute(0,2,3,1) # b X h X w X d
+        m_items_test = model.memory.update(query, m_items_test, False)
+
+    psnr_list[videos_list[video_num].split('/')[-1]].append(psnr(mse_imgs))
+    feature_distance_list[videos_list[video_num].split('/')[-1]].append(mse_feas)
+
+
+# Measuring the abnormality score and the AUC
+anomaly_score_total_list = []
+for video in sorted(videos_new):
+    video_name = video.split('/')[-1]
+    anomaly_score_total_list += score_sum(anomaly_score_list(psnr_list[video_name]), 
+                                     anomaly_score_list_inv(feature_distance_list[video_name]), args.alpha)
+
+anomaly_score_total_list = np.asarray(anomaly_score_total_list)
+
+accuracy = AUC(anomaly_score_total_list, np.expand_dims(1-labels_list, 0))
+
+print('The result of ', args.dataset_type)
+print('AUC: ', accuracy*100, '%')

README.md

@@ -0,0 +1,55 @@
+# PyTorch implementation of "Learning Memory-guided Normality for Anomaly Detection"
+
+<p align="center"><img src="../MNAD_files/overview.png" alt="no_image" width="40%" height="40%" /><img src="../MNAD_files/teaser.png" alt="no_image" width="60%" height="60%" /></p>
+This is the implementation of the paper "Learning Memory-guided Normality for Anomaly Detection (CVPR 2020)".
+
+For more information, checkout the project site [[website](https://cvlab.yonsei.ac.kr/projects/MNAD/)] and the paper [[PDF](http://openaccess.thecvf.com/content_CVPR_2020/papers/Park_Learning_Memory-Guided_Normality_for_Anomaly_Detection_CVPR_2020_paper.pdf)].
+
+## Dependencies
+* Python 3.6
+* PyTorch >= 1.0.0
+* Numpy
+* Sklearn
+
+## Datasets
+* USCD Ped2 [[dataset](http://www.svcl.ucsd.edu/projects/anomaly/dataset.html)]
+* CUHK Avenue [[dataset](http://www.cse.cuhk.edu.hk/leojia/projects/detectabnormal/dataset.html)]
+* ShanghaiTech [[dataset](https://github.com/desenzhou/ShanghaiTechDataset)]
+
+Downlaod the datasets into ``datasets`` folder, like ``./datasets/ped2/``
+
+## Training
+```bash
+git clone https://github.com/cvlab-yonsei/projects
+cd projects/MNAD/code
+python Train.py # for training
+```
+* You can freely define parameters with your own settings like
+```bash
+python Train.py --gpus 1 --dataset_path 'your_dataset_directory' --dataset_type avenue --exp_dir 'your_log_directory'
+```
+
+## Pre-trained model and memory items
+* Download our pre-trained model and memory items <br>Link: [[model and items](https://drive.google.com/file/d/11f65puuljkUa0Z4W0VtkF_2McphS02fq/view?usp=sharing)]
+* Note that, these are from training with the Ped2 dataset
+
+## Evaluation
+* Test the model with our pre-trained model and memory items
+```bash
+python Evaluate.py --model_dir pretrained_model.pth --m_items_dir m_items.pt
+```
+* Test your own model
+```bash
+python Evaluate.py --model_dir your_model.pth --m_items_dir your_m_items.pt
+```
+
+## Bibtex
+```
+@inproceedings{park2020learning,
+  title={Learning Memory-guided Normality for Anomaly Detection},
+  author={Park, Hyunjong and Noh, Jongyoun and Ham, Bumsub},
+  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
+  pages={14372--14381},
+  year={2020}
+}
+```

Train.py

@@ -0,0 +1,148 @@
+import numpy as np
+import os
+import sys
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import torchvision
+import torch.nn.init as init
+import torch.utils.data as data
+import torch.utils.data.dataset as dataset
+import torchvision.datasets as dset
+import torchvision.transforms as transforms
+from torch.autograd import Variable
+import torchvision.utils as v_utils
+import matplotlib.pyplot as plt
+import cv2
+import math
+from collections import OrderedDict
+import copy
+import time
+from model.utils import DataLoader
+from model.Reconstruction import *
+from sklearn.metrics import roc_auc_score
+from utils import *
+import random
+
+import argparse
+
+
+parser = argparse.ArgumentParser(description="MNAD")
+parser.add_argument('--gpus', nargs='+', type=str, help='gpus')
+parser.add_argument('--batch_size', type=int, default=4, help='batch size for training')
+parser.add_argument('--test_batch_size', type=int, default=1, help='batch size for test')
+parser.add_argument('--epochs', type=int, default=60, help='number of epochs for training')
+parser.add_argument('--loss_compact', type=float, default=0.01, help='weight of the feature compactness loss')
+parser.add_argument('--loss_separate', type=float, default=0.01, help='weight of the feature separateness loss')
+parser.add_argument('--h', type=int, default=256, help='height of input images')
+parser.add_argument('--w', type=int, default=256, help='width of input images')
+parser.add_argument('--c', type=int, default=3, help='channel of input images')
+parser.add_argument('--lr', type=float, default=2e-4, help='initial learning rate')
+parser.add_argument('--t_length', type=int, default=5, help='length of the frame sequences')
+parser.add_argument('--fdim', type=int, default=512, help='channel dimension of the features')
+parser.add_argument('--mdim', type=int, default=512, help='channel dimension of the memory items')
+parser.add_argument('--msize', type=int, default=10, help='number of the memory items')
+parser.add_argument('--alpha', type=float, default=0.7, help='weight for the anomality score')
+parser.add_argument('--num_workers', type=int, default=2, help='number of workers for the train loader')
+parser.add_argument('--num_workers_test', type=int, default=1, help='number of workers for the test loader')
+parser.add_argument('--dataset_type', type=str, default='ped2', help='type of dataset: ped2, avenue, shanghai')
+parser.add_argument('--dataset_path', type=str, default='./dataset/', help='directory of data')
+parser.add_argument('--exp_dir', type=str, default='log', help='directory of log')
+
+args = parser.parse_args()
+
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+if args.gpus is None:
+    gpus = "0"
+    os.environ["CUDA_VISIBLE_DEVICES"]= gpus
+else:
+    gpus = ""
+    for i in range(len(args.gpus)):
+        gpus = gpus + args.gpus[i] + ","
+    os.environ["CUDA_VISIBLE_DEVICES"]= gpus[:-1]
+
+torch.backends.cudnn.enabled = True # make sure to use cudnn for computational performance
+
+train_folder = args.dataset_path+args.dataset_type+"/training/frames"
+test_folder = args.dataset_path+args.dataset_type+"/testing/frames"
+
+# Loading dataset
+train_dataset = DataLoader(train_folder, transforms.Compose([
+             transforms.ToTensor(),          
+             ]), resize_height=args.h, resize_width=args.w, time_step=args.t_length-1)
+
+test_dataset = DataLoader(test_folder, transforms.Compose([
+             transforms.ToTensor(),            
+             ]), resize_height=args.h, resize_width=args.w, time_step=args.t_length-1)
+
+train_size = len(train_dataset)
+test_size = len(test_dataset)
+
+train_batch = data.DataLoader(train_dataset, batch_size = args.batch_size, 
+                              shuffle=True, num_workers=args.num_workers, drop_last=True)
+test_batch = data.DataLoader(test_dataset, batch_size = args.test_batch_size, 
+                             shuffle=False, num_workers=args.num_workers_test, drop_last=False)
+
+
+# Model setting
+model = convAE(args.c, args.t_length, args.msize, args.fdim, args.mdim)
+params_encoder =  list(model.encoder.parameters()) 
+params_decoder = list(model.decoder.parameters())
+params = params_encoder + params_decoder
+optimizer = torch.optim.Adam(params, lr = args.lr)
+scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer,T_max =args.epochs)
+model.cuda()
+
+
+# Report the training process
+log_dir = os.path.join('./exp', args.dataset_type, args.exp_dir)
+if not os.path.exists(log_dir):
+    os.makedirs(log_dir)
+orig_stdout = sys.stdout
+f = open(os.path.join(log_dir, 'log.txt'),'w')
+sys.stdout= f
+
+loss_func_mse = nn.MSELoss(reduction='none')
+
+# Training
+
+m_items = F.normalize(torch.rand((args.msize, args.mdim), dtype=torch.float), dim=1).cuda() # Initialize the memory items
+
+for epoch in range(args.epochs):
+    labels_list = []
+    model.train()
+
+    start = time.time()
+    for j,(imgs) in enumerate(train_batch):
+
+        imgs = Variable(imgs).cuda()
+
+        outputs, _, _, m_items, softmax_score_query, softmax_score_memory, compactness_loss, separateness_loss = model.forward(imgs[:,0:12], m_items, True)
+
+
+        optimizer.zero_grad()
+        loss_pixel = torch.mean(loss_func_mse(outputs, imgs[:,12:]))
+        loss = loss_pixel + args.loss_compact * compactness_loss + args.loss_separate * separateness_loss
+        loss.backward(retain_graph=True)
+        optimizer.step()
+
+    scheduler.step()
+
+    print('----------------------------------------')
+    print('Epoch:', epoch)
+    print('Loss: Reconstruction {:.6f}/ Compactness {:.6f}/ Separateness {:.6f}'.format(loss_pixel.item(), compactness_loss.item(), separateness_loss.item()))
+    print('Memory_items:')
+    print(m_items)
+    print('----------------------------------------')
+
+print('Training is finished')
+# Save the model and the memory items
+torch.save(model, os.path.join(log_dir, 'model.pth'))
+torch.save(m_items, os.path.join(log_dir, 'keys.pt'))
+
+sys.stdout = orig_stdout
+f.close()
+
+
+