├── .gitignore ├── LICENSE ├── image └── align_sample │ ├── 1.jpg │ ├── 2.jpg │ └── 3.jpg ├── readme.md └── src ├── CFP ├── ResNet.py ├── eval_cfp.py ├── eval_cfp.sh ├── eval_roc.py └── selfDefine.py ├── IJBA ├── ResNet.py ├── eval_ijba.py ├── eval_ijba.sh ├── selfDefine.py ├── test_recog.py └── test_verify.py ├── end2end ├── ResNet.py ├── main.py ├── selfDefine.py ├── slurm.py └── train.sh ├── preprocess ├── align_cfp.sh ├── align_img_list.txt ├── cfp_alignment.txt ├── estimate_pose.txt ├── pnp.txt ├── pre_cfp_data.py └── test_process_align └── stitching ├── branch_test.py ├── branch_train.py ├── branch_util.py └── train_stitch.sh /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | *.swp 3 | *.zip 4 | 5 | pick_data.py 6 | eval_res.txt 7 | 8 | __pycache__/ 9 | data/ 10 | msceleb/ 11 | model/ 12 | log/ 13 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | BSD 2-Clause License 2 | 3 | Copyright (c) 2018 Yu Rong 4 | All rights reserved. 5 | 6 | Redistribution and use in source and binary forms, with or without 7 | modification, are permitted provided that the following conditions are met: 8 | 9 | * Redistributions of source code must retain the above copyright notice, this 10 | list of conditions and the following disclaimer. 11 | 12 | * Redistributions in binary form must reproduce the above copyright notice, 13 | this list of conditions and the following disclaimer in the documentation 14 | and/or other materials provided with the distribution. 15 | 16 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 17 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 18 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 19 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 20 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 21 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 22 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 23 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 24 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 25 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 26 | -------------------------------------------------------------------------------- /image/align_sample/1.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/penincillin/DREAM/7bea01e5b108546a2b3a809ad17cf478607c9db5/image/align_sample/1.jpg -------------------------------------------------------------------------------- /image/align_sample/2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/penincillin/DREAM/7bea01e5b108546a2b3a809ad17cf478607c9db5/image/align_sample/2.jpg -------------------------------------------------------------------------------- /image/align_sample/3.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/penincillin/DREAM/7bea01e5b108546a2b3a809ad17cf478607c9db5/image/align_sample/3.jpg -------------------------------------------------------------------------------- /readme.md: -------------------------------------------------------------------------------- 1 | # DREAM block for Pose-Robust Face Recognition 2 | This is our implementation for our CVPR 2018 accepted paper *Pose-Robust Face Recognition via Deep Residual Equivariant Mapping* [paper on arxiv](https://arxiv.org/abs/1803.00839). 3 | 4 | The code is wriiten by [Yu Rong](https://github.com/penincillin) and [Kaidi Cao](https://github.com/CarlyleCao) 5 | 6 | ## Prerequisites 7 | - Linux or macOS 8 | - Python 3 9 | - NVIDIA GPU + CUDA CuDNN or CPU (GPU is prefered) 10 | - opencv2.4 (opencv 2.4.13 is preferred) 11 | 12 | ## Getting Started 13 | ### Installation 14 | - Install Anaconda 15 | - [Anaconda3-4.2.0](https://repo.continuum.io/archive/Anaconda3-4.2.0-Linux-x86_64.sh) for Python 3 16 | - Install Pytorch and torchvision through Anaconda (Please follow the guide in [Pytorch](pytorch.org)) 17 | - Clone this repo 18 | ```bash 19 | git clone git@github.com:penincillin/DREAM.git 20 | cd DREAM 21 | ``` 22 | ### Prepare Data and Models 23 | Big files like model.zip could be downloaded both from Google Drive and Baidu Yun. If you have problems with downloading those files, you could contact me :) 24 | #### Face Alignment 25 | All the face images should be aligned. Please follow the align protocol in dataset [CelebA](http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html). After alignment, the face should be in the center of the image, and the size of image should be 178x218. Some aligned samples could be found in image/align_sample. 26 | 27 | #### Datasets 28 | In this paper, we use three face datasets. We train base model and DREAM block on [MS-Celeb-1M](https://www.microsoft.com/en-us/research/project/ms-celeb-1m-challenge-recognizing-one-million-celebrities-real-world/) 29 | We offer a subset of Ms-Celeb-1M with 10 celebrities, you could download from the following link 30 | Ms-Celeb-1M Subset (msceleb.zip): [Google Drive](https://drive.google.com/file/d/1om0pbwBX4RZHVuI3QXVrBj9mLtOK2PV8/view?usp=sharing)     [Baidu Yun](https://pan.baidu.com/s/1Zviee1QXnB7noArpAoy7Iw) 31 | 32 | We evaluate our the performance of our models on CFP and IJB-A. 33 | For CFP, we offer the code to get algined images from the original images (The code could only be runned on Linux). First, you need to download the original CFP dataset, and then download the image list from [Google Drive](https://drive.google.com/file/d/1B9QGThNd_-4Pg8O3si-EUYU9Px748p1C/view?usp=sharing)     [Baidu Yun](https://pan.baidu.com/s/1U_CzmLsJ2OaX4rJeJ7r92g) 34 | 35 | For IJBA, we provide the aligned images here. [Google Drive](https://drive.google.com/file/d/11p1eVSpyHZQUG0uBGyRoFnOXXTuZ501c/view?usp=sharing)　    [Baidu Yun](https://pan.baidu.com/s/1xLi6zDqwAeXEMV4aWi1k3g) 36 | 37 | #### Pretrained Models 38 | We offer several pretrained models. They could be downloaded from [Google Drive](https://drive.google.com/open?id=1CrWbsyAvqTA14ET2wvks_4U2_h1P52qK)     [Baidu Yun](https://pan.baidu.com/s/1LQmWZss0QoRc_chVIHsR_Q) 39 | 40 | ## Train DREAM Block 41 | ### stitch Training 42 | Prepare the feature extracted from any face recognition model (You could use the pretrained model we prepared). 43 | We prepared a piece of sample data (stitching.zip) which could be download from [Google Drive](https://drive.google.com/file/d/1x1K8MxAnVtpfaN3DfO4bdcKH39mmplj-/view?usp=sharing)     [Baidu Yun](https://pan.baidu.com/s/1QIEeE9RxRY6iK3wCpvUh2Q) 44 | - Download the sample data 45 | ```bash 46 | mkdir data 47 | mv stitching.zip data 48 | cd data 49 | unzip stitching.zip 50 | ``` 51 | - Train the model: 52 | ```bash 53 | cd src/stitching 54 | sh train_stitch.sh 55 | ``` 56 | 57 | 58 | ### end2end Training 59 | - Download the Ms-Celeb-1M Subset 60 | ```bash 61 | mkdir data 62 | mv msceleb.zip data 63 | cd data 64 | unzip msceleb.zip 65 | ``` 66 | - Train the model: 67 | ```bash 68 | cd src/end2end 69 | sh train.sh 70 | ``` 71 | ### evaluate CFP 72 | - Download the CFP dataset and preprocess the image. Then download the image list for evaluation 73 | ```bash 74 | # make sure you are in the root directory of DREAM project 75 | mkdir data 76 | cd src/preprocess 77 | sh align_cfp.sh 78 | cd data/CFP 79 | unzip CFP_protocol.zip 80 | ``` 81 | - Download pretrained model 82 | ```bash 83 | # make sure you are in the root directory of DREAM project 84 | cd ../ 85 | mv model.zip data 86 | cd data 87 | unzip model.zip 88 | ``` 89 | - Evaluate the pretrained model on CFP dataset 90 | ```bash 91 | # make sure you are in the root directory of DREAM project 92 | cd src/CFP 93 | sh eval_cfp.sh 94 | ``` 95 | 96 | ### evaluate IJBA 97 | - Download the IJBA dataset(contact me to get the aligned images) 98 | ```bash 99 | # make sure you are in the root directory of DREAM project 100 | mkdir data 101 | mv IJBA.zip data 102 | cd data 103 | unzip IJBA.zip 104 | ``` 105 | - Download pretrained models (If have downloaded the models, skip this step) 106 | ```bash 107 | # make sure you are in the root directory of DREAM project 108 | cd ../ 109 | mv model.zip data 110 | cd data 111 | unzip model.zip 112 | ``` 113 | - Evaluate the pretrained model on IJBA dataset 114 | ```bash 115 | # make sure you are in the root directory of DREAM project 116 | cd src/IJBA 117 | sh eval_ijba.sh 118 | ``` 119 | 120 | ## Citation 121 | Please cite the paper in your publications if it helps your research: 122 | 123 | 124 | 125 | @inproceedings{cao2018Dream, 126 | author = {Kaidi Cao and Yu Rong and Cheng Li and Xiaoou Tang and Chen Change Loy}, 127 | booktitle = {CVPR}, 128 | title = {Pose-Robust Face Recognition via Deep Residual Equivariant Mapping}, 129 | year = {2018} 130 | } 131 | 132 | -------------------------------------------------------------------------------- /src/CFP/ResNet.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import math 3 | import torch.utils.model_zoo as model_zoo 4 | import torch.nn.functional as F 5 | import numpy as np 6 | 7 | 8 | __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 9 | 'resnet152'] 10 | 11 | 12 | model_urls = { 13 | 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth', 14 | 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth', 15 | 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth', 16 | 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth', 17 | 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', 18 | } 19 | 20 | 21 | def conv3x3(in_planes, out_planes, stride=1): 22 | "3x3 convolution with padding" 23 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, 24 | padding=1, bias=False) 25 | 26 | 27 | class BasicBlock(nn.Module): 28 | expansion = 1 29 | 30 | def __init__(self, inplanes, planes, stride=1, downsample=None): 31 | super(BasicBlock, self).__init__() 32 | self.conv1 = conv3x3(inplanes, planes, stride) 33 | self.bn1 = nn.BatchNorm2d(planes) 34 | self.relu = nn.ReLU(inplace=True) 35 | self.conv2 = conv3x3(planes, planes) 36 | self.bn2 = nn.BatchNorm2d(planes) 37 | self.downsample = downsample 38 | self.stride = stride 39 | 40 | def forward(self, x): 41 | residual = x 42 | 43 | out = self.conv1(x) 44 | out = self.bn1(out) 45 | out = self.relu(out) 46 | 47 | out = self.conv2(out) 48 | out = self.bn2(out) 49 | 50 | if self.downsample is not None: 51 | residual = self.downsample(x) 52 | 53 | out += residual 54 | out = self.relu(out) 55 | 56 | return out 57 | 58 | 59 | class Bottleneck(nn.Module): 60 | expansion = 4 61 | 62 | def __init__(self, inplanes, planes, stride=1, downsample=None): 63 | super(Bottleneck, self).__init__() 64 | self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) 65 | self.bn1 = nn.BatchNorm2d(planes) 66 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, 67 | padding=1, bias=False) 68 | self.bn2 = nn.BatchNorm2d(planes) 69 | self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) 70 | self.bn3 = nn.BatchNorm2d(planes * 4) 71 | self.relu = nn.ReLU(inplace=True) 72 | self.downsample = downsample 73 | self.stride = stride 74 | 75 | def forward(self, x): 76 | residual = x 77 | 78 | out = self.conv1(x) 79 | out = self.bn1(out) 80 | out = self.relu(out) 81 | 82 | out = self.conv2(out) 83 | out = self.bn2(out) 84 | out = self.relu(out) 85 | 86 | out = self.conv3(out) 87 | out = self.bn3(out) 88 | 89 | if self.downsample is not None: 90 | residual = self.downsample(x) 91 | 92 | out += residual 93 | out = self.relu(out) 94 | 95 | return out 96 | 97 | 98 | class ResNet(nn.Module): 99 | 100 | def __init__(self, block, layers, num_classes=1000, extract_feature=False, end2end=False): 101 | self.inplanes = 64 102 | self.extract_feature = extract_feature 103 | self.end2end = end2end 104 | super(ResNet, self).__init__() 105 | self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, 106 | bias=False) 107 | self.bn1 = nn.BatchNorm2d(64) 108 | self.relu = nn.ReLU(inplace=True) 109 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) 110 | self.layer1 = self._make_layer(block, 64, layers[0]) 111 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2) 112 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2) 113 | self.layer4 = self._make_layer(block, 512, layers[3], stride=2) 114 | self.avgpool = nn.AvgPool2d(7) 115 | 116 | feat_dim = 256 117 | self.feature = nn.Linear(512 * block.expansion, feat_dim) 118 | if self.end2end: 119 | self.fc1 = nn.Linear(feat_dim, feat_dim) 120 | self.fc2 = nn.Linear(feat_dim, feat_dim) 121 | self.fc = nn.Linear(feat_dim, num_classes) 122 | 123 | for m in self.modules(): 124 | if isinstance(m, nn.Conv2d): 125 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 126 | m.weight.data.normal_(0, math.sqrt(2. / n)) 127 | elif isinstance(m, nn.BatchNorm2d): 128 | m.weight.data.fill_(1) 129 | m.bias.data.zero_() 130 | 131 | def _make_layer(self, block, planes, blocks, stride=1): 132 | downsample = None 133 | if stride != 1 or self.inplanes != planes * block.expansion: 134 | downsample = nn.Sequential( 135 | nn.Conv2d(self.inplanes, planes * block.expansion, 136 | kernel_size=1, stride=stride, bias=False), 137 | nn.BatchNorm2d(planes * block.expansion), 138 | ) 139 | 140 | layers = [] 141 | layers.append(block(self.inplanes, planes, stride, downsample)) 142 | self.inplanes = planes * block.expansion 143 | for i in range(1, blocks): 144 | layers.append(block(self.inplanes, planes)) 145 | 146 | return nn.Sequential(*layers) 147 | 148 | def forward(self, x, yaw): 149 | x = self.conv1(x) 150 | x = self.bn1(x) 151 | x = self.relu(x) 152 | x = self.maxpool(x) 153 | 154 | x = self.layer1(x) 155 | x = self.layer2(x) 156 | x = self.layer3(x) 157 | x = self.layer4(x) 158 | 159 | x = self.avgpool(x) 160 | x = x.view(x.size(0), -1) 161 | mid_feature = self.feature(x) 162 | 163 | if self.end2end: 164 | raw_feature = self.fc1(mid_feature) 165 | raw_feature = self.relu(raw_feature) 166 | raw_feature = self.fc2(raw_feature) 167 | raw_feature = self.relu(raw_feature) 168 | 169 | yaw = yaw.view(yaw.size(0),1) 170 | yaw = yaw.expand_as(raw_feature) 171 | 172 | feature = yaw * raw_feature + mid_feature 173 | 174 | else: 175 | feature = mid_feature 176 | 177 | res_x = self.fc(feature) 178 | 179 | if self.extract_feature: 180 | return feature 181 | else: 182 | return res_x 183 | 184 | 185 | def resnet18(pretrained=False, **kwargs): 186 | """Constructs a ResNet-18 model. 187 | 188 | Args: 189 | pretrained (bool): If True, returns a model pre-trained on ImageNet 190 | """ 191 | model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs) 192 | if pretrained: 193 | model.load_state_dict(model_zoo.load_url(model_urls['resnet18'])) 194 | return model 195 | 196 | 197 | def resnet34(pretrained=False, **kwargs): 198 | """Constructs a ResNet-34 model. 199 | 200 | Args: 201 | pretrained (bool): If True, returns a model pre-trained on ImageNet 202 | """ 203 | model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs) 204 | if pretrained: 205 | model.load_state_dict(model_zoo.load_url(model_urls['resnet34'])) 206 | return model 207 | 208 | 209 | def resnet50(pretrained=False, **kwargs): 210 | """Constructs a ResNet-50 model. 211 | 212 | Args: 213 | pretrained (bool): If True, returns a model pre-trained on ImageNet 214 | """ 215 | model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs) 216 | if pretrained: 217 | model.load_state_dict(model_zoo.load_url(model_urls['resnet50'])) 218 | return model 219 | 220 | 221 | def resnet101(pretrained=False, **kwargs): 222 | """Constructs a ResNet-101 model. 223 | 224 | Args: 225 | pretrained (bool): If True, returns a model pre-trained on ImageNet 226 | """ 227 | model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs) 228 | if pretrained: 229 | model.load_state_dict(model_zoo.load_url(model_urls['resnet101'])) 230 | return model 231 | 232 | 233 | def resnet152(pretrained=False, **kwargs): 234 | """Constructs a ResNet-152 model. 235 | 236 | Args: 237 | pretrained (bool): If True, returns a model pre-trained on ImageNet 238 | """ 239 | model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs) 240 | if pretrained: 241 | model.load_state_dict(model_zoo.load_url(model_urls['resnet152'])) 242 | return model 243 | -------------------------------------------------------------------------------- /src/CFP/eval_cfp.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os,sys,shutil 3 | import time 4 | import struct as st 5 | 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.parallel 9 | import torch.backends.cudnn as cudnn 10 | import torch.optim 11 | import torch.utils.data 12 | #import transforms 13 | import torchvision.datasets as datasets 14 | import torchvision.models as models 15 | import torchvision.transforms as transforms 16 | 17 | from selfDefine import CFPDataset, CaffeCrop 18 | from ResNet import resnet18, resnet50, resnet101 19 | from eval_roc import eval_roc_main 20 | 21 | model_names = sorted(name for name in models.__dict__ 22 | if name.islower() and not name.startswith("__") 23 | and callable(models.__dict__[name])) 24 | 25 | 26 | parser = argparse.ArgumentParser(description='PyTorch CelebA Training') 27 | parser.add_argument('--img_dir', metavar='DIR', default='', help='path to dataset') 28 | parser.add_argument('--arch', '-a', metavar='ARCH', default='resnet18', choices=model_names, 29 | help='model architecture: ' + 30 | ' | '.join(model_names) + 31 | ' (default: alexnet)') 32 | parser.add_argument('-j', '--workers', default=16, type=int, metavar='N', 33 | help='number of data loading workers (default: 4)') 34 | parser.add_argument('--epochs', default=90, type=int, metavar='N', 35 | help='number of total epochs to run') 36 | parser.add_argument('--start-epoch', default=0, type=int, metavar='N', 37 | help='manual epoch number (useful on restarts)') 38 | parser.add_argument('-b', '--batch-size', default=256, type=int, 39 | metavar='N', help='mini-batch size (default: 256)') 40 | parser.add_argument('--lr', '--learning-rate', default=0.1, type=float, 41 | metavar='LR', help='initial learning rate') 42 | parser.add_argument('--momentum', default=0.9, type=float, metavar='M', 43 | help='momentum') 44 | parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float, 45 | metavar='W', help='weight decay (default: 1e-4)') 46 | parser.add_argument('--print-freq', '-p', default=10, type=int, 47 | metavar='N', help='print frequency (default: 10)') 48 | parser.add_argument('--resume', default='./data/resnet18/checkpoint_40.pth.tar', type=str, metavar='PATH', 49 | help='path to latest checkpoint (default: none)') 50 | parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true', 51 | help='evaluate model on validation set') 52 | parser.add_argument('--pretrained', dest='pretrained', action='store_true', 53 | help='use pre-trained model') 54 | parser.add_argument('--model_dir','-m', default='./model', type=str) 55 | 56 | 57 | 58 | def extract_feat(arch, resume): 59 | global args, best_prec1 60 | args = parser.parse_args() 61 | 62 | if arch.find('end2end')>=0: 63 | end2end=True 64 | else: 65 | end2end=False 66 | 67 | arch = arch.split('_')[0] 68 | 69 | # load data and prepare dataset 70 | frontal_list_file = '../../data/CFP/protocol/frontal_list_nonli.txt' 71 | caffe_crop = CaffeCrop('test') 72 | frontal_dataset = CFPDataset(args.img_dir, frontal_list_file, 73 | transforms.Compose([caffe_crop,transforms.ToTensor()])) 74 | frontal_loader = torch.utils.data.DataLoader( 75 | frontal_dataset, 76 | batch_size=args.batch_size, shuffle=False, 77 | num_workers=args.workers, pin_memory=True) 78 | 79 | caffe_crop = CaffeCrop('test') 80 | profile_list_file = '../../data/CFP/protocol/profile_list_nonli.txt' 81 | profile_dataset = CFPDataset(args.img_dir, profile_list_file, 82 | transforms.Compose([caffe_crop,transforms.ToTensor()])) 83 | profile_loader = torch.utils.data.DataLoader( 84 | profile_dataset, 85 | batch_size=args.batch_size, shuffle=False, 86 | num_workers=args.workers, pin_memory=True) 87 | 88 | class_num = 13386 89 | 90 | model = None 91 | assert(arch in ['resnet18','resnet50','resnet101']) 92 | if arch == 'resnet18': 93 | model = resnet18(pretrained=False, num_classes=class_num, \ 94 | extract_feature=True, end2end=end2end) 95 | if arch == 'resnet50': 96 | model = resnet50(pretrained=False, num_classes=class_num,\ 97 | extract_feature=True, end2end=end2end) 98 | if arch == 'resnet101': 99 | model = resnet101(pretrained=False, num_classes=class_num,\ 100 | extract_feature=True, end2end=end2end) 101 | 102 | 103 | model = torch.nn.DataParallel(model).cuda() 104 | model.eval() 105 | 106 | 107 | 108 | assert(os.path.isfile(resume)) 109 | checkpoint = torch.load(resume) 110 | model.load_state_dict(checkpoint['state_dict']) 111 | 112 | cudnn.benchmark = True 113 | 114 | 115 | data_num = len(frontal_dataset) 116 | frontal_feat_file = './data/frontal_feat.bin' 117 | feat_dim = 256 118 | with open(frontal_feat_file, 'wb') as bin_f: 119 | bin_f.write(st.pack('ii', data_num, feat_dim)) 120 | for i, (input, yaw) in enumerate(frontal_loader): 121 | yaw = yaw.float().cuda(async=True) 122 | input_var = torch.autograd.Variable(input, volatile=True) 123 | yaw_var = torch.autograd.Variable(yaw, volatile=True) 124 | output = model(input_var, yaw_var) 125 | output_data = output.cpu().data.numpy() 126 | feat_num = output.size(0) 127 | 128 | for j in range(feat_num): 129 | bin_f.write(st.pack('f'*feat_dim, *tuple(output_data[j,:]))) 130 | 131 | data_num = len(profile_dataset.imgs) 132 | profile_feat_file = './data/profile_feat.bin' 133 | with open(profile_feat_file, 'wb') as bin_f: 134 | bin_f.write(st.pack('ii', data_num, feat_dim)) 135 | for i, (input,yaw) in enumerate(profile_loader): 136 | yaw = yaw.float().cuda(async=True) 137 | input_var = torch.autograd.Variable(input, volatile=True) 138 | yaw_var = torch.autograd.Variable(yaw, volatile=True) 139 | output = model(input_var, yaw_var) 140 | output_data = output.cpu().data.numpy() 141 | feat_num = output.size(0) 142 | 143 | for j in range(feat_num): 144 | bin_f.write(st.pack('f'*feat_dim, *tuple(output_data[j,:]))) 145 | 146 | 147 | if __name__ == '__main__': 148 | 149 | infos = [ ('resnet50_naive', '../../data/model/cfp_res50_naive.pth.tar'), 150 | ('resnet50_end2end', '../../data/model/cfp_res50_end2end.pth.tar'), ] 151 | 152 | 153 | for arch, model_path in infos: 154 | print("{} {}".format(arch, model_path)) 155 | extract_feat(arch, model_path) 156 | eval_roc_main() 157 | print() 158 | -------------------------------------------------------------------------------- /src/CFP/eval_cfp.sh: -------------------------------------------------------------------------------- 1 | data_dir=data 2 | if [ ! -d $data_dir ]; then 3 | mkdir $data_dir 4 | fi 5 | 6 | CUDA_VISIBLE_DEVICES=0,1 python eval_cfp.py \ 7 | --img_dir ../../data/CFP/ \ 8 | -------------------------------------------------------------------------------- /src/CFP/eval_roc.py: -------------------------------------------------------------------------------- 1 | import os, sys, shutil 2 | import numpy as np 3 | import matplotlib.pyplot as plt 4 | import math 5 | import random as rd 6 | import struct as st 7 | from scipy import spatial 8 | from sklearn import metrics 9 | from scipy.optimize import fsolve 10 | from scipy.interpolate import interp1d 11 | from sklearn.metrics.pairwise import cosine_similarity 12 | 13 | 14 | # {{{ draw roc curve 15 | def draw_roc(fpr, tpr, title, img_name="roc.png"): 16 | plt.figure(figsize=(16, 8)) 17 | plt.yticks(np.arange(0.0, 1.05, 0.05)) 18 | plt.xticks(np.arange(0.0, 1.05, 0.05)) 19 | plt.title(title) 20 | plt.plot(fpr, tpr, linewidth=1, color='r') 21 | #plt.xscale('log') 22 | #plt.yscale('log') 23 | plt.savefig(img_name) 24 | 25 | def draw_multiple_roc(tprs, fprs, labels, title, img_name): 26 | plt.figure(figsize=(16,8)) 27 | plt.yticks(np.arange(0.0, 1.05, 0.05)) 28 | plt.xticks(np.arange(0.0, 1.05, 0.05)) 29 | for i in range(len(tprs)): 30 | plt.plot(fprs[i], tprs[i], label=labels[i]) 31 | plt.legend(loc=4) 32 | plt.title(title) 33 | plt.savefig(img_name) 34 | # }}} 35 | 36 | def load_feat(feat_file): 37 | feats = list() 38 | with open(feat_file, 'rb') as in_f: 39 | feat_num, feat_dim = st.unpack('ii', in_f.read(8)) 40 | for i in range(feat_num): 41 | feat = np.array(st.unpack('f'*feat_dim, in_f.read(4*feat_dim))) 42 | feats.append(feat) 43 | return feats 44 | 45 | 46 | #calculate eer{{{ 47 | def calc_eer(fpr, tpr, method=0): 48 | if method == 0: 49 | min_dis, eer = 100.0, 1.0 50 | for i in range(fpr.size): 51 | if(fpr[i]+tpr[i] > 1.0): 52 | break 53 | mid_res = abs(fpr[i]+tpr[i]-1.0) 54 | if(mid_res < min_dis): 55 | min_dis = mid_res 56 | eer = fpr[i] 57 | return eer 58 | else: 59 | f = lambda x: np.interp(x, fpr, tpr)+x-1 60 | return fsolve(f, 0.0) 61 | #}}} 62 | 63 | def eval_roc(protocol_dir, pair_type, split_name, frontal_feats, profile_feats): 64 | labels, scores = [],[] 65 | for idx, pair_file in enumerate(['diff.txt', 'same.txt']): 66 | label = idx 67 | full_pair_file = protocol_dir+'/'+pair_type+'/'+split_name+'/'+pair_file 68 | with open(full_pair_file, 'r') as in_f: 69 | for line in in_f: 70 | record = line.strip().split(',') 71 | pair1, pair2 = int(record[0]),int(record[1]) 72 | if pair_type == 'FF': 73 | vec1, vec2 = frontal_feats[pair1-1], frontal_feats[pair2-1] 74 | else: 75 | vec1, vec2 = frontal_feats[pair1-1], profile_feats[pair2-1] 76 | score = 1-spatial.distance.cosine(vec1,vec2) 77 | scores.append(score) 78 | labels.append(label) 79 | fpr, tpr, thresholds = metrics.roc_curve(labels, scores) 80 | #draw_roc(fpr,tpr,title='center loss') 81 | auc = metrics.auc(fpr, tpr) 82 | eer = calc_eer(fpr, tpr) 83 | return auc, eer 84 | 85 | 86 | def eval_roc_main(): 87 | 88 | frontal_feat_file = './data/frontal_feat.bin' 89 | profile_feat_file = './data/profile_feat.bin' 90 | #frontal_feat_file = './data/center_loss_frontal_feat.bin' 91 | #profile_feat_file = './data/center_loss_profile_feat.bin' 92 | 93 | frontal_feats = load_feat(frontal_feat_file) 94 | profile_feats = load_feat(profile_feat_file) 95 | 96 | protocol_dir = '/mnt/SSD/rongyu/data/cfp/cfp_align/protocol/Split' 97 | pair_types = ['FF', 'FP'] 98 | split_num = 10 99 | 100 | for pair_type in pair_types: 101 | print_info = 'Frontal-Frontal' if pair_type=='FF' else 'Frontal-Profile' 102 | print('----- result for {} -----'.format(print_info)) 103 | aucs, eers = list(), list() 104 | for split_id in range(split_num): 105 | split_name = str(split_id+1) 106 | if len(split_name)<2: split_name = '0'+split_name 107 | auc, eer = eval_roc(protocol_dir, pair_type, split_name, frontal_feats, profile_feats) 108 | #print('{} split {}, auc: {}, eer: {}'.format(pair_type,split_name,auc,eer)) 109 | aucs.append(auc) 110 | eers.append(eer) 111 | print('Average auc:', np.mean(aucs)) 112 | print('Average eer:', np.mean(eers)) 113 | 114 | if __name__ == '__main__': 115 | eval_roc_main() 116 | -------------------------------------------------------------------------------- /src/CFP/selfDefine.py: -------------------------------------------------------------------------------- 1 | import os, sys, shutil 2 | import random as rd 3 | import struct as st 4 | 5 | from PIL import Image 6 | import numpy as np 7 | from scipy import misc 8 | 9 | import torch 10 | import torch.nn as nn 11 | import torch.nn.functional as F 12 | import torch.utils.data as data 13 | from torch.autograd import Variable 14 | from torch.nn.modules.loss import _WeightedLoss 15 | 16 | def load_imgs(img_dir, image_list_file): 17 | imgs = list() 18 | with open(image_list_file, 'r') as imf: 19 | for line in imf: 20 | record = line.strip().split() 21 | img_path, yaw = os.path.join(img_dir,record[0]), float(record[1]) 22 | imgs.append((img_path, yaw)) 23 | return imgs 24 | 25 | 26 | class CFPDataset(data.Dataset): 27 | def __init__(self, img_dir, image_list_file, transform=None): 28 | self.imgs = load_imgs(img_dir, image_list_file) 29 | self.transform = transform 30 | 31 | def __getitem__(self, index): 32 | path, yaw = self.imgs[index] 33 | img = Image.open(path).convert("RGB") 34 | if self.transform is not None: 35 | img = self.transform(img) 36 | return img, yaw 37 | 38 | def __len__(self): 39 | return len(self.imgs) 40 | 41 | 42 | class CaffeCrop(object): 43 | #This class take the same behavior as sensenet 44 | def __init__(self, phase): 45 | assert(phase=='train' or phase=='test') 46 | self.phase = phase 47 | 48 | def __call__(self, img): 49 | # pre determined parameters 50 | final_size = 224 51 | final_width = final_height = final_size 52 | crop_size = 110 53 | crop_height = crop_width = crop_size 54 | crop_center_y_offset = 15 55 | crop_center_x_offset = 0 56 | if self.phase == 'train': 57 | scale_aug = 0.02 58 | trans_aug = 0.01 59 | else: 60 | scale_aug = 0.0 61 | trans_aug = 0.0 62 | 63 | # computed parameters 64 | randint = rd.randint 65 | scale_height_diff = (randint(0,1000)/500-1)*scale_aug 66 | crop_height_aug = crop_height*(1+scale_height_diff) 67 | scale_width_diff = (randint(0,1000)/500-1)*scale_aug 68 | crop_width_aug = crop_width*(1+scale_width_diff) 69 | 70 | 71 | trans_diff_x = (randint(0,1000)/500-1)*trans_aug 72 | trans_diff_y = (randint(0,1000)/500-1)*trans_aug 73 | 74 | 75 | center = ((img.width/2 + crop_center_x_offset)*(1+trans_diff_x), 76 | (img.height/2 + crop_center_y_offset)*(1+trans_diff_y)) 77 | 78 | 79 | if center[0] < crop_width_aug/2: 80 | crop_width_aug = center[0]*2-0.5 81 | print(1) 82 | if center[1] < crop_height_aug/2: 83 | crop_height_aug = center[1]*2-0.5 84 | print(2) 85 | if (center[0]+crop_width_aug/2) >= img.width: 86 | crop_width_aug = (img.width-center[0])*2-0.5 87 | print(3) 88 | if (center[1]+crop_height_aug/2) >= img.height: 89 | crop_height_aug = (img.height-center[1])*2-0.5 90 | print(4) 91 | 92 | crop_box = (center[0]-crop_width_aug/2, center[1]-crop_height_aug/2, 93 | center[0]+crop_width_aug/2, center[1]+crop_width_aug/2) 94 | 95 | mid_img = img.crop(crop_box) 96 | res_img = mid_img.resize( (final_width, final_height) ) 97 | return res_img 98 | -------------------------------------------------------------------------------- /src/IJBA/ResNet.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import math 3 | import torch.utils.model_zoo as model_zoo 4 | import torch.nn.functional as F 5 | import numpy as np 6 | 7 | 8 | __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 9 | 'resnet152'] 10 | 11 | 12 | model_urls = { 13 | 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth', 14 | 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth', 15 | 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth', 16 | 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth', 17 | 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', 18 | } 19 | 20 | 21 | def conv3x3(in_planes, out_planes, stride=1): 22 | "3x3 convolution with padding" 23 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, 24 | padding=1, bias=False) 25 | 26 | 27 | class BasicBlock(nn.Module): 28 | expansion = 1 29 | 30 | def __init__(self, inplanes, planes, stride=1, downsample=None): 31 | super(BasicBlock, self).__init__() 32 | self.conv1 = conv3x3(inplanes, planes, stride) 33 | self.bn1 = nn.BatchNorm2d(planes) 34 | self.relu = nn.ReLU(inplace=True) 35 | self.conv2 = conv3x3(planes, planes) 36 | self.bn2 = nn.BatchNorm2d(planes) 37 | self.downsample = downsample 38 | self.stride = stride 39 | 40 | def forward(self, x): 41 | residual = x 42 | 43 | out = self.conv1(x) 44 | out = self.bn1(out) 45 | out = self.relu(out) 46 | 47 | out = self.conv2(out) 48 | out = self.bn2(out) 49 | 50 | if self.downsample is not None: 51 | residual = self.downsample(x) 52 | 53 | out += residual 54 | out = self.relu(out) 55 | 56 | return out 57 | 58 | 59 | class Bottleneck(nn.Module): 60 | expansion = 4 61 | 62 | def __init__(self, inplanes, planes, stride=1, downsample=None): 63 | super(Bottleneck, self).__init__() 64 | self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) 65 | self.bn1 = nn.BatchNorm2d(planes) 66 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, 67 | padding=1, bias=False) 68 | self.bn2 = nn.BatchNorm2d(planes) 69 | self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) 70 | self.bn3 = nn.BatchNorm2d(planes * 4) 71 | self.relu = nn.ReLU(inplace=True) 72 | self.downsample = downsample 73 | self.stride = stride 74 | 75 | def forward(self, x): 76 | residual = x 77 | 78 | out = self.conv1(x) 79 | out = self.bn1(out) 80 | out = self.relu(out) 81 | 82 | out = self.conv2(out) 83 | out = self.bn2(out) 84 | out = self.relu(out) 85 | 86 | out = self.conv3(out) 87 | out = self.bn3(out) 88 | 89 | if self.downsample is not None: 90 | residual = self.downsample(x) 91 | 92 | out += residual 93 | out = self.relu(out) 94 | 95 | return out 96 | 97 | 98 | class ResNet(nn.Module): 99 | 100 | def __init__(self, block, layers, num_classes=1000, extract_feature=False, end2end=False): 101 | self.inplanes = 64 102 | self.extract_feature = extract_feature 103 | self.end2end = end2end 104 | super(ResNet, self).__init__() 105 | self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, 106 | bias=False) 107 | self.bn1 = nn.BatchNorm2d(64) 108 | self.relu = nn.ReLU(inplace=True) 109 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) 110 | self.layer1 = self._make_layer(block, 64, layers[0]) 111 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2) 112 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2) 113 | self.layer4 = self._make_layer(block, 512, layers[3], stride=2) 114 | self.avgpool = nn.AvgPool2d(7) 115 | 116 | feat_dim = 256 117 | self.feature = nn.Linear(512 * block.expansion, feat_dim) 118 | if self.end2end: 119 | self.fc1 = nn.Linear(feat_dim, feat_dim) 120 | self.fc2 = nn.Linear(feat_dim, feat_dim) 121 | #self.fc = nn.Linear(feat_dim, num_classes) 122 | 123 | for m in self.modules(): 124 | if isinstance(m, nn.Conv2d): 125 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 126 | m.weight.data.normal_(0, math.sqrt(2. / n)) 127 | elif isinstance(m, nn.BatchNorm2d): 128 | m.weight.data.fill_(1) 129 | m.bias.data.zero_() 130 | 131 | def _make_layer(self, block, planes, blocks, stride=1): 132 | downsample = None 133 | if stride != 1 or self.inplanes != planes * block.expansion: 134 | downsample = nn.Sequential( 135 | nn.Conv2d(self.inplanes, planes * block.expansion, 136 | kernel_size=1, stride=stride, bias=False), 137 | nn.BatchNorm2d(planes * block.expansion), 138 | ) 139 | 140 | layers = [] 141 | layers.append(block(self.inplanes, planes, stride, downsample)) 142 | self.inplanes = planes * block.expansion 143 | for i in range(1, blocks): 144 | layers.append(block(self.inplanes, planes)) 145 | 146 | return nn.Sequential(*layers) 147 | 148 | def forward(self, x, yaw): 149 | x = self.conv1(x) 150 | x = self.bn1(x) 151 | x = self.relu(x) 152 | x = self.maxpool(x) 153 | 154 | x = self.layer1(x) 155 | x = self.layer2(x) 156 | x = self.layer3(x) 157 | x = self.layer4(x) 158 | 159 | x = self.avgpool(x) 160 | x = x.view(x.size(0), -1) 161 | mid_feature = self.feature(x) 162 | 163 | if self.end2end: 164 | raw_feature = self.fc1(mid_feature) 165 | raw_feature = self.relu(raw_feature) 166 | raw_feature = self.fc2(raw_feature) 167 | raw_feature = self.relu(raw_feature) 168 | 169 | yaw = yaw.view(yaw.size(0),1) 170 | yaw = yaw.expand_as(raw_feature) 171 | 172 | feature = yaw * raw_feature + mid_feature 173 | 174 | else: 175 | feature = mid_feature 176 | 177 | return feature 178 | 179 | 180 | def resnet18(pretrained=False, **kwargs): 181 | """Constructs a ResNet-18 model. 182 | 183 | Args: 184 | pretrained (bool): If True, returns a model pre-trained on ImageNet 185 | """ 186 | model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs) 187 | if pretrained: 188 | model.load_state_dict(model_zoo.load_url(model_urls['resnet18'])) 189 | return model 190 | 191 | 192 | def resnet34(pretrained=False, **kwargs): 193 | """Constructs a ResNet-34 model. 194 | 195 | Args: 196 | pretrained (bool): If True, returns a model pre-trained on ImageNet 197 | """ 198 | model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs) 199 | if pretrained: 200 | model.load_state_dict(model_zoo.load_url(model_urls['resnet34'])) 201 | return model 202 | 203 | 204 | def resnet50(pretrained=False, **kwargs): 205 | """Constructs a ResNet-50 model. 206 | 207 | Args: 208 | pretrained (bool): If True, returns a model pre-trained on ImageNet 209 | """ 210 | model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs) 211 | if pretrained: 212 | model.load_state_dict(model_zoo.load_url(model_urls['resnet50'])) 213 | return model 214 | 215 | 216 | def resnet101(pretrained=False, **kwargs): 217 | """Constructs a ResNet-101 model. 218 | 219 | Args: 220 | pretrained (bool): If True, returns a model pre-trained on ImageNet 221 | """ 222 | model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs) 223 | if pretrained: 224 | model.load_state_dict(model_zoo.load_url(model_urls['resnet101'])) 225 | return model 226 | 227 | 228 | def resnet152(pretrained=False, **kwargs): 229 | """Constructs a ResNet-152 model. 230 | 231 | Args: 232 | pretrained (bool): If True, returns a model pre-trained on ImageNet 233 | """ 234 | model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs) 235 | if pretrained: 236 | model.load_state_dict(model_zoo.load_url(model_urls['resnet152'])) 237 | return model 238 | -------------------------------------------------------------------------------- /src/IJBA/eval_ijba.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os,sys,shutil 3 | import time 4 | import struct as st 5 | 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.parallel 9 | import torch.backends.cudnn as cudnn 10 | import torch.optim 11 | import torch.utils.data 12 | #import transforms 13 | import torchvision.datasets as datasets 14 | import torchvision.models as models 15 | import torchvision.transforms as transforms 16 | 17 | from selfDefine import CFPDataset, CaffeCrop 18 | 19 | from ResNet import resnet18, resnet50, resnet101 20 | 21 | from test_recog import test_recog 22 | from test_verify import test_verify 23 | 24 | 25 | model_names = sorted(name for name in models.__dict__ 26 | if name.islower() and not name.startswith("__") 27 | and callable(models.__dict__[name])) 28 | 29 | 30 | parser = argparse.ArgumentParser(description='PyTorch CelebA Training') 31 | parser.add_argument('--img_dir', metavar='DIR', default='', help='path to dataset') 32 | parser.add_argument('--arch', '-a', metavar='ARCH', default='resnet18', choices=model_names, 33 | help='model architecture: ' + 34 | ' | '.join(model_names) + 35 | ' (default: alexnet)') 36 | parser.add_argument('-j', '--workers', default=16, type=int, metavar='N', 37 | help='number of data loading workers (default: 4)') 38 | parser.add_argument('--epochs', default=90, type=int, metavar='N', 39 | help='number of total epochs to run') 40 | parser.add_argument('--start-epoch', default=0, type=int, metavar='N', 41 | help='manual epoch number (useful on restarts)') 42 | parser.add_argument('-b', '--batch-size', default=512, type=int, 43 | metavar='N', help='mini-batch size (default: 256)') 44 | parser.add_argument('--lr', '--learning-rate', default=0.1, type=float, 45 | metavar='LR', help='initial learning rate') 46 | parser.add_argument('--momentum', default=0.9, type=float, metavar='M', 47 | help='momentum') 48 | parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float, 49 | metavar='W', help='weight decay (default: 1e-4)') 50 | parser.add_argument('--print-freq', '-p', default=10, type=int, 51 | metavar='N', help='print frequency (default: 10)') 52 | parser.add_argument('--resume', default='./data/resnet18.pth.tar', type=str, metavar='PATH', 53 | help='path to latest checkpoint (default: none)') 54 | parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true', 55 | help='evaluate model on validation set') 56 | parser.add_argument('--pretrained', dest='pretrained', action='store_true', 57 | help='use pre-trained model') 58 | parser.add_argument('--model_dir','-m', default='./model', type=str) 59 | 60 | 61 | 62 | def extract_feat(arch, model_path, yaw_type): 63 | global args, best_prec1 64 | args = parser.parse_args() 65 | 66 | if arch.find('end2end')>=0: 67 | end2end=True 68 | else: 69 | end2end=False 70 | 71 | arch = arch.split('_')[0] 72 | 73 | class_num = 87020 74 | #class_num = 13386 75 | 76 | model = None 77 | assert(arch in ['resnet18','resnet50','resnet101']) 78 | if arch == 'resnet18': 79 | model = resnet18(pretrained=False, num_classes=class_num, \ 80 | extract_feature=True, end2end=end2end) 81 | if arch == 'resnet50': 82 | model = resnet50(pretrained=False, num_classes=class_num, \ 83 | extract_feature=True, end2end=end2end) 84 | if arch == 'resnet101': 85 | model = resnet101(pretrained=False, num_classes=class_num, \ 86 | extract_feature=True, end2end=end2end) 87 | 88 | model = torch.nn.DataParallel(model).cuda() 89 | model.eval() 90 | 91 | 92 | assert(os.path.isfile(model_path)) 93 | checkpoint = torch.load(model_path) 94 | pretrained_state_dict = checkpoint['state_dict'] 95 | model_state_dict = model.state_dict() 96 | for key in pretrained_state_dict: 97 | if key in model_state_dict: 98 | model_state_dict[key] = pretrained_state_dict[key] 99 | model.load_state_dict(model_state_dict) 100 | 101 | print('load trained model complete') 102 | 103 | caffe_crop = CaffeCrop('test') 104 | 105 | infos = [ ('../../data/IJBA/align_image_11', 'ijb_a_11_align_split', 'frame'), 106 | ('../../data/IJBA/align_image_11', 'ijb_a_11_align_split', 'img'), 107 | ('../../data/IJBA/align_image_1N', 'split', 'gallery'), 108 | ('../../data/IJBA/align_image_1N', 'split', 'probe') ] 109 | 110 | 111 | for root_dir, sub_dir, img_type in infos: 112 | 113 | for split in range(1,11): 114 | split_dir = os.path.join(root_dir, sub_dir+str(split)) 115 | img_dir = os.path.join(split_dir, img_type) 116 | img_list_file = os.path.join(split_dir, '{}_list_{}.txt'.format(img_type, yaw_type)) 117 | 118 | img_dataset = CFPDataset(args.img_dir, img_list_file, 119 | transforms.Compose([caffe_crop,transforms.ToTensor()])) 120 | img_loader = torch.utils.data.DataLoader( 121 | img_dataset, 122 | batch_size=args.batch_size, shuffle=False, 123 | num_workers=args.workers, pin_memory=True) 124 | 125 | 126 | data_num = len(img_dataset) 127 | img_feat_file = os.path.join(split_dir, '{}_{}_feat.bin'.format(arch, img_type)) 128 | feat_dim = 256 129 | with open(img_feat_file, 'wb') as bin_f: 130 | bin_f.write(st.pack('ii', data_num, feat_dim)) 131 | for i, (input, yaw) in enumerate(img_loader): 132 | yaw = yaw.float().cuda(async=True) 133 | yaw_var = torch.autograd.Variable(yaw) 134 | input_var = torch.autograd.Variable(input, volatile=True) 135 | output = model(input_var, yaw_var) 136 | output_data = output.cpu().data.numpy() 137 | feat_num = output.size(0) 138 | 139 | for j in range(feat_num): 140 | bin_f.write(st.pack('f'*feat_dim, *tuple(output_data[j,:]))) 141 | 142 | print('we have complete {} {}'.format(img_type, split)) 143 | 144 | 145 | if __name__ == '__main__': 146 | 147 | infos = [ ['resnet18_naive', '../../data/model/ijba_res18_naive.pth.tar', 'nonli'], 148 | ['resnet18_end2end', '../../data/model/ijba_res18_end2end.pth.tar', 'nonli'], ] 149 | 150 | 151 | for arch, model_path, yaw_type in infos: 152 | print(model_path.split('/')[-1]) 153 | extract_feat(arch, model_path, yaw_type) 154 | arch = arch.split('_')[0] 155 | test_recog(arch) 156 | test_verify(arch) 157 | -------------------------------------------------------------------------------- /src/IJBA/eval_ijba.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env sh 2 | CUDA_VISIBLE_DEVICES=2,3 python eval_ijba.py 2>&1 | tee eval_res.txt 3 | -------------------------------------------------------------------------------- /src/IJBA/selfDefine.py: -------------------------------------------------------------------------------- 1 | import os, sys, shutil 2 | import random as rd 3 | import struct as st 4 | 5 | from PIL import Image 6 | import numpy as np 7 | from scipy import misc 8 | 9 | import torch 10 | import torch.nn as nn 11 | import torch.nn.functional as F 12 | import torch.utils.data as data 13 | from torch.autograd import Variable 14 | from torch.nn.modules.loss import _WeightedLoss 15 | 16 | def load_imgs(img_dir, image_list_file): 17 | imgs = list() 18 | with open(image_list_file, 'r') as imf: 19 | for line in imf: 20 | record = line.strip().split() 21 | img_path, yaw = os.path.join(img_dir,record[0]), float(record[1]) 22 | imgs.append((img_path, yaw)) 23 | return imgs 24 | 25 | 26 | class CFPDataset(data.Dataset): 27 | def __init__(self, img_dir, image_list_file, transform=None): 28 | self.imgs = load_imgs(img_dir, image_list_file) 29 | self.transform = transform 30 | 31 | def __getitem__(self, index): 32 | path, yaw = self.imgs[index] 33 | img = Image.open(path).convert("RGB") 34 | if self.transform is not None: 35 | img = self.transform(img) 36 | return img, yaw 37 | 38 | def __len__(self): 39 | return len(self.imgs) 40 | 41 | 42 | class CaffeCrop(object): 43 | #This class take the same behavior as sensenet 44 | def __init__(self, phase): 45 | assert(phase=='train' or phase=='test') 46 | self.phase = phase 47 | 48 | def __call__(self, img): 49 | # pre determined parameters 50 | final_size = 224 51 | final_width = final_height = final_size 52 | crop_size = 110 53 | crop_height = crop_width = crop_size 54 | crop_center_y_offset = 15 55 | crop_center_x_offset = 0 56 | if self.phase == 'train': 57 | scale_aug = 0.02 58 | trans_aug = 0.01 59 | else: 60 | scale_aug = 0.0 61 | trans_aug = 0.0 62 | 63 | # computed parameters 64 | randint = rd.randint 65 | scale_height_diff = (randint(0,1000)/500-1)*scale_aug 66 | crop_height_aug = crop_height*(1+scale_height_diff) 67 | scale_width_diff = (randint(0,1000)/500-1)*scale_aug 68 | crop_width_aug = crop_width*(1+scale_width_diff) 69 | 70 | 71 | trans_diff_x = (randint(0,1000)/500-1)*trans_aug 72 | trans_diff_y = (randint(0,1000)/500-1)*trans_aug 73 | 74 | 75 | center = ((img.width/2 + crop_center_x_offset)*(1+trans_diff_x), 76 | (img.height/2 + crop_center_y_offset)*(1+trans_diff_y)) 77 | 78 | 79 | if center[0] < crop_width_aug/2: 80 | crop_width_aug = center[0]*2-0.5 81 | print(1) 82 | if center[1] < crop_height_aug/2: 83 | crop_height_aug = center[1]*2-0.5 84 | print(2) 85 | if (center[0]+crop_width_aug/2) >= img.width: 86 | crop_width_aug = (img.width-center[0])*2-0.5 87 | print(3) 88 | if (center[1]+crop_height_aug/2) >= img.height: 89 | crop_height_aug = (img.height-center[1])*2-0.5 90 | print(4) 91 | 92 | crop_box = (center[0]-crop_width_aug/2, center[1]-crop_height_aug/2, 93 | center[0]+crop_width_aug/2, center[1]+crop_width_aug/2) 94 | 95 | mid_img = img.crop(crop_box) 96 | res_img = mid_img.resize( (final_width, final_height) ) 97 | return res_img 98 | -------------------------------------------------------------------------------- /src/IJBA/test_recog.py: -------------------------------------------------------------------------------- 1 | import os, sys, shutil 2 | import struct as st 3 | import numpy as np 4 | from scipy import spatial 5 | from sklearn import metrics 6 | from scipy.optimize import fsolve 7 | from scipy.interpolate import interp1d 8 | from sklearn.metrics.pairwise import cosine_similarity 9 | 10 | def load_meta_data(meta_file, sub_dir): 11 | meta_data = dict() 12 | with open(meta_file, 'r') as in_f: 13 | in_f.readline() # the first line is not data 14 | for idx, line in enumerate(in_f): 15 | record = line.strip().split(',') 16 | template, class_id = int(record[0]), int(record[1]) 17 | img_path = '{}/{}.jpg'.format(sub_dir,idx+1) 18 | if template not in meta_data: 19 | meta_data[template] = ( class_id, [img_path,] ) 20 | else: 21 | meta_data[template][1].append(img_path) 22 | return meta_data 23 | 24 | def load_feat(list_file, bin_file): 25 | mid_feats = dict() 26 | with open(list_file, 'r') as list_f, open(bin_file, 'rb') as bin_f: 27 | (data_num, feat_dim) = st.unpack('ii', bin_f.read(8)) 28 | for line in list_f: 29 | record = line.strip().split('/') 30 | img_name = '/'.join(record[-2:]) 31 | feat = np.array(st.unpack('f'*feat_dim, bin_f.read(4*feat_dim))) 32 | mid_feats[img_name] = feat 33 | return mid_feats, feat_dim 34 | 35 | def update_meta_data(meta_data, feats, feat_dim): 36 | new_meta_data = dict() 37 | for template in meta_data.keys(): 38 | class_id, img_names = meta_data[template] 39 | feat = np.zeros(feat_dim) 40 | feat_num = 0 41 | for img_name in img_names: 42 | if img_name in feats: 43 | feat += feats[img_name] 44 | feat_num += 1 45 | #else: 46 | #print(img_name) 47 | if feat_num > 0: feat /= feat_num 48 | if feat_num > 0: new_meta_data[template] = (class_id, feat) 49 | return new_meta_data 50 | 51 | def get_top(probe, gallery_data): 52 | score_info = list() 53 | probe_id, probe_feat = probe 54 | for template in gallery_data.keys(): 55 | gallery_id, gallery_feat = gallery_data[template] 56 | score = 1-spatial.distance.cosine(probe_feat, gallery_feat) 57 | score_info.append((gallery_id, score)) 58 | score_info = sorted(score_info, key=lambda a:a[1], reverse=True) 59 | top5_id = [item[0] for item in score_info[:5]] 60 | return top5_id 61 | 62 | 63 | def eval_recog(probe_data, gallery_data): 64 | gallery_ids = set() 65 | for template in gallery_data.keys(): 66 | gallery_ids.add(gallery_data[template][0]) 67 | 68 | top1_num, top5_num, tot_num = 0, 0, 0 69 | for template in probe_data.keys(): 70 | class_id = probe_data[template][0] 71 | if class_id not in gallery_ids: continue 72 | top5_id = get_top(probe_data[template], gallery_data) 73 | if class_id==top5_id[0]: 74 | top1_num += 1 75 | top5_num += 1 76 | elif class_id in top5_id: 77 | top5_num += 1 78 | tot_num += 1 79 | return top1_num/tot_num, top5_num/tot_num 80 | 81 | 82 | def test_recog(model_type): 83 | 84 | #model_type = 'resnet18' 85 | protocol_dir = '../../data/IJBA/protocol_1N' 86 | align_img_dir = '../../data/IJBA/align_image_1N' 87 | split_num = 10 88 | top1s, top5s = list(), list() 89 | for split in range(1, split_num+1): 90 | 91 | # load meta data first 92 | split_protocol_dir = os.path.join(protocol_dir, 'split{}'.format(split)) 93 | #split_img_dir = os.path.join(align_img_dir, 'split{}'.format(split)) 94 | probe_file = os.path.join(split_protocol_dir, 'search_probe_{}.csv'.format(split)) 95 | 96 | probe_data = load_meta_data(probe_file, 'probe') 97 | gallery_file = os.path.join(split_protocol_dir, 'search_gallery_{}.csv'.format(split)) 98 | gallery_data = load_meta_data(gallery_file, 'gallery') 99 | 100 | # load extract feat 101 | feats = dict() 102 | feat_dim = 0 103 | split_img_dir = os.path.join(align_img_dir, 'split{}'.format(split)) 104 | for img_type in ['gallery', 'probe']: 105 | list_file = os.path.join(split_img_dir, '{}_list.txt'.format(img_type)) 106 | bin_file = os.path.join(split_img_dir, '{}_{}_feat.bin'.format(model_type, img_type)) 107 | mid_feats, feat_dim = load_feat(list_file, bin_file) 108 | feats.update(mid_feats) 109 | 110 | # update meta data 111 | probe_data = update_meta_data(probe_data, feats, feat_dim) 112 | gallery_data = update_meta_data(gallery_data, feats, feat_dim) 113 | top1, top5 = eval_recog(probe_data, gallery_data) 114 | top1s.append(top1) 115 | top5s.append(top5) 116 | print('split {}, top1: {}, top5: {}'.format(split,top1,top5)) 117 | 118 | print('top1: {} +/- {}'.format(np.mean(top1s), np.std(top1s))) 119 | print('top5: {} +/- {}'.format(np.mean(top5s), np.std(top5s))) 120 | 121 | return np.mean(top1s), np.std(top1s), np.mean(top5s), np.std(top5s) 122 | 123 | 124 | if __name__ == '__main__': 125 | test_recog(model_type='resnet50') 126 | -------------------------------------------------------------------------------- /src/IJBA/test_verify.py: -------------------------------------------------------------------------------- 1 | import os, sys, shutil 2 | import struct as st 3 | import numpy as np 4 | import bisect 5 | import pickle 6 | from scipy import spatial 7 | from sklearn import metrics 8 | from scipy.optimize import fsolve 9 | from scipy.interpolate import interp1d 10 | from sklearn.metrics.pairwise import cosine_similarity 11 | 12 | def load_meta_data(meta_file): 13 | meta_data = dict() 14 | with open(meta_file, 'r') as in_f: 15 | in_f.readline() # the first line is not data 16 | for line in in_f: 17 | record = line.strip().split(',') 18 | template, class_id, img_path = int(record[0]), int(record[1]), record[2] 19 | if template not in meta_data: 20 | meta_data[template] = ( class_id, [img_path,] ) 21 | else: 22 | meta_data[template][1].append(img_path) 23 | return meta_data 24 | 25 | def load_feat(list_file, bin_file): 26 | mid_feats = dict() 27 | with open(list_file, 'r') as list_f, open(bin_file, 'rb') as bin_f: 28 | (data_num, feat_dim) = st.unpack('ii', bin_f.read(8)) 29 | for line in list_f: 30 | record = line.strip().split('/') 31 | img_name = '/'.join(record[-2:]) 32 | feat = np.array(st.unpack('f'*feat_dim, bin_f.read(4*feat_dim))) 33 | mid_feats[img_name] = feat 34 | return mid_feats, feat_dim 35 | 36 | def update_meta_data(meta_data, feats, feat_dim): 37 | new_meta_data = dict() 38 | for template in meta_data.keys(): 39 | class_id, img_names = meta_data[template] 40 | feat = np.zeros(feat_dim) 41 | feat_num = 0 42 | for img_name in img_names: 43 | if img_name in feats: 44 | feat += feats[img_name] 45 | feat_num += 1 46 | if feat_num > 0: feat /= feat_num 47 | if feat_num > 0: new_meta_data[template] = (class_id, feat) 48 | return new_meta_data 49 | 50 | # calc_tar2{{{ 51 | def calc_tar2(fpr, tpr, method=0): 52 | if method == 0: 53 | min_dis, tar2 = 100.0, 1.0 54 | for i in range(fpr.size): 55 | if(fpr[i]+tpr[i] > 1.0): 56 | break 57 | mid_res = abs(fpr[i]+tpr[i]-1.0) 58 | if(mid_res < min_dis): 59 | min_dis = mid_res 60 | tar2 = fpr[i] 61 | return tar2 62 | else: 63 | f = lambda x: np.interp(x, fpr, tpr)+x-1 64 | return fsolve(f, 0.0) 65 | #}}} 66 | 67 | def eval_roc(pair_file, meta_data, split): 68 | labels, scores = [], [] 69 | with open(pair_file, 'r') as in_f: 70 | for line in in_f: 71 | record = line.strip().split(',') 72 | temp1, temp2 = int(record[0]),int(record[1]) 73 | if not(temp1 in meta_data and temp2 in meta_data): 74 | continue 75 | info1, info2 = meta_data[temp1], meta_data[temp2] 76 | label = int(info1[0] == info2[0]) 77 | score = 1-spatial.distance.cosine(info1[1],info2[1]) 78 | labels.append(label) 79 | scores.append(score) 80 | fpr, tpr, thresholds = metrics.roc_curve(labels, scores) 81 | idx1 = bisect.bisect_left(fpr,0.01) 82 | tar1 = tpr[ bisect.bisect_left(fpr,0.001) ] 83 | tar2 = tpr[ bisect.bisect_left(fpr,0.01) ] 84 | return tar1, tar2 85 | 86 | 87 | def test_verify(model_type): 88 | 89 | #model_type = 'resnet18' 90 | protocol_dir = '../../data/IJBA/protocol_11' 91 | align_img_dir = '../../data/IJBA/align_image_11' 92 | tar1s, tar2s = [],[] 93 | split_num = 10 94 | for split in range(1, split_num+1): 95 | 96 | # load meta data first 97 | split_protocol_dir = os.path.join(protocol_dir, 'split'+str(split)) 98 | meta_file = os.path.join(split_protocol_dir, 'verify_metadata_{}.csv'.format(split)) 99 | meta_data = load_meta_data(meta_file) 100 | 101 | # load extract feat 102 | feats = dict() 103 | feat_dim = 0 104 | split_img_dir = os.path.join(align_img_dir, 'ijb_a_11_align_split{}'.format(split)) 105 | for img_type in ['frame', 'img']: 106 | list_file = os.path.join(split_img_dir, '{}_list.txt'.format(img_type)) 107 | bin_file = os.path.join(split_img_dir, '{}_{}_feat.bin'.format(model_type,img_type)) 108 | mid_feats, feat_dim = load_feat(list_file, bin_file) 109 | feats.update(mid_feats) 110 | 111 | # update meta data 112 | meta_data = update_meta_data(meta_data, feats, feat_dim) 113 | 114 | # eval roc 115 | pair_file = os.path.join(split_protocol_dir, 'verify_comparisons_{}.csv'.format(split)) 116 | tar1, tar2 = eval_roc(pair_file, meta_data, split) 117 | tar1s.append(tar1) 118 | tar2s.append(tar2) 119 | print('split {}, tar1: {}, tar2: {}'.format(split,tar1,tar2)) 120 | print('tar1: {} +/- {}'.format(np.mean(tar1s), np.std(tar1s))) 121 | print('tar2: {} +/- {}'.format(np.mean(tar2s), np.std(tar2s))) 122 | 123 | return np.mean(tar1s), np.std(tar1s), np.mean(tar2s), np.std(tar2s) 124 | 125 | if __name__ == '__main__': 126 | test_verify(model_type='resnet18') 127 | -------------------------------------------------------------------------------- /src/end2end/ResNet.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import math 3 | import torch.utils.model_zoo as model_zoo 4 | import torch.nn.functional as F 5 | 6 | 7 | __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 8 | 'resnet152'] 9 | 10 | 11 | model_urls = { 12 | 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth', 13 | 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth', 14 | 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth', 15 | 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth', 16 | 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', 17 | } 18 | 19 | 20 | def conv3x3(in_planes, out_planes, stride=1): 21 | "3x3 convolution with padding" 22 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, 23 | padding=1, bias=False) 24 | 25 | 26 | class BasicBlock(nn.Module): 27 | expansion = 1 28 | 29 | def __init__(self, inplanes, planes, stride=1, downsample=None): 30 | super(BasicBlock, self).__init__() 31 | self.conv1 = conv3x3(inplanes, planes, stride) 32 | self.bn1 = nn.BatchNorm2d(planes) 33 | self.relu = nn.ReLU(inplace=True) 34 | self.conv2 = conv3x3(planes, planes) 35 | self.bn2 = nn.BatchNorm2d(planes) 36 | self.downsample = downsample 37 | self.stride = stride 38 | 39 | def forward(self, x): 40 | residual = x 41 | 42 | out = self.conv1(x) 43 | out = self.bn1(out) 44 | out = self.relu(out) 45 | 46 | out = self.conv2(out) 47 | out = self.bn2(out) 48 | 49 | if self.downsample is not None: 50 | residual = self.downsample(x) 51 | 52 | out += residual 53 | out = self.relu(out) 54 | 55 | return out 56 | 57 | 58 | class Bottleneck(nn.Module): 59 | expansion = 4 60 | 61 | def __init__(self, inplanes, planes, stride=1, downsample=None): 62 | super(Bottleneck, self).__init__() 63 | self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) 64 | self.bn1 = nn.BatchNorm2d(planes) 65 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, 66 | padding=1, bias=False) 67 | self.bn2 = nn.BatchNorm2d(planes) 68 | self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) 69 | self.bn3 = nn.BatchNorm2d(planes * 4) 70 | self.relu = nn.ReLU(inplace=True) 71 | self.downsample = downsample 72 | self.stride = stride 73 | 74 | def forward(self, x): 75 | residual = x 76 | 77 | out = self.conv1(x) 78 | out = self.bn1(out) 79 | out = self.relu(out) 80 | 81 | out = self.conv2(out) 82 | out = self.bn2(out) 83 | out = self.relu(out) 84 | 85 | out = self.conv3(out) 86 | out = self.bn3(out) 87 | 88 | if self.downsample is not None: 89 | residual = self.downsample(x) 90 | 91 | out += residual 92 | out = self.relu(out) 93 | 94 | return out 95 | 96 | 97 | class ResNet(nn.Module): 98 | 99 | def __init__(self, block, layers, num_classes=1000, end2end=True): 100 | self.inplanes = 64 101 | self.end2end = end2end 102 | super(ResNet, self).__init__() 103 | self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, 104 | bias=False) 105 | self.bn1 = nn.BatchNorm2d(64) 106 | self.relu = nn.ReLU(inplace=True) 107 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) 108 | self.layer1 = self._make_layer(block, 64, layers[0]) 109 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2) 110 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2) 111 | self.layer4 = self._make_layer(block, 512, layers[3], stride=2) 112 | self.avgpool = nn.AvgPool2d(7) 113 | self.feature = nn.Linear(512 * block.expansion, 256) 114 | if self.end2end: 115 | self.fc1 = nn.Linear(256, 256) 116 | self.fc2 = nn.Linear(256, 256) 117 | self.fc = nn.Linear(256, num_classes) 118 | 119 | for m in self.modules(): 120 | if isinstance(m, nn.Conv2d): 121 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 122 | m.weight.data.normal_(0, math.sqrt(2. / n)) 123 | elif isinstance(m, nn.BatchNorm2d): 124 | m.weight.data.fill_(1) 125 | m.bias.data.zero_() 126 | 127 | def _make_layer(self, block, planes, blocks, stride=1): 128 | downsample = None 129 | if stride != 1 or self.inplanes != planes * block.expansion: 130 | downsample = nn.Sequential( 131 | nn.Conv2d(self.inplanes, planes * block.expansion, 132 | kernel_size=1, stride=stride, bias=False), 133 | nn.BatchNorm2d(planes * block.expansion), 134 | ) 135 | 136 | layers = [] 137 | layers.append(block(self.inplanes, planes, stride, downsample)) 138 | self.inplanes = planes * block.expansion 139 | for i in range(1, blocks): 140 | layers.append(block(self.inplanes, planes)) 141 | 142 | return nn.Sequential(*layers) 143 | 144 | def forward(self, x, yaw): 145 | x = self.conv1(x) 146 | x = self.bn1(x) 147 | x = self.relu(x) 148 | x = self.maxpool(x) 149 | 150 | x = self.layer1(x) 151 | x = self.layer2(x) 152 | x = self.layer3(x) 153 | x = self.layer4(x) 154 | 155 | x = self.avgpool(x) 156 | x = x.view(x.size(0), -1) 157 | mid_feature = self.feature(x) 158 | # code above is the original arch of resnet 159 | 160 | if self.end2end: 161 | raw_feature = self.fc1(mid_feature) 162 | raw_feature = self.relu(raw_feature) 163 | raw_feature = self.fc2(raw_feature) 164 | raw_feature = self.relu(raw_feature) 165 | 166 | yaw = yaw.view(yaw.size(0),1) 167 | yaw = yaw.expand_as(raw_feature) 168 | 169 | feature = yaw * raw_feature + mid_feature 170 | else: 171 | feature = mid_feature 172 | 173 | 174 | feature = F.dropout(feature, p=0.7, training=self.training) 175 | pred_score = self.fc(feature) 176 | 177 | return pred_score 178 | 179 | 180 | def resnet18(pretrained=False, **kwargs): 181 | """Constructs a ResNet-18 model. 182 | 183 | Args: 184 | pretrained (bool): If True, returns a model pre-trained on ImageNet 185 | """ 186 | model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs) 187 | if pretrained: 188 | model.load_state_dict(model_zoo.load_url(model_urls['resnet18'])) 189 | return model 190 | 191 | 192 | def resnet34(pretrained=False, **kwargs): 193 | """Constructs a ResNet-34 model. 194 | 195 | Args: 196 | pretrained (bool): If True, returns a model pre-trained on ImageNet 197 | """ 198 | model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs) 199 | if pretrained: 200 | model.load_state_dict(model_zoo.load_url(model_urls['resnet34'])) 201 | return model 202 | 203 | 204 | def resnet50(pretrained=False, **kwargs): 205 | """Constructs a ResNet-50 model. 206 | 207 | Args: 208 | pretrained (bool): If True, returns a model pre-trained on ImageNet 209 | """ 210 | model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs) 211 | if pretrained: 212 | model.load_state_dict(model_zoo.load_url(model_urls['resnet50'])) 213 | return model 214 | 215 | 216 | def resnet101(pretrained=False, **kwargs): 217 | """Constructs a ResNet-101 model. 218 | 219 | Args: 220 | pretrained (bool): If True, returns a model pre-trained on ImageNet 221 | """ 222 | model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs) 223 | if pretrained: 224 | model.load_state_dict(model_zoo.load_url(model_urls['resnet101'])) 225 | return model 226 | 227 | 228 | def resnet152(pretrained=False, **kwargs): 229 | """Constructs a ResNet-152 model. 230 | 231 | Args: 232 | pretrained (bool): If True, returns a model pre-trained on ImageNet 233 | """ 234 | model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs) 235 | if pretrained: 236 | model.load_state_dict(model_zoo.load_url(model_urls['resnet152'])) 237 | return model 238 | -------------------------------------------------------------------------------- /src/end2end/main.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os,sys,shutil 3 | import time 4 | 5 | import torch 6 | import torch.nn as nn 7 | import torch.nn.parallel 8 | import torch.backends.cudnn as cudnn 9 | import torch.optim 10 | import torch.utils.data 11 | import torch.nn.functional as F 12 | #import transforms 13 | import torchvision.datasets as datasets 14 | import torchvision.models as models 15 | import torchvision.transforms as transforms 16 | 17 | from ResNet import resnet18, resnet50, resnet101 18 | from selfDefine import MsCelebDataset, CaffeCrop 19 | 20 | 21 | model_names = sorted(name for name in models.__dict__ 22 | if name.islower() and not name.startswith("__") 23 | and callable(models.__dict__[name])) 24 | 25 | 26 | parser = argparse.ArgumentParser(description='PyTorch CelebA Training') 27 | parser.add_argument('--img_dir', metavar='DIR', default='', help='path to dataset') 28 | parser.add_argument('--arch', '-a', metavar='ARCH', default='resnet50', choices=model_names, 29 | help='model architecture: ' + 30 | ' | '.join(model_names) + 31 | ' (default: alexnet)') 32 | parser.add_argument('-j', '--workers', default=16, type=int, metavar='N', 33 | help='number of data loading workers (default: 4)') 34 | parser.add_argument('--epochs', default=60, type=int, metavar='N', 35 | help='number of total epochs to run') 36 | parser.add_argument('--start-epoch', default=0, type=int, metavar='N', 37 | help='manual epoch number (useful on restarts)') 38 | parser.add_argument('-b', '--batch-size', default=256, type=int, 39 | metavar='N', help='mini-batch size (default: 256)') 40 | parser.add_argument('--lr', '--learning-rate', default=0.1, type=float, 41 | metavar='LR', help='initial learning rate') 42 | parser.add_argument('--momentum', default=0.9, type=float, metavar='M', 43 | help='momentum') 44 | parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float, 45 | metavar='W', help='weight decay (default: 1e-4)') 46 | parser.add_argument('--print-freq', '-p', default=10, type=int, 47 | metavar='N', help='print frequency (default: 10)') 48 | parser.add_argument('--resume', default='', type=str, metavar='PATH', 49 | help='path to latest checkpoint (default: none)') 50 | parser.add_argument('--pretrained', default='', type=str, metavar='PATH', 51 | help='path to latest checkpoint (default: none)') 52 | parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true', 53 | help='evaluate model on validation set') 54 | parser.add_argument('--model_dir','-m', default='./model', type=str) 55 | parser.add_argument('--end2end', action='store_true',\ 56 | help='if true, using end2end with dream block, else, using naive architecture') 57 | 58 | best_prec1 = 0 59 | 60 | 61 | def main(): 62 | global args, best_prec1 63 | args = parser.parse_args() 64 | 65 | print('img_dir:', args.img_dir) 66 | print('end2end?:', args.end2end) 67 | 68 | # load data and prepare dataset 69 | train_list_file = '../../data/msceleb/train_list.txt' 70 | train_label_file = '../../data/msceleb/train_label.txt' 71 | caffe_crop = CaffeCrop('train') 72 | train_dataset = MsCelebDataset(args.img_dir, train_list_file, train_label_file, 73 | transforms.Compose([caffe_crop,transforms.ToTensor()])) 74 | train_loader = torch.utils.data.DataLoader( 75 | train_dataset, 76 | batch_size=args.batch_size, shuffle=True, 77 | num_workers=args.workers, pin_memory=True) 78 | 79 | caffe_crop = CaffeCrop('test') 80 | val_list_file = '../../data/msceleb/test_list.txt' 81 | val_label_file = '../../data/msceleb/test_label.txt' 82 | val_dataset = MsCelebDataset(args.img_dir, val_list_file, val_label_file, 83 | transforms.Compose([caffe_crop,transforms.ToTensor()])) 84 | val_loader = torch.utils.data.DataLoader( 85 | val_dataset, 86 | batch_size=args.batch_size, shuffle=False, 87 | num_workers=args.workers, pin_memory=True) 88 | 89 | assert(train_dataset.max_label == val_dataset.max_label) 90 | class_num = train_dataset.max_label + 1 91 | 92 | print('class_num: ',class_num) 93 | 94 | # prepare model 95 | model = None 96 | assert(args.arch in ['resnet18','resnet50','resnet101']) 97 | if args.arch == 'resnet18': 98 | model = resnet18(pretrained=False, num_classes=class_num, end2end=args.end2end) 99 | if args.arch == 'resnet50': 100 | model = resnet50(pretrained=False, num_classes=class_num, end2end=args.end2end) 101 | if args.arch == 'resnet101': 102 | model = resnet101(pretrained=False, num_classes=class_num, end2end=args.end2end) 103 | model = torch.nn.DataParallel(model).cuda() 104 | 105 | 106 | # define loss function (criterion) and optimizer 107 | criterion = nn.CrossEntropyLoss().cuda() 108 | 109 | optimizer = torch.optim.SGD(model.parameters(), args.lr, 110 | momentum=args.momentum, 111 | weight_decay=args.weight_decay) 112 | 113 | # optionally resume from a checkpoint 114 | if args.pretrained: 115 | checkpoint = torch.load(args.pretrained) 116 | pretrained_state_dict = checkpoint['state_dict'] 117 | model_state_dict = model.state_dict() 118 | 119 | for key in pretrained_state_dict: 120 | model_state_dict[key] = pretrained_state_dict[key] 121 | model.load_state_dict(model_state_dict) 122 | 123 | 124 | if args.resume: 125 | if os.path.isfile(args.resume): 126 | print("=> loading checkpoint '{}'".format(args.resume)) 127 | checkpoint = torch.load(args.resume) 128 | args.start_epoch = checkpoint['epoch'] 129 | best_prec1 = checkpoint['best_prec1'] 130 | model.load_state_dict(checkpoint['state_dict']) 131 | optimizer.load_state_dict(checkpoint['optimizer']) 132 | print("=> loaded checkpoint '{}' (epoch {})" 133 | .format(args.resume, checkpoint['epoch'])) 134 | else: 135 | print("=> no checkpoint found at '{}'".format(args.resume)) 136 | 137 | cudnn.benchmark = True 138 | 139 | 140 | 141 | if args.evaluate: 142 | validate(val_loader, model, criterion) 143 | return 144 | 145 | for epoch in range(args.start_epoch, args.epochs): 146 | adjust_learning_rate(optimizer, epoch) 147 | 148 | # train for one epoch 149 | train(train_loader, model, criterion, optimizer, epoch) 150 | 151 | # evaluate on validation set 152 | prec1 = validate(val_loader, model, criterion) 153 | 154 | # remember best prec@1 and save checkpoint 155 | is_best = prec1 > best_prec1 156 | best_prec1 = max(prec1, best_prec1) 157 | save_checkpoint({ 158 | 'epoch': epoch + 1, 159 | 'arch': args.arch, 160 | 'state_dict': model.state_dict(), 161 | 'best_prec1': best_prec1, 162 | 'optimizer' : optimizer.state_dict(), 163 | }, is_best) 164 | 165 | 166 | def train(train_loader, model, criterion, optimizer, epoch): 167 | batch_time = AverageMeter() 168 | data_time = AverageMeter() 169 | cla_losses = AverageMeter() 170 | yaw_losses = AverageMeter() 171 | losses = AverageMeter() 172 | top1 = AverageMeter() 173 | top5 = AverageMeter() 174 | 175 | # switch to train mode 176 | model.train() 177 | 178 | end = time.time() 179 | for i, (input, target, yaw) in enumerate(train_loader): 180 | # measure data loading time 181 | data_time.update(time.time() - end) 182 | 183 | target = target.cuda(async=True) 184 | yaw = yaw.float().cuda(async=True) 185 | input_var = torch.autograd.Variable(input) 186 | yaw_var = torch.autograd.Variable(yaw) 187 | target_var = torch.autograd.Variable(target) 188 | 189 | # compute output 190 | pred_score = model(input_var, yaw_var) 191 | 192 | loss = criterion(pred_score, target_var) 193 | 194 | # measure accuracy and record loss 195 | prec1, prec5 = accuracy(pred_score.data, target, topk=(1, 5)) 196 | losses.update(loss.data[0], input.size(0)) 197 | top1.update(prec1[0], input.size(0)) 198 | top5.update(prec5[0], input.size(0)) 199 | 200 | # compute gradient and do SGD step 201 | optimizer.zero_grad() 202 | loss.backward() 203 | optimizer.step() 204 | 205 | # measure elapsed time 206 | batch_time.update(time.time() - end) 207 | end = time.time() 208 | 209 | if i % args.print_freq == 0: 210 | print('Epoch: [{0}][{1}/{2}]\t' 211 | 'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' 212 | 'Data {data_time.val:.3f} ({data_time.avg:.3f})\t' 213 | 'Loss {loss.val:.4f} ({loss.avg:.4f})\t' 214 | 'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t' 215 | 'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format( 216 | epoch, i, len(train_loader), batch_time=batch_time, 217 | data_time=data_time, loss=losses, top1=top1, top5=top5)) 218 | 219 | 220 | def validate(val_loader, model, criterion): 221 | batch_time = AverageMeter() 222 | cla_losses = AverageMeter() 223 | yaw_losses = AverageMeter() 224 | losses = AverageMeter() 225 | top1 = AverageMeter() 226 | top5 = AverageMeter() 227 | 228 | # switch to evaluate mode 229 | model.eval() 230 | 231 | end = time.time() 232 | for i, (input, target, yaw) in enumerate(val_loader): 233 | target = target.cuda(async=True) 234 | yaw = yaw.float().cuda(async=True) 235 | input_var = torch.autograd.Variable(input, volatile=True) 236 | target_var = torch.autograd.Variable(target, volatile=True) 237 | yaw_var = torch.autograd.Variable(yaw) 238 | 239 | # compute output 240 | pred_score = model(input_var, yaw_var) 241 | 242 | loss = criterion(pred_score, target_var) 243 | 244 | # measure accuracy and record loss 245 | prec1, prec5 = accuracy(pred_score.data, target, topk=(1, 5)) 246 | losses.update(loss.data[0], input.size(0)) 247 | top1.update(prec1[0], input.size(0)) 248 | top5.update(prec5[0], input.size(0)) 249 | 250 | # measure elapsed time 251 | batch_time.update(time.time() - end) 252 | end = time.time() 253 | 254 | if i % args.print_freq == 0: 255 | print('Test: [{0}/{1}]\t' 256 | 'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' 257 | 'Loss {loss.val:.4f} ({loss.avg:.4f})\t' 258 | 'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t' 259 | 'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format( 260 | i, len(val_loader), batch_time=batch_time, loss=losses, 261 | top1=top1, top5=top5)) 262 | 263 | print(' * Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f}' 264 | .format(top1=top1, top5=top5)) 265 | 266 | return top1.avg 267 | 268 | 269 | def save_checkpoint(state, is_best, filename='checkpoint.pth.tar'): 270 | 271 | full_filename = os.path.join(args.model_dir, filename) 272 | full_bestname = os.path.join(args.model_dir, 'model_best.pth.tar') 273 | torch.save(state, full_filename) 274 | epoch_num = state['epoch'] 275 | if epoch_num%5==0 and epoch_num>=0: 276 | torch.save(state, full_filename.replace('checkpoint','checkpoint_'+str(epoch_num))) 277 | if is_best: 278 | shutil.copyfile(full_filename, full_bestname) 279 | 280 | 281 | class AverageMeter(object): 282 | """Computes and stores the average and current value""" 283 | def __init__(self): 284 | self.reset() 285 | 286 | def reset(self): 287 | self.val = 0 288 | self.avg = 0 289 | self.sum = 0 290 | self.count = 0 291 | 292 | def update(self, val, n=1): 293 | self.val = val 294 | self.sum += val * n 295 | self.count += n 296 | self.avg = self.sum / self.count 297 | 298 | 299 | def adjust_learning_rate(optimizer, epoch): 300 | """Sets the learning rate to the initial LR decayed by 10 every 30 epochs""" 301 | #lr = args.lr * (0.1 ** (epoch // 30)) 302 | if epoch in [int(args.epochs*0.8), int(args.epochs*0.9), int(args.epochs*0.95)]: 303 | for param_group in optimizer.param_groups: 304 | param_group['lr'] *= 0.1 305 | 306 | 307 | def accuracy(output, target, topk=(1,)): 308 | """Computes the precision@k for the specified values of k""" 309 | maxk = max(topk) 310 | batch_size = target.size(0) 311 | 312 | _, pred = output.topk(maxk, 1, True, True) 313 | pred = pred.t() 314 | correct = pred.eq(target.view(1, -1).expand_as(pred)) 315 | 316 | res = [] 317 | for k in topk: 318 | correct_k = correct[:k].view(-1).float().sum(0) 319 | res.append(correct_k.mul_(100.0 / batch_size)) 320 | return res 321 | 322 | 323 | if __name__ == '__main__': 324 | main() 325 | -------------------------------------------------------------------------------- /src/end2end/selfDefine.py: -------------------------------------------------------------------------------- 1 | import os, sys, shutil 2 | import random as rd 3 | 4 | from PIL import Image 5 | import numpy as np 6 | 7 | import torch 8 | import torch.nn.functional as F 9 | import torch.utils.data as data 10 | from torch.autograd import Variable 11 | from torch.nn.modules.loss import _WeightedLoss 12 | 13 | def load_imgs(img_dir, image_list_file, label_file): 14 | imgs = list() 15 | max_label = 0 16 | with open(image_list_file, 'r') as imf: 17 | with open(label_file, 'r') as laf: 18 | record = laf.readline().strip().split() 19 | total_num, label_num = int(record[0]), int(record[1]) 20 | for line in imf: 21 | img_path = os.path.join(img_dir, line.strip()) 22 | record = laf.readline().strip().split() 23 | label,yaw = int(record[0]), float(record[1]) 24 | max_label = max(max_label, label) 25 | imgs.append((img_path, label, yaw)) 26 | assert(total_num == len(imgs)) 27 | assert(label_num == max_label+1) 28 | return imgs, max_label 29 | 30 | 31 | class MsCelebDataset(data.Dataset): 32 | def __init__(self, img_dir, image_list_file, label_file, transform=None): 33 | self.imgs, self.max_label = load_imgs(img_dir, image_list_file, label_file) 34 | self.transform = transform 35 | 36 | def __getitem__(self, index): 37 | path, target, yaw = self.imgs[index] 38 | img = Image.open(path).convert("RGB") 39 | if self.transform is not None: 40 | img = self.transform(img) 41 | return img, target, yaw 42 | 43 | def __len__(self): 44 | return len(self.imgs) 45 | 46 | 47 | 48 | 49 | class CaffeCrop(object): 50 | """ 51 | This class take the same behavior as sensenet 52 | """ 53 | def __init__(self, phase): 54 | assert(phase=='train' or phase=='test') 55 | self.phase = phase 56 | 57 | def __call__(self, img): 58 | # pre determined parameters 59 | final_size = 224 60 | final_width = final_height = final_size 61 | crop_size = 110 62 | crop_height = crop_width = crop_size 63 | crop_center_y_offset = 15 64 | crop_center_x_offset = 0 65 | if self.phase == 'train': 66 | scale_aug = 0.02 67 | trans_aug = 0.01 68 | else: 69 | scale_aug = 0.0 70 | trans_aug = 0.0 71 | 72 | # computed parameters 73 | randint = rd.randint 74 | scale_height_diff = (randint(0,1000)/500-1)*scale_aug 75 | crop_height_aug = crop_height*(1+scale_height_diff) 76 | scale_width_diff = (randint(0,1000)/500-1)*scale_aug 77 | crop_width_aug = crop_width*(1+scale_width_diff) 78 | 79 | 80 | trans_diff_x = (randint(0,1000)/500-1)*trans_aug 81 | trans_diff_y = (randint(0,1000)/500-1)*trans_aug 82 | 83 | 84 | center = ((img.width/2 + crop_center_x_offset)*(1+trans_diff_x), 85 | (img.height/2 + crop_center_y_offset)*(1+trans_diff_y)) 86 | 87 | 88 | if center[0] < crop_width_aug/2: 89 | crop_width_aug = center[0]*2-0.5 90 | if center[1] < crop_height_aug/2: 91 | crop_height_aug = center[1]*2-0.5 92 | if (center[0]+crop_width_aug/2) >= img.width: 93 | crop_width_aug = (img.width-center[0])*2-0.5 94 | if (center[1]+crop_height_aug/2) >= img.height: 95 | crop_height_aug = (img.height-center[1])*2-0.5 96 | 97 | crop_box = (center[0]-crop_width_aug/2, center[1]-crop_height_aug/2, 98 | center[0]+crop_width_aug/2, center[1]+crop_width_aug/2) 99 | 100 | mid_img = img.crop(crop_box) 101 | res_img = mid_img.resize( (final_width, final_height) ) 102 | return res_img 103 | -------------------------------------------------------------------------------- /src/end2end/slurm.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This code could allocate the GPU automatically. 3 | Please run this code in Python 3 only!!! 4 | ''' 5 | 6 | import os, sys, shutil 7 | import subprocess as sbp 8 | from datetime import datetime 9 | 10 | def get_all_gpu(): 11 | all_gpus = list() 12 | nvidia = sbp.run(['nvidia-smi', '-L'], stdout=sbp.PIPE) 13 | for line in nvidia.stdout.decode('utf-8').split('\n')[:-1]: 14 | all_gpus.append(int(line[4])) 15 | return all_gpus 16 | 17 | def get_idle_gpu(all_gpus): 18 | idle_gpus = list() 19 | for gpu_id in all_gpus: 20 | nvidia = sbp.run(['nvidia-smi', '-i', str(gpu_id)], stdout=sbp.PIPE) 21 | output = nvidia.stdout.decode('utf-8') 22 | if output.find('No running processes found')>=0: 23 | idle_gpus.append(gpu_id) 24 | return idle_gpus 25 | 26 | def allocate_gpu(process_name, gpu_num): 27 | all_gpus = get_all_gpu() 28 | 29 | idle_gpus = get_idle_gpu(all_gpus) 30 | 31 | if gpu_num > len(idle_gpus): 32 | print('GPU is not enough') 33 | else: 34 | allocate_gpus = idle_gpus[:gpu_num] 35 | print('GPU {} is allocated'.format(','.join(map(str,allocate_gpus)))) 36 | CUDA_command = 'CUDA_VISIBLE_DEVICES='+','.join(map(str,allocate_gpus)) 37 | exec_command = CUDA_command + ' ' + process_name 38 | #print(exec_command) 39 | os.system(exec_command) 40 | 41 | 42 | 43 | if __name__ == '__main__': 44 | 45 | if not os.path.exists('./log'): 46 | os.mkdir('./log/') 47 | log_file = './log/{}.log'.format(datetime.now().strftime('%m_%d_%H_%M')) 48 | 49 | process_name = 'python main.py 2>&1 | tee {}'.format(log_file) 50 | gpu_num = 8 51 | 52 | allocate_gpu(process_name, gpu_num) 53 | -------------------------------------------------------------------------------- /src/end2end/train.sh: -------------------------------------------------------------------------------- 1 | model_dir=model 2 | if [ ! -d $model_dir ]; then 3 | mkdir $model_dir 4 | fi 5 | log_dir=log 6 | if [ ! -d $log_dir ]; then 7 | mkdir $log_dir 8 | fi 9 | 10 | curr_date=$(date +'%m_%d_%H_%M') 11 | log_file="./log/$curr_date.log" 12 | 13 | # train the model with GPUs 0, 1, 2, and 3 14 | CUDA_VISIBLE_DEVICES=0,1,2,3 python main.py \ 15 | --end2end --img_dir ../../data/msceleb/image \ 16 | 2>&1 | tee $log_file 17 | -------------------------------------------------------------------------------- /src/preprocess/align_cfp.sh: -------------------------------------------------------------------------------- 1 | # prepare CFP data first 2 | dst_dir=../../data/CFP 3 | list_file_name=/mnt/SSD/rongyu/data/cfp/cfp_dataset/Data/list_name.txt 4 | dst_file_name=align_img_list.txt 5 | python pre_cfp_data.py $dst_dir $list_file_name $dst_file_name 6 | 7 | 8 | # align CFP image 9 | pnp_file=pnp.txt 10 | image_prefix=/mnt/SSD/rongyu/data/cfp/cfp_dataset/ 11 | alignment_file=cfp_alignment.txt 12 | aligned_img_file=align_img_list.txt 13 | pose_file=estimate_pose.txt 14 | 15 | export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib 16 | ./test_process_align $pnp_file $image_prefix $alignment_file $aligned_img_file $pose_file 17 | -------------------------------------------------------------------------------- /src/preprocess/pnp.txt: -------------------------------------------------------------------------------- 1 | -55.1199760000000 -20.0425830000000 9.26549820000000 2 | -36.0639000000000 -23.7608930000000 0.728753030000000 3 | -16.0913370000000 -19.2479590000000 -4.81830070000000 4 | 15.4050260000000 -19.1353490000000 -5.54768660000000 5 | 35.4269750000000 -23.3949600000000 -0.720136230000000 6 | 54.4413220000000 -19.5296100000000 7.50984000000000 7 | -45.9189300000000 0.570106390000000 5.00845860000000 8 | -18.5672740000000 1.87615780000000 -0.0870856420000000 9 | 19.2185540000000 2.11073450000000 -0.110125390000000 10 | 45.9701840000000 0.671661260000000 4.28201870000000 11 | -18.8383200000000 40.6377790000000 -8.96802620000000 12 | 0.102908830000000 47.0601810000000 -15.3461590000000 13 | 19.2440450000000 41.0716250000000 -8.78937150000000 14 | 0.280321930000000 62.4372220000000 -14.0781120000000 15 | 0.295113470000000 71.2905200000000 -12.2051410000000 16 | 0.178925720000000 81.1921310000000 -11.5306380000000 17 | -32.2462040000000 -2.73430400000000e-06 1.43103760000000e-06 18 | 32.2462120000000 -2.71543260000000e-06 1.47307650000000e-06 19 | -0.393499640000000 33.9113660000000 -21.3970640000000 20 | -23.9587060000000 72.2165150000000 0.314607980000000 21 | 24.4262730000000 72.0572280000000 -0.314598620000000 22 | 23 | 444.963900000000 203.439600000000 24 | 445.093200000000 205.134100000000 25 | 444.159900000000 206.514600000000 26 | 447.642100000000 207.561100000000 27 | 457.300300000000 208.527700000000 28 | 468.232000000000 214.503900000000 29 | 443.422900000000 220.752000000000 30 | 444.431000000000 222.945000000000 31 | 452.958800000000 227.616000000000 32 | 464.050000000000 230.028000000000 33 | 423.289400000000 258.213000000000 34 | 432.440600000000 265.409000000000 35 | 442.742900000000 262.585000000000 36 | 430.786800000000 276.878000000000 37 | 431.181300000000 286.866000000000 38 | 431.004900000000 297.214000000000 39 | 443.853400000000 222.072000000000 40 | 457.082000000000 228.419000000000 41 | 423.954300000000 253.167000000000 42 | 434.671200000000 285.710000000000 43 | 447.233700000000 290.635000000000 44 | 45 | 312.914000000000 82.8232000000000 46 | 332.733000000000 85.8278000000000 47 | 348.926000000000 91.9491000000000 48 | 364.921000000000 93.5363000000000 49 | 372.699000000000 90.0181000000000 50 | 378.209000000000 86.7211000000000 51 | 313.862000000000 98.0127000000000 52 | 333.290000000000 102.212000000000 53 | 355.834000000000 103.272000000000 54 | 364.405000000000 101.379000000000 55 | 319.731000000000 133.602000000000 56 | 331.934000000000 144.645000000000 57 | 347.972000000000 143.384000000000 58 | 323.027000000000 151.133000000000 59 | 318.324000000000 156.543000000000 60 | 312.870000000000 163.397000000000 61 | 323.990000000000 101.823000000000 62 | 360.150000000000 103.176000000000 63 | 346.616000000000 141.085000000000 64 | 295.006000000000 146.603000000000 65 | 323.828000000000 152.583000000000 66 | -------------------------------------------------------------------------------- /src/preprocess/pre_cfp_data.py: -------------------------------------------------------------------------------- 1 | import os, sys, shutil 2 | 3 | if __name__ == '__main__': 4 | 5 | assert(len(sys.argv)==4) 6 | 7 | dst_dir = sys.argv[1] # ../data/cfp 8 | if not os.path.exists(dst_dir): 9 | os.makedirs(dst_dir) 10 | 11 | celeb_name_file = sys.argv[2] # /ssd/data/cfp/cfp_dataset/Data/list_name.txt 12 | with open(celeb_name_file, 'r') as in_f: 13 | name_mapping = {idx+1:'_'.join(line.strip().split()) \ 14 | for idx, line in enumerate(in_f)} 15 | 16 | dst_data_dir = os.path.join(dst_dir, 'data') 17 | if not os.path.exists(dst_data_dir): 18 | os.mkdir(dst_data_dir) 19 | 20 | 21 | align_img_list = sys.argv[3] 22 | with open(align_img_list, 'w') as out_f: 23 | 24 | celeb_num = len(name_mapping) 25 | for i in range(celeb_num): 26 | 27 | celeb_name = name_mapping[i+1] 28 | 29 | sub_dir = os.path.join(dst_data_dir, celeb_name) 30 | if not os.path.exists(sub_dir): 31 | os.mkdir(sub_dir) 32 | 33 | frontal_sub_dir = os.path.join(sub_dir, 'frontal') 34 | profile_sub_dir = os.path.join(sub_dir, 'profile') 35 | if not os.path.exists(frontal_sub_dir): 36 | os.mkdir(frontal_sub_dir) 37 | if not os.path.exists(profile_sub_dir): 38 | os.mkdir(profile_sub_dir) 39 | 40 | frontal_num, profile_num = 10, 4 41 | infos = [(frontal_sub_dir, frontal_num), (profile_sub_dir, profile_num)] 42 | for sub_dir_name, img_num in infos: 43 | for im_n in range(1, img_num+1): 44 | img_id = '0'+str(im_n) if len(str(im_n))<2 else str(im_n) 45 | img_name = img_id + '.jpg' 46 | out_f.write(os.path.join(sub_dir_name, img_name)+'\n') 47 | -------------------------------------------------------------------------------- /src/preprocess/test_process_align: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/penincillin/DREAM/7bea01e5b108546a2b3a809ad17cf478607c9db5/src/preprocess/test_process_align -------------------------------------------------------------------------------- /src/stitching/branch_test.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import struct as st 3 | import torch 4 | import torch.nn as nn 5 | import math 6 | from branch_util import * 7 | 8 | def load_feat(feat_file): 9 | feats = list() 10 | with open(feat_file, 'rb') as in_f: 11 | feat_num, feat_dim = st.unpack('ii', in_f.read(8)) 12 | for i in range(feat_num): 13 | feat = np.array(st.unpack('f'*feat_dim, in_f.read(4*feat_dim))) 14 | feats.append(feat) 15 | return feats 16 | 17 | if __name__ == '__main__': 18 | frontal_feat_file = '/home/sensetime/Downloads/ext_feat/data/frontal_feat_naive.bin' 19 | profile_feat_file = '/home/sensetime/Downloads/ext_feat/data/profile_feat_naive.bin' 20 | 21 | frontal_feats = np.vstack(load_feat(frontal_feat_file)) 22 | profile_feats = np.vstack(load_feat(profile_feat_file)) 23 | 24 | frontal_angles = [] 25 | profile_angles = [] 26 | with open('../../data/pose_output.txt', 'r') as fin: 27 | tmp_count = 0 28 | for line in fin: 29 | line_split = line.strip().split() 30 | if tmp_count % 14 < 10: 31 | frontal_angles.append(float(line_split[2])) 32 | else: 33 | profile_angles.append(float(line_split[2])) 34 | tmp_count = tmp_count + 1 35 | frontal_angles = np.vstack(frontal_angles) 36 | profile_angles = np.vstack(profile_angles) 37 | 38 | model = Branch(feat_dim=256) 39 | model.cuda() 40 | checkpoint = torch.load('checkpoint.pth') 41 | model.load_state_dict(checkpoint['state_dict']) 42 | model.eval() 43 | 44 | frontal_feat_file = './frontal_feat.bin' 45 | feat_dim = 256 46 | data_num = frontal_feats.shape[0] 47 | with open(frontal_feat_file, 'wb') as bin_f: 48 | bin_f.write(st.pack('ii', data_num, feat_dim)) 49 | for i in range(data_num): 50 | feat = frontal_feats[i, :].reshape([1, -1]) 51 | yaw = np.zeros([1, 1]) 52 | yaw[0,0] = norm_angle(frontal_angles[i, :]) 53 | feat = torch.autograd.Variable(torch.from_numpy(feat.astype(np.float32)), volatile=True).cuda() 54 | yaw = torch.autograd.Variable(torch.from_numpy(yaw.astype(np.float32)), volatile=True).cuda() 55 | output = model(feat, yaw) 56 | output_data = output.cpu().data.numpy() 57 | feat_num = output.size(0) 58 | for j in range(feat_num): 59 | bin_f.write(st.pack('f'*feat_dim, *tuple(output_data[j,:]))) 60 | 61 | data_num = profile_feats.shape[0] 62 | profile_feat_file = './profile_feat.bin' 63 | with open(profile_feat_file, 'wb') as bin_f: 64 | bin_f.write(st.pack('ii', data_num, feat_dim)) 65 | for i in range(data_num): 66 | feat = profile_feats[i, :].reshape([1, -1]) 67 | yaw = np.zeros([1, 1]) 68 | yaw[0,0] = norm_angle(profile_angles[i, :]) 69 | feat = torch.autograd.Variable(torch.from_numpy(feat.astype(np.float32)), volatile=True).cuda() 70 | yaw = torch.autograd.Variable(torch.from_numpy(yaw.astype(np.float32)), volatile=True).cuda() 71 | output = model(feat, yaw) 72 | output_data = output.cpu().data.numpy() 73 | feat_num = output.size(0) 74 | for j in range(feat_num): 75 | bin_f.write(st.pack('f'*feat_dim, *tuple(output_data[j,:]))) 76 | -------------------------------------------------------------------------------- /src/stitching/branch_train.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn as nn 4 | import math 5 | import random 6 | import argparse 7 | import shutil 8 | import struct as st 9 | from branch_util import * 10 | 11 | 12 | parser = argparse.ArgumentParser(description='Pytorch Branch Finetuning') 13 | parser.add_argument('-ilf', '--image-list-file', default='../../data/stitching/sample_img_list.txt', 14 | type=str, metavar='N', help='image list file') 15 | parser.add_argument('-ff', '--feat-file', default='../../data/stitching/ext_feat.bin', 16 | type=str, metavar='N', help='extracted feature') 17 | parser.add_argument('-b', '--batch-size', default=256, type=int, metavar='N', help='mini-batch size (default: 256)') 18 | parser.add_argument('-f', '--feat-len', default=256, type=int, metavar='F', help='feature length (default: 256)') 19 | parser.add_argument('--iters', default=10000, type=int, metavar='N', help='number of total iters to run') 20 | parser.add_argument('--lr', '--learning-rate', default=1e-3, type=float, metavar='LR', help='initial learning rate') 21 | parser.add_argument('--momentum', default=0.9, type=float, metavar='M', help='momentum') 22 | parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float, metavar='W', help='weight decay (default: 1e-4)') 23 | parser.add_argument('--print-freq', '-p', default=10, type=int, metavar='N', help='print frequency (default: 10)') 24 | 25 | def main(): 26 | args = parser.parse_args() 27 | # Load data 28 | img_list_file = args.image_list_file 29 | feat_file = args.feat_file 30 | 31 | train_map = get_dict(img_list_file) 32 | 33 | 34 | data = load_feat(feat_file) 35 | data = np.vstack(data) 36 | model = Branch(feat_dim=256) 37 | model.cuda() 38 | criterion = nn.MSELoss().cuda() 39 | optimizer = torch.optim.SGD(model.parameters(), args.lr, 40 | momentum=args.momentum, 41 | weight_decay=args.weight_decay) 42 | model.train() 43 | losses = AverageMeter() 44 | for iter in range(args.iters): 45 | batch_train_feat, batch_target_feat, batch_norm_angle = gen_batch(train_map, data, args.batch_size, args.feat_len) 46 | batch_train_feat = torch.autograd.Variable(torch.from_numpy(batch_train_feat.astype(np.float32))).cuda() 47 | batch_norm_angle = torch.autograd.Variable(torch.from_numpy(batch_norm_angle.astype(np.float32))).cuda() 48 | batch_target_feat = torch.autograd.Variable(torch.from_numpy(batch_target_feat.astype(np.float32))).cuda() 49 | 50 | output = model(batch_train_feat, batch_norm_angle) 51 | loss = criterion(output, batch_target_feat) 52 | losses.update(loss.data[0], loss.size(0)) 53 | 54 | # compute gradient and do SGD step 55 | optimizer.zero_grad() 56 | loss.backward() 57 | optimizer.step() 58 | 59 | if iter % args.print_freq == 0: 60 | print('Iter: [{0}/{1}]\t' 61 | 'Loss {loss.val:.4f} ({loss.avg:.4f})\t'.format( 62 | iter, args.iters, loss=losses)) 63 | 64 | torch.save({'state_dict': model.state_dict()}, 'checkpoint.pth') 65 | 66 | 67 | def get_dict(file_name): 68 | res_dict = dict() 69 | with open(file_name) as f: 70 | for idx, line in enumerate(f): 71 | record = line.strip().split() 72 | img_name, yaw, img_id = record[0],float(record[1]), idx 73 | celeb_name = img_name.split('/')[1] 74 | if celeb_name not in res_dict: 75 | res_dict[celeb_name] = (list(), list()) 76 | yaw = abs(yaw) 77 | if yaw <= 20: 78 | res_dict[celeb_name][0].append((img_name, yaw, img_id)) 79 | elif yaw >= 50: 80 | res_dict[celeb_name][1].append((img_name, yaw, img_id)) 81 | pop_list = [] 82 | for key in res_dict: 83 | if len(res_dict[key][0]) == 0 or len(res_dict[key][1]) == 0: 84 | pop_list.append(key) 85 | for key in pop_list: 86 | res_dict.pop(key) 87 | return res_dict 88 | 89 | def gen_batch(train_map, data, batch_size, feat_len): 90 | batch_train_feat = np.zeros([batch_size, feat_len]) 91 | batch_target_feat = np.zeros([batch_size, feat_len]) 92 | batch_yaw = np.zeros([batch_size, 1]) 93 | keys = train_map.keys() 94 | for i in range(batch_size): 95 | this_key = random.sample(keys, 1)[0] 96 | frontal_set = train_map[this_key][0] 97 | profile_set = train_map[this_key][1] 98 | frontals_index = [a[2] for a in frontal_set] 99 | frontals_feat = data[frontals_index, ...] 100 | profile_selec = random.sample(profile_set, 1) 101 | batch_train_feat[i, ...] = data[profile_selec[0][2], ...] 102 | batch_target_feat[i, ...] = np.mean(frontals_feat, axis = 0) 103 | batch_yaw[i, ...] = norm_angle(profile_selec[0][1]) 104 | return batch_train_feat, batch_target_feat, batch_yaw 105 | 106 | 107 | class AverageMeter(object): 108 | """Computes and stores the average and current value""" 109 | def __init__(self): 110 | self.reset() 111 | 112 | def reset(self): 113 | self.val = 0 114 | self.avg = 0 115 | self.sum = 0 116 | self.count = 0 117 | 118 | def update(self, val, n=1): 119 | self.val = val 120 | self.sum += val * n 121 | self.count += n 122 | self.avg = self.sum / self.count 123 | 124 | def load_feat(feat_file): 125 | feats = list() 126 | with open(feat_file, 'rb') as in_f: 127 | feat_num, feat_dim = st.unpack('ii', in_f.read(8)) 128 | for i in range(feat_num): 129 | feat = np.array(st.unpack('f'*feat_dim, in_f.read(4*feat_dim))) 130 | feats.append(feat) 131 | return feats 132 | 133 | if __name__ == '__main__': 134 | main() 135 | -------------------------------------------------------------------------------- /src/stitching/branch_util.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import struct as st 3 | import torch 4 | import torch.nn as nn 5 | import math 6 | 7 | def sigmoid(x): 8 | return 1 / (1 + math.exp(-x)) 9 | 10 | def norm_angle(angle): 11 | norm_angle = sigmoid(10 * (abs(angle) / 45.0 - 1)) 12 | return norm_angle 13 | 14 | class Branch(nn.Module): 15 | def __init__(self, feat_dim): 16 | super(Branch, self).__init__() 17 | self.fc1 = nn.Linear(feat_dim, feat_dim) 18 | self.fc2 = nn.Linear(feat_dim, feat_dim) 19 | self.relu = nn.ReLU(inplace=True) 20 | 21 | def forward(self, input, yaw): 22 | x = self.fc1(input) 23 | x = self.relu(x) 24 | x = self.fc2(x) 25 | x = self.relu(x) 26 | 27 | yaw = yaw.view(yaw.size(0),1) 28 | yaw = yaw.expand_as(x) 29 | 30 | feature = yaw * x + input 31 | return feature 32 | -------------------------------------------------------------------------------- /src/stitching/train_stitch.sh: -------------------------------------------------------------------------------- 1 | curr_date=$(date +'%m_%d_%H_%M') 2 | mkdir log 3 | log_file="./log/$curr_date.log" 4 | CUDA_VISIBLE_DEVICES=0 python branch_train.py 2>&1 | tee $log_file 5 | --------------------------------------------------------------------------------