[ { "path": "Evaluate.py", "content": "import numpy as np\nimport os\nimport sys\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nimport torchvision\nimport torch.nn.init as init\nimport torch.utils.data as data\nimport torch.utils.data.dataset as dataset\nimport torchvision.datasets as dset\nimport torchvision.transforms as transforms\nfrom torch.autograd import Variable\nimport torchvision.utils as v_utils\nimport matplotlib.pyplot as plt\nimport cv2\nimport math\nfrom collections import OrderedDict\nimport copy\nimport time\nfrom model.utils import DataLoader\nfrom model.final_future_prediction_with_memory_spatial_sumonly_weight_ranking_top1 import *\nfrom model.Reconstruction import *\nfrom sklearn.metrics import roc_auc_score\nfrom utils import *\nimport random\nimport glob\n\nimport argparse\n\n\nparser = argparse.ArgumentParser(description=\"MNAD\")\nparser.add_argument('--gpus', nargs='+', type=str, help='gpus')\nparser.add_argument('--batch_size', type=int, default=4, help='batch size for training')\nparser.add_argument('--test_batch_size', type=int, default=1, help='batch size for test')\nparser.add_argument('--h', type=int, default=256, help='height of input images')\nparser.add_argument('--w', type=int, default=256, help='width of input images')\nparser.add_argument('--c', type=int, default=3, help='channel of input images')\nparser.add_argument('--method', type=str, default='pred', help='The target task for anoamly detection')\nparser.add_argument('--t_length', type=int, default=5, help='length of the frame sequences')\nparser.add_argument('--fdim', type=int, default=512, help='channel dimension of the features')\nparser.add_argument('--mdim', type=int, default=512, help='channel dimension of the memory items')\nparser.add_argument('--msize', type=int, default=10, help='number of the memory items')\nparser.add_argument('--alpha', type=float, default=0.6, help='weight for the anomality score')\nparser.add_argument('--th', type=float, default=0.01, help='threshold for test updating')\nparser.add_argument('--num_workers', type=int, default=2, help='number of workers for the train loader')\nparser.add_argument('--num_workers_test', type=int, default=1, help='number of workers for the test loader')\nparser.add_argument('--dataset_type', type=str, default='ped2', help='type of dataset: ped2, avenue, shanghai')\nparser.add_argument('--dataset_path', type=str, default='./dataset', help='directory of data')\nparser.add_argument('--model_dir', type=str, help='directory of model')\nparser.add_argument('--m_items_dir', type=str, help='directory of model')\n\nargs = parser.parse_args()\n\nos.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"\nif args.gpus is None:\n gpus = \"0\"\n os.environ[\"CUDA_VISIBLE_DEVICES\"]= gpus\nelse:\n gpus = \"\"\n for i in range(len(args.gpus)):\n gpus = gpus + args.gpus[i] + \",\"\n os.environ[\"CUDA_VISIBLE_DEVICES\"]= gpus[:-1]\n\ntorch.backends.cudnn.enabled = True # make sure to use cudnn for computational performance\n\ntest_folder = args.dataset_path+\"/\"+args.dataset_type+\"/testing/frames\"\n\n# Loading dataset\ntest_dataset = DataLoader(test_folder, transforms.Compose([\n transforms.ToTensor(), \n ]), resize_height=args.h, resize_width=args.w, time_step=args.t_length-1)\n\ntest_size = len(test_dataset)\n\ntest_batch = data.DataLoader(test_dataset, batch_size = args.test_batch_size, \n shuffle=False, num_workers=args.num_workers_test, drop_last=False)\n\nloss_func_mse = nn.MSELoss(reduction='none')\n\n# Loading the trained model\nmodel = torch.load(args.model_dir)\nmodel.cuda()\nm_items = torch.load(args.m_items_dir)\nlabels = np.load('./data/frame_labels_'+args.dataset_type+'.npy')\n\nvideos = OrderedDict()\nvideos_list = sorted(glob.glob(os.path.join(test_folder, '*')))\nfor video in videos_list:\n video_name = video.split('/')[-1]\n videos[video_name] = {}\n videos[video_name]['path'] = video\n videos[video_name]['frame'] = glob.glob(os.path.join(video, '*.jpg'))\n videos[video_name]['frame'].sort()\n videos[video_name]['length'] = len(videos[video_name]['frame'])\n\nlabels_list = []\nlabel_length = 0\npsnr_list = {}\nfeature_distance_list = {}\n\nprint('Evaluation of', args.dataset_type)\n\n# Setting for video anomaly detection\nfor video in sorted(videos_list):\n video_name = video.split('/')[-1]\n if args.method == 'pred':\n labels_list = np.append(labels_list, labels[0][4+label_length:videos[video_name]['length']+label_length])\n else:\n labels_list = np.append(labels_list, labels[0][label_length:videos[video_name]['length']+label_length])\n label_length += videos[video_name]['length']\n psnr_list[video_name] = []\n feature_distance_list[video_name] = []\n\nlabel_length = 0\nvideo_num = 0\nlabel_length += videos[videos_list[video_num].split('/')[-1]]['length']\nm_items_test = m_items.clone()\n\nmodel.eval()\n\nfor k,(imgs) in enumerate(test_batch):\n \n if args.method == 'pred':\n if k == label_length-4*(video_num+1):\n video_num += 1\n label_length += videos[videos_list[video_num].split('/')[-1]]['length']\n else:\n if k == label_length:\n video_num += 1\n label_length += videos[videos_list[video_num].split('/')[-1]]['length']\n\n imgs = Variable(imgs).cuda()\n \n if args.method == 'pred':\n 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)\n mse_imgs = torch.mean(loss_func_mse((outputs[0]+1)/2, (imgs[0,3*4:]+1)/2)).item()\n mse_feas = compactness_loss.item()\n\n # Calculating the threshold for updating at the test time\n point_sc = point_score(outputs, imgs[:,3*4:])\n \n else:\n outputs, feas, updated_feas, m_items_test, softmax_score_query, softmax_score_memory, compactness_loss = model.forward(imgs, m_items_test, False)\n mse_imgs = torch.mean(loss_func_mse((outputs[0]+1)/2, (imgs[0]+1)/2)).item()\n mse_feas = compactness_loss.item()\n\n # Calculating the threshold for updating at the test time\n point_sc = point_score(outputs, imgs)\n\n if point_sc < args.th:\n query = F.normalize(feas, dim=1)\n query = query.permute(0,2,3,1) # b X h X w X d\n m_items_test = model.memory.update(query, m_items_test, False)\n\n psnr_list[videos_list[video_num].split('/')[-1]].append(psnr(mse_imgs))\n feature_distance_list[videos_list[video_num].split('/')[-1]].append(mse_feas)\n\n\n# Measuring the abnormality score and the AUC\nanomaly_score_total_list = []\nfor video in sorted(videos_list):\n video_name = video.split('/')[-1]\n anomaly_score_total_list += score_sum(anomaly_score_list(psnr_list[video_name]), \n anomaly_score_list_inv(feature_distance_list[video_name]), args.alpha)\n\nanomaly_score_total_list = np.asarray(anomaly_score_total_list)\n\naccuracy = AUC(anomaly_score_total_list, np.expand_dims(1-labels_list, 0))\n\nprint('The result of ', args.dataset_type)\nprint('AUC: ', accuracy*100, '%')\n" }, { "path": "MNAD_files/style.css", "content": "/* Space out content a bit */\n\n@import url('https://fonts.googleapis.com/css?family=Baloo|Bungee+Inline|Lato|Righteous|Shojumaru');\n\nbody {\n padding-top: 20px;\n padding-bottom: 20px;\n font-family: 'Lato', cursive;\n font-size: 14px;\n}\n\n/* Everything but the jumbotron gets side spacing for mobile first views */\n.header,\n.row,\n.footer {\n padding-left: 15px;\n padding-right: 15px;\n}\n\n/* Custom page header */\n.header {\n border-bottom: 1px solid #e5e5e5;\n}\n/* Make the masthead heading the same height as the navigation */\n.header h1 {\n margin-top: 0;\n margin-bottom: 0;\n line-height: 40px;\n padding-bottom: 19px;\n font-size: 30px;\n font-weight: bold;\n}\n.header h3 {\n margin-top: 0;\n margin-bottom: 0;\n line-height: 40px;\n padding-bottom: 19px;\n font-size: 20px;\n}\n.header h4 {\n font-family: 'Baloo', cursive;\n}\n\n/* Custom page footer */\n.footer {\n padding-top: 19px;\n color: #777;\n border-top: 1px solid #e5e5e5;\n}\n\n/* Customize container */\n@media (min-width: 938px) {\n .container {\n max-width: 900px;\n }\n}\n.container-narrow > hr {\n margin: 20px 0;\n}\n\n/* Main marketing message and sign up button */\n.container .jumbotron {\n text-align: center;\n border-bottom: 1px solid #e5e5e5;\n padding-left: 20px;\n padding: 30px;\n}\n.jumbotron .btn {\n font-size: 21px;\n padding: 14px 24px;\n}\n\n.row p + h3 {\n margin-top: 28px;\n}\n\ndiv.row h3 {\n padding-bottom: 5px;\n border-bottom: 1px solid #ccc;\n}\n\n/* Responsive: Portrait tablets and up */\n@media screen and (min-width: 938px) {\n /* Remove the padding we set earlier */\n .header,\n .marketing,\n .footer {\n padding-left: 0;\n padding-right: 0;\n }\n /* Space out the masthead */\n .header {\n margin-bottom: 30px;\n }\n /* Remove the bottom border on the jumbotron for visual effect */\n .jumbotron {\n border-bottom: 0;\n }\n}\n\n.readme h1 {\n display: none;\n}\n\n.left_column{\n float:middle;\n \n}\n\n.right_column{\n float:middle;\n \n}" }, { "path": "README.md", "content": "# PyTorch implementation of \"Learning Memory-guided Normality for Anomaly Detection\"\n\n

\"no_image\"\"no_image\"

\nThis is the implementation of the paper \"Learning Memory-guided Normality for Anomaly Detection (CVPR 2020)\".\n\nFor 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)].\n\n## Dependencies\n* Python 3.6\n* PyTorch 1.1.0\n* Numpy\n* Sklearn\n\n## Datasets\n* USCD Ped2 [[dataset](https://github.com/StevenLiuWen/ano_pred_cvpr2018)]\n* CUHK Avenue [[dataset](https://github.com/StevenLiuWen/ano_pred_cvpr2018)]\n* ShanghaiTech [[dataset](https://github.com/StevenLiuWen/ano_pred_cvpr2018)]\n\nThese datasets are from an official github of \"Future Frame Prediction for Anomaly Detection - A New Baseline (CVPR 2018)\".\n\nDownload the datasets into ``dataset`` folder, like ``./dataset/ped2/``\n\n## Update\n* 02/04/21: We uploaded the codes based on reconstruction method, and pretrained wieghts for Ped2 reconstruction, Avenue prediction and Avenue reconstruction.\n\n\n## Training\n* ~~The training and testing codes are based on prediction method~~\n* Now you can implemnet the codes based on both prediction and reconstruction methods.\n* The codes are basically based on the prediction method, and you can easily implement this as\n```bash\ngit clone https://github.com/cvlab-yonsei/projects\ncd projects/MNAD/code\npython Train.py # for training\n```\n* You can freely define parameters with your own settings like\n```bash\npython Train.py --gpus 1 --dataset_path 'your_dataset_directory' --dataset_type avenue --exp_dir 'your_log_directory'\n```\n* For the reconstruction task, you need to newly set the parameters, *e.g,*, the target task, the weights of the losses and the number of the time sequence.\n```bash\npython Train.py --method recon --loss_compact 0.01 --loss_separate 0.01 --t_length 1 # for training\n```\n\n## Evaluation\n* Test your own model\n* Check your dataset_type (ped2, avenue or shanghai)\n```bash\npython Evaluate.py --dataset_type ped2 --model_dir your_model.pth --m_items_dir your_m_items.pt\n```\n* For the reconstruction task, you need to set the parameters as\n```bash\npython Evaluate.py --method recon --t_length 1 --alpha 0.7 --th 0.015 --dataset_type ped2 --model_dir your_model.pth --m_items_dir your_m_items.pt\n```\n* Test the model with our pre-trained model and memory items\n```bash\npython Evaluate.py --dataset_type ped2 --model_dir pretrained_model.pth --m_items_dir m_items.pt\n```\n\n## Pre-trained model and memory items\n\nWill be released soon.\n\n\n## Bibtex\n```\n@inproceedings{park2020learning,\n title={Learning Memory-guided Normality for Anomaly Detection},\n author={Park, Hyunjong and Noh, Jongyoun and Ham, Bumsub},\n booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},\n pages={14372--14381},\n year={2020}\n}\n```\n" }, { "path": "Train.py", "content": "import numpy as np\nimport os\nimport sys\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nimport torchvision\nimport torch.nn.init as init\nimport torch.utils.data as data\nimport torch.utils.data.dataset as dataset\nimport torchvision.datasets as dset\nimport torchvision.transforms as transforms\nfrom torch.autograd import Variable\nimport torchvision.utils as v_utils\nimport matplotlib.pyplot as plt\nimport cv2\nimport math\nfrom collections import OrderedDict\nimport copy\nimport time\nfrom model.utils import DataLoader\nfrom sklearn.metrics import roc_auc_score\nfrom utils import *\nimport random\n\nimport argparse\n\n\nparser = argparse.ArgumentParser(description=\"MNAD\")\nparser.add_argument('--gpus', nargs='+', type=str, help='gpus')\nparser.add_argument('--batch_size', type=int, default=4, help='batch size for training')\nparser.add_argument('--test_batch_size', type=int, default=1, help='batch size for test')\nparser.add_argument('--epochs', type=int, default=60, help='number of epochs for training')\nparser.add_argument('--loss_compact', type=float, default=0.1, help='weight of the feature compactness loss')\nparser.add_argument('--loss_separate', type=float, default=0.1, help='weight of the feature separateness loss')\nparser.add_argument('--h', type=int, default=256, help='height of input images')\nparser.add_argument('--w', type=int, default=256, help='width of input images')\nparser.add_argument('--c', type=int, default=3, help='channel of input images')\nparser.add_argument('--lr', type=float, default=2e-4, help='initial learning rate')\nparser.add_argument('--method', type=str, default='pred', help='The target task for anoamly detection')\nparser.add_argument('--t_length', type=int, default=5, help='length of the frame sequences')\nparser.add_argument('--fdim', type=int, default=512, help='channel dimension of the features')\nparser.add_argument('--mdim', type=int, default=512, help='channel dimension of the memory items')\nparser.add_argument('--msize', type=int, default=10, help='number of the memory items')\nparser.add_argument('--num_workers', type=int, default=2, help='number of workers for the train loader')\nparser.add_argument('--num_workers_test', type=int, default=1, help='number of workers for the test loader')\nparser.add_argument('--dataset_type', type=str, default='ped2', help='type of dataset: ped2, avenue, shanghai')\nparser.add_argument('--dataset_path', type=str, default='./dataset', help='directory of data')\nparser.add_argument('--exp_dir', type=str, default='log', help='directory of log')\n\nargs = parser.parse_args()\n\nos.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"\nif args.gpus is None:\n gpus = \"0\"\n os.environ[\"CUDA_VISIBLE_DEVICES\"]= gpus\nelse:\n gpus = \"\"\n for i in range(len(args.gpus)):\n gpus = gpus + args.gpus[i] + \",\"\n os.environ[\"CUDA_VISIBLE_DEVICES\"]= gpus[:-1]\n\ntorch.backends.cudnn.enabled = True # make sure to use cudnn for computational performance\n\ntrain_folder = args.dataset_path+\"/\"+args.dataset_type+\"/training/frames\"\ntest_folder = args.dataset_path+\"/\"+args.dataset_type+\"/testing/frames\"\n\n# Loading dataset\ntrain_dataset = DataLoader(train_folder, transforms.Compose([\n transforms.ToTensor(), \n ]), resize_height=args.h, resize_width=args.w, time_step=args.t_length-1)\n\ntest_dataset = DataLoader(test_folder, transforms.Compose([\n transforms.ToTensor(), \n ]), resize_height=args.h, resize_width=args.w, time_step=args.t_length-1)\n\ntrain_size = len(train_dataset)\ntest_size = len(test_dataset)\n\ntrain_batch = data.DataLoader(train_dataset, batch_size = args.batch_size, \n shuffle=True, num_workers=args.num_workers, drop_last=True)\ntest_batch = data.DataLoader(test_dataset, batch_size = args.test_batch_size, \n shuffle=False, num_workers=args.num_workers_test, drop_last=False)\n\n\n# Model setting\nassert args.method == 'pred' or args.method == 'recon', 'Wrong task name'\nif args.method == 'pred':\n from model.final_future_prediction_with_memory_spatial_sumonly_weight_ranking_top1 import *\n model = convAE(args.c, args.t_length, args.msize, args.fdim, args.mdim)\nelse:\n from model.Reconstruction import *\n model = convAE(args.c, memory_size = args.msize, feature_dim = args.fdim, key_dim = args.mdim)\nparams_encoder = list(model.encoder.parameters()) \nparams_decoder = list(model.decoder.parameters())\nparams = params_encoder + params_decoder\noptimizer = torch.optim.Adam(params, lr = args.lr)\nscheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer,T_max =args.epochs)\nmodel.cuda()\n\n\n# Report the training process\nlog_dir = os.path.join('./exp', args.dataset_type, args.method, args.exp_dir)\nif not os.path.exists(log_dir):\n os.makedirs(log_dir)\norig_stdout = sys.stdout\nf = open(os.path.join(log_dir, 'log.txt'),'w')\nsys.stdout= f\n\nloss_func_mse = nn.MSELoss(reduction='none')\n\n# Training\n\nm_items = F.normalize(torch.rand((args.msize, args.mdim), dtype=torch.float), dim=1).cuda() # Initialize the memory items\n\nfor epoch in range(args.epochs):\n labels_list = []\n model.train()\n \n start = time.time()\n for j,(imgs) in enumerate(train_batch):\n \n imgs = Variable(imgs).cuda()\n \n if args.method == 'pred':\n outputs, _, _, m_items, softmax_score_query, softmax_score_memory, separateness_loss, compactness_loss = model.forward(imgs[:,0:12], m_items, True)\n \n else:\n outputs, _, _, m_items, softmax_score_query, softmax_score_memory, separateness_loss, compactness_loss = model.forward(imgs, m_items, True)\n \n \n optimizer.zero_grad()\n if args.method == 'pred':\n loss_pixel = torch.mean(loss_func_mse(outputs, imgs[:,12:]))\n else:\n loss_pixel = torch.mean(loss_func_mse(outputs, imgs))\n \n loss = loss_pixel + args.loss_compact * compactness_loss + args.loss_separate * separateness_loss\n loss.backward(retain_graph=True)\n optimizer.step()\n \n scheduler.step()\n \n print('----------------------------------------')\n print('Epoch:', epoch+1)\n if args.method == 'pred':\n print('Loss: Prediction {:.6f}/ Compactness {:.6f}/ Separateness {:.6f}'.format(loss_pixel.item(), compactness_loss.item(), separateness_loss.item()))\n else:\n print('Loss: Reconstruction {:.6f}/ Compactness {:.6f}/ Separateness {:.6f}'.format(loss_pixel.item(), compactness_loss.item(), separateness_loss.item()))\n print('Memory_items:')\n print(m_items)\n print('----------------------------------------')\n \nprint('Training is finished')\n# Save the model and the memory items\ntorch.save(model, os.path.join(log_dir, 'model.pth'))\ntorch.save(m_items, os.path.join(log_dir, 'keys.pt'))\n \nsys.stdout = orig_stdout\nf.close()\n\n\n\n" }, { "path": "data/data_seqkey_all.py", "content": "import numpy as np\nimport os\nimport torch\nimport torch.utils.data as data\nimport torchvision.transforms as transforms\nfrom torch.utils.data import DataLoader\nfrom PIL import Image\nimport os.path\nimport sys\n\n\ndef pil_loader(path):\n # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)\n with open(path, 'rb') as f:\n img = Image.open(f)\n return img.convert('RGB')\n\n\ndef accimage_loader(path):\n import accimage\n try:\n return accimage.Image(path)\n except IOError:\n # Potentially a decoding problem, fall back to PIL.Image\n return pil_loader(path)\n\n\ndef default_loader(path):\n from torchvision import get_image_backend\n if get_image_backend() == 'accimage':\n return accimage_loader(path)\n else:\n return pil_loader(path)\n \ndef make_dataset(dir, class_to_idx):\n frames = []\n print(sorted(class_to_idx.keys()))\n dir = os.path.expanduser(dir)\n for target in sorted(class_to_idx.keys()):\n print(target)\n d = os.path.join(dir, target)\n if not os.path.isdir(d):\n continue\n# new_fnames = []\n \n for root, _, fnames in sorted(os.walk(d)):\n for fname in sorted(fnames):\n# fname = fname.split('.')[0]\n# seq = fname.split('_')[0][1:]\n# fname = fname.split('_')[1]\n# fname = fname.zfill(4)\n# new_fnames.append('V'+seq+'_'+fname+'.png')\n \n path = os.path.join(root, fname)\n frames.append(path)\n \n return frames\n\n\nclass DatasetFolder(data.Dataset):\n \n\n def __init__(self, root, loader=default_loader,transform=None, target_transform=None, length=5):\n classes, class_to_idx = self._find_classes(root)\n samples = make_dataset(root, class_to_idx)\n if len(samples) == 0:\n raise(RuntimeError(\"Found 0 files in subfolders of: \" + root))\n \n self.root = root\n self.loader = loader\n self.length = length\n# self.stride = np.random.choice(3,1) + 1\n self.classes = classes\n self.class_to_idx = class_to_idx\n# self.samples_gt = samples[self.length:]\n self.samples = samples[:-(self.length-1)]\n \n self.samples_all = samples\n self.samples_pool = samples[1:] \n# self.targets = [s[1] for s in samples]\n\n self.transform = transform\n self.target_transform = target_transform\n\n def _find_classes(self, dir):\n \"\"\"\n Finds the class folders in a dataset.\n Args:\n dir (string): Root directory path.\n Returns:\n tuple: (classes, class_to_idx) where classes are relative to (dir), and class_to_idx is a dictionary.\n Ensures:\n No class is a subdirectory of another.\n \"\"\"\n if sys.version_info >= (3, 5):\n # Faster and available in Python 3.5 and above\n classes = [d.name for d in os.scandir(dir) if d.is_dir()]\n else:\n classes = [d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))]\n classes.sort()\n class_to_idx = {classes[i]: i for i in range(len(classes))}\n return classes, class_to_idx\n\n def __getitem__(self, index):\n \"\"\"\n Args:\n index (int): Index\n Returns:\n tuple: (samples, gt(+length)) \n \n \"\"\"\n \n \n sample = [] \n \n path_start = self.samples[index]\n sample_start = self.loader(path_start)\n if self.transform is not None:\n sample_start = self.transform(sample_start)\n \n \n sample.append(sample_start) \n \n for i in range(self.length - 1):\n path = self.samples_all[index + (i+1)]\n sample_immediate = self.loader(path)\n if self.transform is not None:\n sample_immediate = self.transform(sample_immediate)\n \n sample.append(sample_immediate)\n \n \n# path_gt = self.samples_gt[index]\n# sample_gt = self.loader(path_gt)\n \n# if self.transform is not None:\n# sample_gt = self.transform(sample_gt)\n \n sample_input = sample[0]\n for i in range(self.length-1):\n sample_input = torch.cat((sample_input,sample[i+1]), dim=0)\n\n return sample_input\n \n def _stride(self):\n \n stride = int(np.random.choice(3,1) + 1)\n #if stride != 1:\n# self.samples_gt = self.samples_all[self.length*stride:]\n # self.samples = self.samples_all[:-(self.length*stride)]\n \n return stride\n\n def __len__(self):\n return len(self.samples)\n\n def __repr__(self):\n fmt_str = 'Dataset ' + self.__class__.__name__ + '\\n'\n fmt_str += ' Number of datapoints: {}\\n'.format(self.__len__())\n fmt_str += ' Root Location: {}\\n'.format(self.root)\n tmp = ' Transforms (if any): '\n fmt_str += '{0}{1}\\n'.format(tmp, self.transform.__repr__().replace('\\n', '\\n' + ' ' * len(tmp)))\n tmp = ' Target Transforms (if any): '\n fmt_str += '{0}{1}'.format(tmp, self.target_transform.__repr__().replace('\\n', '\\n' + ' ' * len(tmp)))\n return fmt_str\n\n\nIMG_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', 'webp']\n\n\n\n\n\nclass ImageFolder(DatasetFolder):\n \n \n \n def __init__(self, root, transform=None, target_transform=None,\n loader=default_loader, length=5):\n super(ImageFolder, self).__init__(root, loader,\n transform=transform,\n target_transform=target_transform)\n self.imgs = self.samples\n \n\n\n" }, { "path": "model/Memory.py", "content": "import torch\nimport torch.autograd as ag\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport numpy as np\nimport math\nimport functools\nimport random\nfrom torch.nn import functional as F\n\ndef random_uniform(shape, low, high, cuda):\n x = torch.rand(*shape)\n result_cpu = (high - low) * x + low\n if cuda:\n return result_cpu.cuda()\n else:\n return result_cpu\n \ndef distance(a, b):\n return torch.sqrt(((a - b) ** 2).sum()).unsqueeze(0)\n\ndef distance_batch(a, b):\n bs, _ = a.shape\n result = distance(a[0], b)\n for i in range(bs-1):\n result = torch.cat((result, distance(a[i], b)), 0)\n \n return result\n\ndef multiply(x): #to flatten matrix into a vector \n return functools.reduce(lambda x,y: x*y, x, 1)\n\ndef flatten(x):\n \"\"\" Flatten matrix into a vector \"\"\"\n count = multiply(x.size())\n return x.resize_(count)\n\ndef index(batch_size, x):\n idx = torch.arange(0, batch_size).long() \n idx = torch.unsqueeze(idx, -1)\n return torch.cat((idx, x), dim=1)\n\ndef MemoryLoss(memory):\n\n m, d = memory.size()\n memory_t = torch.t(memory)\n similarity = (torch.matmul(memory, memory_t))/2 + 1/2 # 30X30\n identity_mask = torch.eye(m).cuda()\n sim = torch.abs(similarity - identity_mask)\n \n return torch.sum(sim)/(m*(m-1))\n\n\nclass Memory(nn.Module):\n def __init__(self, memory_size, feature_dim, key_dim, temp_update, temp_gather):\n super(Memory, self).__init__()\n # Constants\n self.memory_size = memory_size\n self.feature_dim = feature_dim\n self.key_dim = key_dim\n self.temp_update = temp_update\n self.temp_gather = temp_gather\n \n def hard_neg_mem(self, mem, i):\n similarity = torch.matmul(mem,torch.t(self.keys_var))\n similarity[:,i] = -1\n _, max_idx = torch.topk(similarity, 1, dim=1)\n \n \n return self.keys_var[max_idx]\n \n def random_pick_memory(self, mem, max_indices):\n \n m, d = mem.size()\n output = []\n for i in range(m):\n flattened_indices = (max_indices==i).nonzero()\n a, _ = flattened_indices.size()\n if a != 0:\n number = np.random.choice(a, 1)\n output.append(flattened_indices[number, 0])\n else:\n output.append(-1)\n \n return torch.tensor(output)\n \n def get_update_query(self, mem, max_indices, update_indices, score, query, train):\n \n m, d = mem.size()\n if train:\n query_update = torch.zeros((m,d)).cuda()\n random_update = torch.zeros((m,d)).cuda()\n for i in range(m):\n idx = torch.nonzero(max_indices.squeeze(1)==i)\n a, _ = idx.size()\n #ex = update_indices[0][i]\n if a != 0:\n #random_idx = torch.randperm(a)[0]\n #idx = idx[idx != ex]\n# query_update[i] = torch.sum(query[idx].squeeze(1), dim=0)\n query_update[i] = torch.sum(((score[idx,i] / torch.max(score[:,i])) *query[idx].squeeze(1)), dim=0)\n #random_update[i] = query[random_idx] * (score[random_idx,i] / torch.max(score[:,i]))\n else:\n query_update[i] = 0 \n #random_update[i] = 0\n \n \n return query_update \n \n else:\n query_update = torch.zeros((m,d)).cuda()\n for i in range(m):\n idx = torch.nonzero(max_indices.squeeze(1)==i)\n a, _ = idx.size()\n #ex = update_indices[0][i]\n if a != 0:\n #idx = idx[idx != ex]\n query_update[i] = torch.sum(((score[idx,i] / torch.max(score[:,i])) *query[idx].squeeze(1)), dim=0)\n# query_update[i] = torch.sum(query[idx].squeeze(1), dim=0)\n else:\n query_update[i] = 0 \n \n return query_update\n\n def get_score(self, mem, query):\n bs, h,w,d = query.size()\n m, d = mem.size()\n \n score = torch.matmul(query, torch.t(mem))# b X h X w X m\n score = score.view(bs*h*w, m)# (b X h X w) X m\n \n score_query = F.softmax(score, dim=0)\n score_memory = F.softmax(score,dim=1)\n \n return score_query, score_memory\n \n def forward(self, query, keys, train=True):\n\n batch_size, dims,h,w = query.size() # b X d X h X w\n query = F.normalize(query, dim=1)\n query = query.permute(0,2,3,1) # b X h X w X d\n \n #train\n if train:\n #gathering loss\n gathering_loss = self.gather_loss(query,keys, train)\n #spreading_loss\n spreading_loss = self.spread_loss(query, keys, train)\n # read\n updated_query, softmax_score_query,softmax_score_memory = self.read(query, keys)\n #update\n updated_memory = self.update(query, keys, train)\n \n return updated_query, updated_memory, softmax_score_query, softmax_score_memory, gathering_loss, spreading_loss\n \n #test\n else:\n #gathering loss\n gathering_loss = self.gather_loss(query,keys, train)\n \n # read\n updated_query, softmax_score_query,softmax_score_memory = self.read(query, keys)\n \n #update\n updated_memory = keys\n \n \n return updated_query, updated_memory, softmax_score_query, softmax_score_memory, gathering_loss\n \n \n \n def update(self, query, keys,train):\n \n batch_size, h,w,dims = query.size() # b X h X w X d \n \n softmax_score_query, softmax_score_memory = self.get_score(keys, query)\n \n query_reshape = query.contiguous().view(batch_size*h*w, dims)\n \n _, gathering_indices = torch.topk(softmax_score_memory, 1, dim=1)\n _, updating_indices = torch.topk(softmax_score_query, 1, dim=0)\n \n if train:\n # top-1 queries (of each memory) update (weighted sum) & random pick \n query_update = self.get_update_query(keys, gathering_indices, updating_indices, softmax_score_query, query_reshape,train)\n updated_memory = F.normalize(query_update + keys, dim=1)\n \n else:\n # only weighted sum update when test \n query_update = self.get_update_query(keys, gathering_indices, updating_indices, softmax_score_query, query_reshape, train)\n updated_memory = F.normalize(query_update + keys, dim=1)\n \n # top-1 update\n #query_update = query_reshape[updating_indices][0]\n #updated_memory = F.normalize(query_update + keys, dim=1)\n \n return updated_memory.detach()\n \n \n def pointwise_gather_loss(self, query_reshape, keys, gathering_indices, train):\n n,dims = query_reshape.size() # (b X h X w) X d\n loss_mse = torch.nn.MSELoss(reduction='none')\n \n pointwise_loss = loss_mse(query_reshape, keys[gathering_indices].squeeze(1).detach())\n \n return pointwise_loss\n \n def spread_loss(self,query, keys, train):\n batch_size, h,w,dims = query.size() # b X h X w X d\n\n loss = torch.nn.TripletMarginLoss(margin=1.0)\n\n softmax_score_query, softmax_score_memory = self.get_score(keys, query)\n\n query_reshape = query.contiguous().view(batch_size*h*w, dims)\n\n _, gathering_indices = torch.topk(softmax_score_memory, 2, dim=1)\n\n #1st, 2nd closest memories\n pos = keys[gathering_indices[:,0]]\n neg = keys[gathering_indices[:,1]]\n\n spreading_loss = loss(query_reshape,pos.detach(), neg.detach())\n\n return spreading_loss\n \n def gather_loss(self, query, keys, train):\n \n batch_size, h,w,dims = query.size() # b X h X w X d\n\n loss_mse = torch.nn.MSELoss()\n\n softmax_score_query, softmax_score_memory = self.get_score(keys, query)\n\n query_reshape = query.contiguous().view(batch_size*h*w, dims)\n\n _, gathering_indices = torch.topk(softmax_score_memory, 1, dim=1)\n\n gathering_loss = loss_mse(query_reshape, keys[gathering_indices].squeeze(1).detach())\n\n return gathering_loss\n \n \n \n \n def read(self, query, updated_memory):\n batch_size, h,w,dims = query.size() # b X h X w X d\n\n softmax_score_query, softmax_score_memory = self.get_score(updated_memory, query)\n\n query_reshape = query.contiguous().view(batch_size*h*w, dims)\n \n concat_memory = torch.matmul(softmax_score_memory.detach(), updated_memory) # (b X h X w) X d\n updated_query = torch.cat((query_reshape, concat_memory), dim = 1) # (b X h X w) X 2d\n updated_query = updated_query.view(batch_size, h, w, 2*dims)\n updated_query = updated_query.permute(0,3,1,2)\n \n return updated_query, softmax_score_query, softmax_score_memory\n \n " }, { "path": "model/Reconstruction.py", "content": "import numpy as np\nimport os\nimport sys\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom .Memory import *\n\nclass Encoder(torch.nn.Module):\n def __init__(self, t_length = 2, n_channel =3):\n super(Encoder, self).__init__()\n \n def Basic(intInput, intOutput):\n return torch.nn.Sequential(\n torch.nn.Conv2d(in_channels=intInput, out_channels=intOutput, kernel_size=3, stride=1, padding=1),\n torch.nn.BatchNorm2d(intOutput),\n torch.nn.ReLU(inplace=False),\n torch.nn.Conv2d(in_channels=intOutput, out_channels=intOutput, kernel_size=3, stride=1, padding=1),\n torch.nn.BatchNorm2d(intOutput),\n torch.nn.ReLU(inplace=False)\n )\n \n def Basic_(intInput, intOutput):\n return torch.nn.Sequential(\n torch.nn.Conv2d(in_channels=intInput, out_channels=intOutput, kernel_size=3, stride=1, padding=1),\n torch.nn.BatchNorm2d(intOutput),\n torch.nn.ReLU(inplace=False),\n torch.nn.Conv2d(in_channels=intOutput, out_channels=intOutput, kernel_size=3, stride=1, padding=1),\n )\n \n self.moduleConv1 = Basic(n_channel*(t_length-1), 64)\n self.modulePool1 = torch.nn.MaxPool2d(kernel_size=2, stride=2)\n\n self.moduleConv2 = Basic(64, 128)\n self.modulePool2 = torch.nn.MaxPool2d(kernel_size=2, stride=2)\n \n self.moduleConv3 = Basic(128, 256)\n self.modulePool3 = torch.nn.MaxPool2d(kernel_size=2, stride=2)\n\n self.moduleConv4 = Basic_(256, 512)\n self.moduleBatchNorm = torch.nn.BatchNorm2d(512)\n self.moduleReLU = torch.nn.ReLU(inplace=False)\n \n def forward(self, x):\n\n tensorConv1 = self.moduleConv1(x)\n tensorPool1 = self.modulePool1(tensorConv1)\n\n tensorConv2 = self.moduleConv2(tensorPool1)\n tensorPool2 = self.modulePool2(tensorConv2)\n\n tensorConv3 = self.moduleConv3(tensorPool2)\n tensorPool3 = self.modulePool3(tensorConv3)\n\n tensorConv4 = self.moduleConv4(tensorPool3)\n \n return tensorConv4\n\n \n \nclass Decoder(torch.nn.Module):\n def __init__(self, t_length = 2, n_channel =3):\n super(Decoder, self).__init__()\n \n def Basic(intInput, intOutput):\n return torch.nn.Sequential(\n torch.nn.Conv2d(in_channels=intInput, out_channels=intOutput, kernel_size=3, stride=1, padding=1),\n torch.nn.BatchNorm2d(intOutput),\n torch.nn.ReLU(inplace=False),\n torch.nn.Conv2d(in_channels=intOutput, out_channels=intOutput, kernel_size=3, stride=1, padding=1),\n torch.nn.BatchNorm2d(intOutput),\n torch.nn.ReLU(inplace=False)\n )\n \n \n def Gen(intInput, intOutput, nc):\n return torch.nn.Sequential(\n torch.nn.Conv2d(in_channels=intInput, out_channels=nc, kernel_size=3, stride=1, padding=1),\n torch.nn.BatchNorm2d(nc),\n torch.nn.ReLU(inplace=False),\n torch.nn.Conv2d(in_channels=nc, out_channels=nc, kernel_size=3, stride=1, padding=1),\n torch.nn.BatchNorm2d(nc),\n torch.nn.ReLU(inplace=False),\n torch.nn.Conv2d(in_channels=nc, out_channels=intOutput, kernel_size=3, stride=1, padding=1),\n torch.nn.Tanh()\n )\n \n def Upsample(nc, intOutput):\n return torch.nn.Sequential(\n torch.nn.ConvTranspose2d(in_channels = nc, out_channels=intOutput, kernel_size = 3, stride = 2, padding = 1, output_padding = 1),\n torch.nn.BatchNorm2d(intOutput),\n torch.nn.ReLU(inplace=False)\n )\n \n self.moduleConv = Basic(1024, 512)\n self.moduleUpsample4 = Upsample(512, 512)\n\n self.moduleDeconv3 = Basic(512, 256)\n self.moduleUpsample3 = Upsample(256, 256)\n\n self.moduleDeconv2 = Basic(256, 128)\n self.moduleUpsample2 = Upsample(128, 128)\n\n self.moduleDeconv1 = Gen(128,n_channel,64)\n \n \n \n def forward(self, x):\n \n tensorConv = self.moduleConv(x)\n\n tensorUpsample4 = self.moduleUpsample4(tensorConv)\n \n tensorDeconv3 = self.moduleDeconv3(tensorUpsample4)\n tensorUpsample3 = self.moduleUpsample3(tensorDeconv3)\n \n tensorDeconv2 = self.moduleDeconv2(tensorUpsample3)\n tensorUpsample2 = self.moduleUpsample2(tensorDeconv2)\n \n output = self.moduleDeconv1(tensorUpsample2)\n\n \n return output\n \n\n\nclass convAE(torch.nn.Module):\n def __init__(self, n_channel =3, t_length = 2, memory_size = 10, feature_dim = 512, key_dim = 512, temp_update = 0.1, temp_gather=0.1):\n super(convAE, self).__init__()\n\n self.encoder = Encoder(t_length, n_channel)\n self.decoder = Decoder(t_length, n_channel)\n self.memory = Memory(memory_size,feature_dim, key_dim, temp_update, temp_gather)\n \n\n def forward(self, x, keys,train=True):\n\n fea = self.encoder(x)\n if train:\n updated_fea, keys, softmax_score_query, softmax_score_memory, gathering_loss, spreading_loss = self.memory(fea, keys, train)\n output = self.decoder(updated_fea)\n \n return output, fea, updated_fea, keys, softmax_score_query, softmax_score_memory, gathering_loss, spreading_loss\n \n #test\n else:\n updated_fea, keys, softmax_score_query, softmax_score_memory, gathering_loss = self.memory(fea, keys, train)\n output = self.decoder(updated_fea)\n \n return output, fea, updated_fea, keys, softmax_score_query, softmax_score_memory, gathering_loss\n \n \n\n\n\n " }, { "path": "model/final_future_prediction_with_memory_spatial_sumonly_weight_ranking_top1.py", "content": "import numpy as np\nimport os\nimport sys\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom .memory_final_spatial_sumonly_weight_ranking_top1 import *\n\nclass Encoder(torch.nn.Module):\n def __init__(self, t_length = 5, n_channel =3):\n super(Encoder, self).__init__()\n \n def Basic(intInput, intOutput):\n return torch.nn.Sequential(\n torch.nn.Conv2d(in_channels=intInput, out_channels=intOutput, kernel_size=3, stride=1, padding=1),\n torch.nn.BatchNorm2d(intOutput),\n torch.nn.ReLU(inplace=False),\n torch.nn.Conv2d(in_channels=intOutput, out_channels=intOutput, kernel_size=3, stride=1, padding=1),\n torch.nn.BatchNorm2d(intOutput),\n torch.nn.ReLU(inplace=False)\n )\n \n def Basic_(intInput, intOutput):\n return torch.nn.Sequential(\n torch.nn.Conv2d(in_channels=intInput, out_channels=intOutput, kernel_size=3, stride=1, padding=1),\n torch.nn.BatchNorm2d(intOutput),\n torch.nn.ReLU(inplace=False),\n torch.nn.Conv2d(in_channels=intOutput, out_channels=intOutput, kernel_size=3, stride=1, padding=1),\n )\n \n self.moduleConv1 = Basic(n_channel*(t_length-1), 64)\n self.modulePool1 = torch.nn.MaxPool2d(kernel_size=2, stride=2)\n\n self.moduleConv2 = Basic(64, 128)\n self.modulePool2 = torch.nn.MaxPool2d(kernel_size=2, stride=2)\n \n self.moduleConv3 = Basic(128, 256)\n self.modulePool3 = torch.nn.MaxPool2d(kernel_size=2, stride=2)\n\n self.moduleConv4 = Basic_(256, 512)\n self.moduleBatchNorm = torch.nn.BatchNorm2d(512)\n self.moduleReLU = torch.nn.ReLU(inplace=False)\n \n def forward(self, x):\n\n tensorConv1 = self.moduleConv1(x)\n tensorPool1 = self.modulePool1(tensorConv1)\n\n tensorConv2 = self.moduleConv2(tensorPool1)\n tensorPool2 = self.modulePool2(tensorConv2)\n\n tensorConv3 = self.moduleConv3(tensorPool2)\n tensorPool3 = self.modulePool3(tensorConv3)\n\n tensorConv4 = self.moduleConv4(tensorPool3)\n \n return tensorConv4, tensorConv1, tensorConv2, tensorConv3\n\n \n \nclass Decoder(torch.nn.Module):\n def __init__(self, t_length = 5, n_channel =3):\n super(Decoder, self).__init__()\n \n def Basic(intInput, intOutput):\n return torch.nn.Sequential(\n torch.nn.Conv2d(in_channels=intInput, out_channels=intOutput, kernel_size=3, stride=1, padding=1),\n torch.nn.BatchNorm2d(intOutput),\n torch.nn.ReLU(inplace=False),\n torch.nn.Conv2d(in_channels=intOutput, out_channels=intOutput, kernel_size=3, stride=1, padding=1),\n torch.nn.BatchNorm2d(intOutput),\n torch.nn.ReLU(inplace=False)\n )\n \n \n def Gen(intInput, intOutput, nc):\n return torch.nn.Sequential(\n torch.nn.Conv2d(in_channels=intInput, out_channels=nc, kernel_size=3, stride=1, padding=1),\n torch.nn.BatchNorm2d(nc),\n torch.nn.ReLU(inplace=False),\n torch.nn.Conv2d(in_channels=nc, out_channels=nc, kernel_size=3, stride=1, padding=1),\n torch.nn.BatchNorm2d(nc),\n torch.nn.ReLU(inplace=False),\n torch.nn.Conv2d(in_channels=nc, out_channels=intOutput, kernel_size=3, stride=1, padding=1),\n torch.nn.Tanh()\n )\n \n def Upsample(nc, intOutput):\n return torch.nn.Sequential(\n torch.nn.ConvTranspose2d(in_channels = nc, out_channels=intOutput, kernel_size = 3, stride = 2, padding = 1, output_padding = 1),\n torch.nn.BatchNorm2d(intOutput),\n torch.nn.ReLU(inplace=False)\n )\n \n self.moduleConv = Basic(1024, 512)\n self.moduleUpsample4 = Upsample(512, 256)\n\n self.moduleDeconv3 = Basic(512, 256)\n self.moduleUpsample3 = Upsample(256, 128)\n\n self.moduleDeconv2 = Basic(256, 128)\n self.moduleUpsample2 = Upsample(128, 64)\n\n self.moduleDeconv1 = Gen(128,n_channel,64)\n \n \n \n def forward(self, x, skip1, skip2, skip3):\n \n tensorConv = self.moduleConv(x)\n\n tensorUpsample4 = self.moduleUpsample4(tensorConv)\n cat4 = torch.cat((skip3, tensorUpsample4), dim = 1)\n \n tensorDeconv3 = self.moduleDeconv3(cat4)\n tensorUpsample3 = self.moduleUpsample3(tensorDeconv3)\n cat3 = torch.cat((skip2, tensorUpsample3), dim = 1)\n \n tensorDeconv2 = self.moduleDeconv2(cat3)\n tensorUpsample2 = self.moduleUpsample2(tensorDeconv2)\n cat2 = torch.cat((skip1, tensorUpsample2), dim = 1)\n \n output = self.moduleDeconv1(cat2)\n\n \n return output\n \n\n\nclass convAE(torch.nn.Module):\n def __init__(self, n_channel =3, t_length = 5, memory_size = 10, feature_dim = 512, key_dim = 512, temp_update = 0.1, temp_gather=0.1):\n super(convAE, self).__init__()\n\n self.encoder = Encoder(t_length, n_channel)\n self.decoder = Decoder(t_length, n_channel)\n self.memory = Memory(memory_size,feature_dim, key_dim, temp_update, temp_gather)\n \n\n def forward(self, x, keys,train=True):\n\n fea, skip1, skip2, skip3 = self.encoder(x)\n if train:\n updated_fea, keys, softmax_score_query, softmax_score_memory, separateness_loss, compactness_loss = self.memory(fea, keys, train)\n output = self.decoder(updated_fea, skip1, skip2, skip3)\n \n return output, fea, updated_fea, keys, softmax_score_query, softmax_score_memory, separateness_loss, compactness_loss\n \n #test\n else:\n updated_fea, keys, softmax_score_query, softmax_score_memory,query, top1_keys, keys_ind, compactness_loss = self.memory(fea, keys, train)\n output = self.decoder(updated_fea, skip1, skip2, skip3)\n \n return output, fea, updated_fea, keys, softmax_score_query, softmax_score_memory, query, top1_keys, keys_ind, compactness_loss\n \n \n\n\n\n \n \n" }, { "path": "model/memory_final_spatial_sumonly_weight_ranking_top1.py", "content": "import torch\nimport torch.autograd as ag\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport numpy as np\nimport math\nimport functools\nimport random\nfrom torch.nn import functional as F\n\ndef random_uniform(shape, low, high, cuda):\n x = torch.rand(*shape)\n result_cpu = (high - low) * x + low\n if cuda:\n return result_cpu.cuda()\n else:\n return result_cpu\n \ndef distance(a, b):\n return torch.sqrt(((a - b) ** 2).sum()).unsqueeze(0)\n\ndef distance_batch(a, b):\n bs, _ = a.shape\n result = distance(a[0], b)\n for i in range(bs-1):\n result = torch.cat((result, distance(a[i], b)), 0)\n \n return result\n\ndef multiply(x): #to flatten matrix into a vector \n return functools.reduce(lambda x,y: x*y, x, 1)\n\ndef flatten(x):\n \"\"\" Flatten matrix into a vector \"\"\"\n count = multiply(x.size())\n return x.resize_(count)\n\ndef index(batch_size, x):\n idx = torch.arange(0, batch_size).long() \n idx = torch.unsqueeze(idx, -1)\n return torch.cat((idx, x), dim=1)\n\ndef MemoryLoss(memory):\n\n m, d = memory.size()\n memory_t = torch.t(memory)\n similarity = (torch.matmul(memory, memory_t))/2 + 1/2 # 30X30\n identity_mask = torch.eye(m).cuda()\n sim = torch.abs(similarity - identity_mask)\n \n return torch.sum(sim)/(m*(m-1))\n\n\nclass Memory(nn.Module):\n def __init__(self, memory_size, feature_dim, key_dim, temp_update, temp_gather):\n super(Memory, self).__init__()\n # Constants\n self.memory_size = memory_size\n self.feature_dim = feature_dim\n self.key_dim = key_dim\n self.temp_update = temp_update\n self.temp_gather = temp_gather\n \n def hard_neg_mem(self, mem, i):\n similarity = torch.matmul(mem,torch.t(self.keys_var))\n similarity[:,i] = -1\n _, max_idx = torch.topk(similarity, 1, dim=1)\n \n \n return self.keys_var[max_idx]\n \n def random_pick_memory(self, mem, max_indices):\n \n m, d = mem.size()\n output = []\n for i in range(m):\n flattened_indices = (max_indices==i).nonzero()\n a, _ = flattened_indices.size()\n if a != 0:\n number = np.random.choice(a, 1)\n output.append(flattened_indices[number, 0])\n else:\n output.append(-1)\n \n return torch.tensor(output)\n \n def get_update_query(self, mem, max_indices, update_indices, score, query, train):\n \n m, d = mem.size()\n if train:\n query_update = torch.zeros((m,d)).cuda()\n # random_update = torch.zeros((m,d)).cuda()\n for i in range(m):\n idx = torch.nonzero(max_indices.squeeze(1)==i)\n a, _ = idx.size()\n if a != 0:\n query_update[i] = torch.sum(((score[idx,i] / torch.max(score[:,i])) *query[idx].squeeze(1)), dim=0)\n else:\n query_update[i] = 0 \n \n \n return query_update \n \n else:\n query_update = torch.zeros((m,d)).cuda()\n for i in range(m):\n idx = torch.nonzero(max_indices.squeeze(1)==i)\n a, _ = idx.size()\n if a != 0:\n query_update[i] = torch.sum(((score[idx,i] / torch.max(score[:,i])) *query[idx].squeeze(1)), dim=0)\n else:\n query_update[i] = 0 \n \n return query_update\n\n def get_score(self, mem, query):\n bs, h,w,d = query.size()\n m, d = mem.size()\n \n score = torch.matmul(query, torch.t(mem))# b X h X w X m\n score = score.view(bs*h*w, m)# (b X h X w) X m\n \n score_query = F.softmax(score, dim=0)\n score_memory = F.softmax(score,dim=1)\n \n return score_query, score_memory\n \n def forward(self, query, keys, train=True):\n\n batch_size, dims,h,w = query.size() # b X d X h X w\n query = F.normalize(query, dim=1)\n query = query.permute(0,2,3,1) # b X h X w X d\n \n #train\n if train:\n #losses\n separateness_loss, compactness_loss = self.gather_loss(query,keys, train)\n # read\n updated_query, softmax_score_query,softmax_score_memory = self.read(query, keys)\n #update\n updated_memory = self.update(query, keys, train)\n \n return updated_query, updated_memory, softmax_score_query, softmax_score_memory, separateness_loss, compactness_loss\n \n #test\n else:\n # loss\n compactness_loss, query_re, top1_keys, keys_ind = self.gather_loss(query,keys, train)\n \n # read\n updated_query, softmax_score_query,softmax_score_memory = self.read(query, keys)\n \n #update\n updated_memory = keys\n \n \n return updated_query, updated_memory, softmax_score_query, softmax_score_memory, query_re, top1_keys,keys_ind, compactness_loss\n \n \n \n def update(self, query, keys,train):\n \n batch_size, h,w,dims = query.size() # b X h X w X d \n \n softmax_score_query, softmax_score_memory = self.get_score(keys, query)\n \n query_reshape = query.contiguous().view(batch_size*h*w, dims)\n \n _, gathering_indices = torch.topk(softmax_score_memory, 1, dim=1)\n _, updating_indices = torch.topk(softmax_score_query, 1, dim=0)\n \n if train:\n \n query_update = self.get_update_query(keys, gathering_indices, updating_indices, softmax_score_query, query_reshape,train)\n updated_memory = F.normalize(query_update + keys, dim=1)\n \n else:\n query_update = self.get_update_query(keys, gathering_indices, updating_indices, softmax_score_query, query_reshape, train)\n updated_memory = F.normalize(query_update + keys, dim=1)\n \n return updated_memory.detach()\n \n \n def pointwise_gather_loss(self, query_reshape, keys, gathering_indices, train):\n n,dims = query_reshape.size() # (b X h X w) X d\n loss_mse = torch.nn.MSELoss(reduction='none')\n \n pointwise_loss = loss_mse(query_reshape, keys[gathering_indices].squeeze(1).detach())\n \n return pointwise_loss\n \n def gather_loss(self,query, keys, train):\n batch_size, h,w,dims = query.size() # b X h X w X d\n if train:\n loss = torch.nn.TripletMarginLoss(margin=1.0)\n loss_mse = torch.nn.MSELoss()\n softmax_score_query, softmax_score_memory = self.get_score(keys, query)\n \n query_reshape = query.contiguous().view(batch_size*h*w, dims)\n \n _, gathering_indices = torch.topk(softmax_score_memory, 2, dim=1)\n \n #1st, 2nd closest memories\n pos = keys[gathering_indices[:,0]]\n neg = keys[gathering_indices[:,1]]\n top1_loss = loss_mse(query_reshape, pos.detach())\n gathering_loss = loss(query_reshape,pos.detach(), neg.detach())\n \n return gathering_loss, top1_loss\n \n \n else:\n loss_mse = torch.nn.MSELoss()\n \n softmax_score_query, softmax_score_memory = self.get_score(keys, query)\n \n query_reshape = query.contiguous().view(batch_size*h*w, dims)\n \n _, gathering_indices = torch.topk(softmax_score_memory, 1, dim=1)\n \n gathering_loss = loss_mse(query_reshape, keys[gathering_indices].squeeze(1).detach())\n \n return gathering_loss, query_reshape, keys[gathering_indices].squeeze(1).detach(), gathering_indices[:,0]\n \n \n \n \n def read(self, query, updated_memory):\n batch_size, h,w,dims = query.size() # b X h X w X d\n\n softmax_score_query, softmax_score_memory = self.get_score(updated_memory, query)\n\n query_reshape = query.contiguous().view(batch_size*h*w, dims)\n \n concat_memory = torch.matmul(softmax_score_memory.detach(), updated_memory) # (b X h X w) X d\n updated_query = torch.cat((query_reshape, concat_memory), dim = 1) # (b X h X w) X 2d\n updated_query = updated_query.view(batch_size, h, w, 2*dims)\n updated_query = updated_query.permute(0,3,1,2)\n \n return updated_query, softmax_score_query, softmax_score_memory\n \n \n" }, { "path": "model/utils.py", "content": "import numpy as np\nfrom collections import OrderedDict\nimport os\nimport glob\nimport cv2\nimport torch.utils.data as data\n\n\nrng = np.random.RandomState(2020)\n\ndef np_load_frame(filename, resize_height, resize_width):\n \"\"\"\n Load image path and convert it to numpy.ndarray. Notes that the color channels are BGR and the color space\n is normalized from [0, 255] to [-1, 1].\n\n :param filename: the full path of image\n :param resize_height: resized height\n :param resize_width: resized width\n :return: numpy.ndarray\n \"\"\"\n image_decoded = cv2.imread(filename)\n image_resized = cv2.resize(image_decoded, (resize_width, resize_height))\n image_resized = image_resized.astype(dtype=np.float32)\n image_resized = (image_resized / 127.5) - 1.0\n return image_resized\n\n\n\n\nclass DataLoader(data.Dataset):\n def __init__(self, video_folder, transform, resize_height, resize_width, time_step=4, num_pred=1):\n self.dir = video_folder\n self.transform = transform\n self.videos = OrderedDict()\n self._resize_height = resize_height\n self._resize_width = resize_width\n self._time_step = time_step\n self._num_pred = num_pred\n self.setup()\n self.samples = self.get_all_samples()\n \n \n def setup(self):\n videos = glob.glob(os.path.join(self.dir, '*'))\n for video in sorted(videos):\n video_name = video.split('/')[-1]\n self.videos[video_name] = {}\n self.videos[video_name]['path'] = video\n self.videos[video_name]['frame'] = glob.glob(os.path.join(video, '*.jpg'))\n self.videos[video_name]['frame'].sort()\n self.videos[video_name]['length'] = len(self.videos[video_name]['frame'])\n \n \n def get_all_samples(self):\n frames = []\n videos = glob.glob(os.path.join(self.dir, '*'))\n for video in sorted(videos):\n video_name = video.split('/')[-1]\n for i in range(len(self.videos[video_name]['frame'])-self._time_step):\n frames.append(self.videos[video_name]['frame'][i])\n \n return frames \n \n \n def __getitem__(self, index):\n video_name = self.samples[index].split('/')[-2]\n frame_name = int(self.samples[index].split('/')[-1].split('.')[-2])\n \n batch = []\n for i in range(self._time_step+self._num_pred):\n image = np_load_frame(self.videos[video_name]['frame'][frame_name+i], self._resize_height, self._resize_width)\n if self.transform is not None:\n batch.append(self.transform(image))\n\n return np.concatenate(batch, axis=0)\n \n \n def __len__(self):\n return len(self.samples)\n" }, { "path": "utils.py", "content": "import numpy as np\nimport os\nimport sys\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision\nimport torchvision.utils as v_utils\nimport matplotlib.pyplot as plt\nimport cv2\nimport math\nfrom collections import OrderedDict\nimport copy\nimport time\nfrom sklearn.metrics import roc_auc_score\n\ndef rmse(predictions, targets):\n return np.sqrt(((predictions - targets) ** 2).mean())\n\ndef psnr(mse):\n\n return 10 * math.log10(1 / mse)\n\ndef get_lr(optimizer):\n for param_group in optimizer.param_groups:\n return param_group['lr']\n\n\ndef normalize_img(img):\n\n img_re = copy.copy(img)\n \n img_re = (img_re - np.min(img_re)) / (np.max(img_re) - np.min(img_re))\n \n return img_re\n\ndef point_score(outputs, imgs):\n \n loss_func_mse = nn.MSELoss(reduction='none')\n error = loss_func_mse((outputs[0]+1)/2,(imgs[0]+1)/2)\n normal = (1-torch.exp(-error))\n score = (torch.sum(normal*loss_func_mse((outputs[0]+1)/2,(imgs[0]+1)/2)) / torch.sum(normal)).item()\n return score\n \ndef anomaly_score(psnr, max_psnr, min_psnr):\n return ((psnr - min_psnr) / (max_psnr-min_psnr))\n\ndef anomaly_score_inv(psnr, max_psnr, min_psnr):\n return (1.0 - ((psnr - min_psnr) / (max_psnr-min_psnr)))\n\ndef anomaly_score_list(psnr_list):\n anomaly_score_list = list()\n for i in range(len(psnr_list)):\n anomaly_score_list.append(anomaly_score(psnr_list[i], np.max(psnr_list), np.min(psnr_list)))\n \n return anomaly_score_list\n\ndef anomaly_score_list_inv(psnr_list):\n anomaly_score_list = list()\n for i in range(len(psnr_list)):\n anomaly_score_list.append(anomaly_score_inv(psnr_list[i], np.max(psnr_list), np.min(psnr_list)))\n \n return anomaly_score_list\n\ndef AUC(anomal_scores, labels):\n frame_auc = roc_auc_score(y_true=np.squeeze(labels, axis=0), y_score=np.squeeze(anomal_scores))\n return frame_auc\n\ndef score_sum(list1, list2, alpha):\n list_result = []\n for i in range(len(list1)):\n list_result.append((alpha*list1[i]+(1-alpha)*list2[i]))\n \n return list_result\n" } ]