forked from cleardusk/3DDFA_V2
-
Notifications
You must be signed in to change notification settings - Fork 0
/
demo.py
executable file
·86 lines (70 loc) · 3.19 KB
/
demo.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
# coding: utf-8
__author__ = 'cleardusk'
import sys
import argparse
import cv2
import yaml
from FaceBoxes import FaceBoxes
from TDDFA import TDDFA
from utils.render import render
from utils.depth import depth
from utils.pncc import pncc
from utils.uv import uv_tex
from utils.pose import viz_pose
from utils.serialization import ser_to_ply, ser_to_obj
from utils.functions import draw_landmarks, get_suffix
from utils.tddfa_util import str2bool
def main(args):
# Init TDDFA
cfg = yaml.load(open(args.config), Loader=yaml.SafeLoader)
gpu_mode = args.mode == 'gpu'
tddfa = TDDFA(gpu_mode=gpu_mode, **cfg)
# Init FaceBoxes
face_boxes = FaceBoxes()
# Given a still image path and load to BGR channel
img = cv2.imread(args.img_fp)
# Detect faces, get 3DMM params and roi boxes
boxes = face_boxes(img)
n = len(boxes)
print(f'Detect {n} faces')
if n == 0:
print(f'No face detected, exit')
sys.exit(-1)
param_lst, roi_box_lst = tddfa(img, boxes)
# Visualization and serialization
dense_flag = args.opt in ('2d_dense', '3d', 'depth', 'pncc', 'uv_tex', 'ply', 'obj')
old_suffix = get_suffix(args.img_fp)
new_suffix = f'.{args.opt}' if args.opt in ('ply', 'obj') else '.jpg'
wfp = f'examples/results/{args.img_fp.split("/")[-1].replace(old_suffix, "")}_{args.opt}' + new_suffix
ver_lst = tddfa.recon_vers(param_lst, roi_box_lst, dense_flag=dense_flag)
if args.opt == '2d_sparse':
draw_landmarks(img, ver_lst, show_flag=args.show_flag, dense_flag=dense_flag, wfp=wfp)
elif args.opt == '2d_dense':
draw_landmarks(img, ver_lst, show_flag=args.show_flag, dense_flag=dense_flag, wfp=wfp)
elif args.opt == '3d':
render(img, ver_lst, alpha=0.6, show_flag=args.show_flag, wfp=wfp)
elif args.opt == 'depth':
# if `with_bf_flag` is False, the background is black
depth(img, ver_lst, show_flag=args.show_flag, wfp=wfp, with_bg_flag=True)
elif args.opt == 'pncc':
pncc(img, ver_lst, show_flag=args.show_flag, wfp=wfp, with_bg_flag=True)
elif args.opt == 'uv_tex':
uv_tex(img, ver_lst, show_flag=args.show_flag, wfp=wfp)
elif args.opt == 'pose':
viz_pose(img, param_lst, ver_lst, show_flag=args.show_flag, wfp=wfp)
elif args.opt == 'ply':
ser_to_ply(ver_lst, height=img.shape[0], wfp=wfp)
elif args.opt == 'obj':
ser_to_obj(img, ver_lst, height=img.shape[0], wfp=wfp)
else:
raise ValueError(f'Unknown opt {args.opt}')
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='The demo of still image of 3DDFA_V2')
parser.add_argument('-c', '--config', type=str, default='configs/mb1_120x120.yml')
parser.add_argument('-f', '--img_fp', type=str, default='examples/inputs/trump_hillary.jpg')
parser.add_argument('-m', '--mode', type=str, default='cpu', help='gpu or cpu mode')
parser.add_argument('-o', '--opt', type=str, default='2d_sparse',
choices=['2d_sparse', '2d_dense', '3d', 'depth', 'pncc', 'uv_tex', 'pose', 'ply', 'obj'])
parser.add_argument('--show_flag', type=str2bool, default='true', help='whether to show the visualization result')
args = parser.parse_args()
main(args)