Update readme.md

Author: cleardusk

.gitignore

@@ -1,4 +1,9 @@
 .idea/
 *.pyc
+__pycache__/
 test.data/
-
+training/snapshot/
+training/logs/
+train.configs/param_all_norm.pkl
+train.configs/param_all_norm_val.pkl
+train.configs/train_aug_120x120.list.*

benchmark.py

@@ -17,6 +17,7 @@
 
 from ddfa_utils import ToTensorGjz, NormalizeGjz, DDFATestDataset
 from params import *
+import argparse
 
 
 def _reconstruct_vertex(param, whitening=True, dense=False):
@@ -75,7 +76,7 @@ def extract_param(checkpoint_fp, root='', filelists=None, arch='mobilenet_1', nu
                 outputs.append(param_prediction)
         outputs = np.array(outputs, dtype=np.float32)
 
-    print(f'{time.time() - end: .3f}s')
+    print(f'Extracting params take {time.time() - end: .3f}s')
     return outputs
 
 
@@ -103,10 +104,8 @@ def benchmark_aflw2000_params(params):
     return _benchmark_aflw2000(outputs)
 
 
-def benchmark_pipeline():
+def benchmark_pipeline(arch, checkpoint_fp):
     device_ids = [0]
-    checkpoint_fp = 'models/phase1_wpdc_vdc.pth.tar'
-    arch = 'mobilenet_1'
 
     def aflw():
         params = extract_param(
@@ -135,7 +134,12 @@ def aflw2000():
 
 
 def main():
-    benchmark_pipeline()
+    parser = argparse.ArgumentParser(description='3DDFA Benchmark')
+    parser.add_argument('--arch', default='mobilenet_1', type=str)
+    parser.add_argument('-c', '--checkpoint-fp', default='models/phase1_wpdc_vdc.pth.tar', type=str)
+    args = parser.parse_args()
+
+    benchmark_pipeline(args.arch, args.checkpoint_fp)
 
 
 if __name__ == '__main__':

params.py

@@ -5,7 +5,12 @@
 import numpy as np
 from io_utils import _load, _numpy_to_cuda, _numpy_to_tensor, _load_gpu
 
-d = 'train.configs'
+
+def make_abs_path(d):
+    return osp.join(osp.dirname(osp.realpath(__file__)), d)
+
+
+d = make_abs_path('train.configs')
 keypoints = _load(osp.join(d, 'keypoints_sim.npy'))
 w_shp = _load(osp.join(d, 'w_shp_sim.npy'))
 w_exp = _load(osp.join(d, 'w_exp_sim.npy'))  # simplified version

readme.md

@@ -1,5 +1,40 @@
-## Face Alignment in Full Pose Range: A 3D Total Solution
+# Face Alignment in Full Pose Range: A 3D Total Solution
 
+## Introduction
 The pytorch implementation of paper [Face Alignment in Full Pose Range: A 3D Total Solution](https://arxiv.org/abs/1804.01005).
 
-This work is in progress.
+## Citation
+    @article{zhu2017face,
+      title={Face Alignment in Full Pose Range: A 3D Total Solution},
+      author={Zhu, Xiangyu and Lei, Zhen and Li, Stan Z and others},
+      journal={IEEE Transactions on Pattern Analysis and Machine Intelligence},
+      year={2017},
+      publisher={IEEE}
+    }
+
+
+## Requirements
+ - PyTorch >= 0.4.0
+ - Python3.6
+
+I strongly recommend using Python3.6 instead of older version for its better design.
+
+## Evaluation
+First, you should download the cropped testset ALFW and ALFW-2000-3D in [test.data.zip](https://pan.baidu.com/s/1DTVGCG5k0jjjhOc8GcSLOw), then unzip it and put it in the root directory.
+Next, run the benchmark code by providing trained model path.
+I have already provided four pre-trained models in `models` directory. These models are trained using different loss in the first stage. The model size is about 13M due to the high efficiency of mobilenet-v1 structure.
+```
+python3 ./benchmark.py -c models/phase1_wpdc_vdc.pth.tar
+```
+
+The performances of pre-trained models are shown below. In the first stage, the effectiveness of different loss is in order: WPDC > VDC > PDC. While the strategy using VDC to finetune WPDC achieves the best result.
+
+| Model | AFLW (21 pts) | AFLW 2000-3D (68 pts) |
+|:-:|:-:|:-:|
+| phase1_pdc.pth.tar  | 6.956±0.981 | 5.644±1.323 |
+| phase1_vdc.pth.tar  | 6.717±0.924 | 5.030±1.044 |
+| phase1_wpdc.pth.tar | 6.348±0.929 | 4.759±0.996 |
+| phase1_wpdc_vdc.pth.tar | **5.401±0.754** | **4.252±0.976** |
+
+## Training
+[todo]

train.py

@@ -1,8 +1,6 @@
 #!/usr/bin/env python3
 # coding: utf-8
 
-import sys
-import os
 import os.path as osp
 from pathlib import Path
 import numpy as np
@@ -34,7 +32,7 @@ def parse_args():
     parser.add_argument('-j', '--workers', default=6, type=int)
     parser.add_argument('--epochs', default=40, type=int)
     parser.add_argument('--start-epoch', default=1, type=int)
-    parser.add_argument('-b', '--batch-size', default=128)
+    parser.add_argument('-b', '--batch-size', default=128, type=int)
     parser.add_argument('-vb', '--val-batch-size', default=32, type=int)
     parser.add_argument('--base-lr', '--learning-rate', default=0.001, type=float)
     parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
@@ -58,12 +56,13 @@ def parse_args():
     parser.add_argument('--frozen', default='false', type=str2bool)
     parser.add_argument('--milestones', default='15,25,30', type=str)
     parser.add_argument('--task', default='all', type=str)
+    parser.add_argument('--test_initial', default='false', type=str2bool)
     parser.add_argument('--warmup', default=-1, type=int)
     parser.add_argument('--param-fp-train',
-                        default='configs/param_all_norm.pkl',
+                        default='',
                         type=str)
     parser.add_argument('--param-fp-val',
-                        default='configs/param_all_norm_val.pkl')
+                        default='')
     parser.add_argument('--opt-style', default='resample', type=str)  # resample
     parser.add_argument('--resample-num', default=132, type=int)
     parser.add_argument('--loss', default='vdc', type=str)
@@ -168,8 +167,8 @@ def validate(val_loader, model, criterion, epoch):
             target = target.cuda(non_blocking=True)
             output = model(input)
 
-        loss = criterion(output, target)
-        losses.append(loss)
+            loss = criterion(output, target)
+            losses.append(loss)
 
         elapse = time.time() - end
         loss = np.mean(losses)
@@ -274,7 +273,6 @@ def main():
             filename
         )
 
-        # [todo:] add evaluation for AFLW and ALFW2000
         validate(val_loader, model, criterion, epoch)
 
 

training/train.py

@@ -0,0 +1 @@
+../train.py

training/train_pdc.sh

@@ -0,0 +1,28 @@
+#!/usr/bin/env bash
+
+LOG_ALIAS=$1
+LOG_DIR="logs"
+mkdir -p ${LOG_DIR}
+
+LOG_FILE="${LOG_DIR}/${LOG_ALIAS}_`date +'%Y-%m-%d_%H:%M.%S'`.log"
+#echo $LOG_FILE
+
+./train.py --arch="mobilenet_1" \
+      --start-epoch=1 \
+	  --loss=vdc \
+	  --snapshot="snapshot/phase1_pdc" \
+	  --param-fp-train='../train.configs/param_all_norm.pkl' \
+	  --param-fp-val='../train.configs/param_all_norm_val.pkl' \
+	  --warmup=5 \
+	  --opt-style=resample \
+      --batch-size=256 \
+      --base-lr=0.01 \
+      --epochs=50 \
+      --milestones=30,40 \
+      --print-freq=50 \
+      --devices-id=0,1,2,3 \
+      --workers=8 \
+      --filelists-train="../train.configs/augtrain_aug_120x120.list.train" \
+      --filelists-val="../train.configs/train_aug_120x120.list.val" \
+      --root="/mnt/ramdisk/augmentation_crop_120x120_v2" \
+      --log-file="${LOG_FILE}"

training/train_vdc.sh

@@ -0,0 +1,29 @@
+#!/usr/bin/env bash
+
+LOG_ALIAS=$1
+LOG_DIR="logs"
+mkdir -p ${LOG_DIR}
+
+LOG_FILE="${LOG_DIR}/${LOG_ALIAS}_`date +'%Y-%m-%d_%H:%M.%S'`.log"
+#echo $LOG_FILE
+
+./train.py --arch="mobilenet_1" \
+      --start-epoch=1 \
+	  --loss=vdc \
+	  --snapshot="snapshot/phase1_vdc" \
+	  --param-fp-train='../train.configs/param_all_norm.pkl' \
+	  --param-fp-val='../train.configs/param_all_norm_val.pkl' \
+	  --warmup=-1 \
+	  --opt-style=resample \
+      --resample-num=232 \
+      --batch-size=512 \
+      --base-lr=0.00001 \
+      --epochs=50 \
+      --milestones=30,40 \
+      --print-freq=50 \
+      --devices-id=0,1,2,3 \
+      --workers=8 \
+      --filelists-train="../train.configs/augtrain_aug_120x120.list.train" \
+      --filelists-val="../train.configs/train_aug_120x120.list.val" \
+      --root="/mnt/ramdisk/augmentation_crop_120x120_v2" \
+      --log-file="${LOG_FILE}"

training/train_wpdc.sh

@@ -0,0 +1,29 @@
+#!/usr/bin/env bash
+
+LOG_ALIAS=$1
+LOG_DIR="logs"
+mkdir -p ${LOG_DIR}
+
+LOG_FILE="${LOG_DIR}/${LOG_ALIAS}_`date +'%Y-%m-%d_%H:%M.%S'`.log"
+#echo $LOG_FILE
+
+./train.py --arch="mobilenet_1" \
+      --start-epoch=1 \
+	  --loss=wpdc \
+	  --snapshot="snapshot/phase1_wpdc" \
+	  --param-fp-train='../train.configs/param_all_norm.pkl' \
+	  --param-fp-val='../train.configs/param_all_norm_val.pkl' \
+	  --warmup=5 \
+	  --opt-style=resample \
+      --resample-num=232 \
+      --batch-size=512 \
+      --base-lr=0.02 \
+      --epochs=50 \
+      --milestones=30,40 \
+      --print-freq=50 \
+      --devices-id=0,1,2,3 \
+      --workers=8 \
+      --filelists-train="../train.configs/augtrain_aug_120x120.list.train" \
+      --filelists-val="../train.configs/train_aug_120x120.list.val" \
+      --root="/mnt/ramdisk/augmentation_crop_120x120_v2" \
+      --log-file="${LOG_FILE}"

vdc_loss.py

@@ -1,34 +1,15 @@
 #!/usr/bin/env python3
 # coding: utf-8
 
-import sys
-import os.path as osp
-from pathlib import Path
-import numpy as np
 import torch
 import torch.nn as nn
-import torch.cuda as cuda
 from io_utils import _load, _numpy_to_cuda, _numpy_to_tensor
-import time
-import random
 from params import *
 
-# d = 'configs'
-# keypoints = _load(osp.join(d, 'keypoints_sim.npy'))
-# w_shp = _load(osp.join(d, 'w_shp_sim.npy'))
-# w_exp = _load(osp.join(d, 'w_exp_sim.npy'))  # simplified version
-# meta = _load(osp.join(d, 'param_whitening.pkl'))
-#
-# param_mean = meta.get('param_mean')
-# param_std = meta.get('param_std')
-# u_shp = _load(osp.join(d, 'u_shp.npy'))
-# u_exp = _load(osp.join(d, 'u_exp.npy'))
-# u = u_shp + u_exp
-
-# _to_tensor = _numpy_to_tensor  # cpu
 _to_tensor = _numpy_to_cuda  # gpu
 
 
+
 def _parse_param_batch(param):
     """Work for both numpy and tensor"""
     N = param.shape[0]

wpdc_loss.py

@@ -1,35 +1,12 @@
 #!/usr/bin/env python3
 # coding: utf-8
 
-import sys
-import os.path as osp
-from pathlib import Path
-import numpy as np
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
-import pickle
 from math import sqrt
-import time
 from io_utils import _load, _numpy_to_cuda, _numpy_to_tensor, _load_gpu
 from params import *
 
-# d = 'configs'
-# keypoints = _load(osp.join(d, 'keypoints_sim.npy'))
-# w_shp = _load(osp.join(d, 'w_shp_sim.npy'))
-# w_exp = _load(osp.join(d, 'w_exp_sim.npy'))  # simplified version
-#
-# meta = _load(osp.join(d, 'param_whitening.pkl'))
-# param_mean = meta.get('param_mean')
-# param_std = meta.get('param_std')
-# u_shp = _load(osp.join(d, 'u_shp.npy'))
-# u_exp = _load(osp.join(d, 'u_exp.npy'))
-# u = u_shp + u_exp
-# w = np.concatenate((w_shp, w_exp), axis=1)
-# w_base = w[keypoints]
-# w_norm = np.linalg.norm(w, axis=0)
-# w_base_norm = np.linalg.norm(w_base, axis=0)
-
 _to_tensor = _numpy_to_cuda  # gpu