-
Notifications
You must be signed in to change notification settings - Fork 441
/
Copy pathdemo_reconstruct.py
executable file
·131 lines (123 loc) · 7.51 KB
/
demo_reconstruct.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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
# -*- coding: utf-8 -*-
#
# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# Using this computer program means that you agree to the terms
# in the LICENSE file included with this software distribution.
# Any use not explicitly granted by the LICENSE is prohibited.
#
# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# For comments or questions, please email us at [email protected]
# For commercial licensing contact, please contact [email protected]
import os, sys
import cv2
import numpy as np
from time import time
from scipy.io import savemat
import argparse
from tqdm import tqdm
import torch
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from decalib.deca import DECA
from decalib.datasets import datasets
from decalib.utils import util
from decalib.utils.config import cfg as deca_cfg
from decalib.utils.tensor_cropper import transform_points
def main(args):
# if args.rasterizer_type != 'standard':
# args.render_orig = False
savefolder = args.savefolder
device = args.device
os.makedirs(savefolder, exist_ok=True)
# load test images
testdata = datasets.TestData(args.inputpath, iscrop=args.iscrop, face_detector=args.detector, sample_step=args.sample_step)
# run DECA
deca_cfg.model.use_tex = args.useTex
deca_cfg.rasterizer_type = args.rasterizer_type
deca_cfg.model.extract_tex = args.extractTex
deca = DECA(config = deca_cfg, device=device)
# for i in range(len(testdata)):
for i in tqdm(range(len(testdata))):
name = testdata[i]['imagename']
images = testdata[i]['image'].to(device)[None,...]
with torch.no_grad():
codedict = deca.encode(images)
opdict, visdict = deca.decode(codedict) #tensor
if args.render_orig:
tform = testdata[i]['tform'][None, ...]
tform = torch.inverse(tform).transpose(1,2).to(device)
original_image = testdata[i]['original_image'][None, ...].to(device)
_, orig_visdict = deca.decode(codedict, render_orig=True, original_image=original_image, tform=tform)
orig_visdict['inputs'] = original_image
if args.saveDepth or args.saveKpt or args.saveObj or args.saveMat or args.saveImages:
os.makedirs(os.path.join(savefolder, name), exist_ok=True)
# -- save results
if args.saveDepth:
depth_image = deca.render.render_depth(opdict['trans_verts']).repeat(1,3,1,1)
visdict['depth_images'] = depth_image
cv2.imwrite(os.path.join(savefolder, name, name + '_depth.jpg'), util.tensor2image(depth_image[0]))
if args.saveKpt:
np.savetxt(os.path.join(savefolder, name, name + '_kpt2d.txt'), opdict['landmarks2d'][0].cpu().numpy())
np.savetxt(os.path.join(savefolder, name, name + '_kpt3d.txt'), opdict['landmarks3d'][0].cpu().numpy())
if args.saveObj:
deca.save_obj(os.path.join(savefolder, name, name + '.obj'), opdict)
if args.saveMat:
opdict = util.dict_tensor2npy(opdict)
savemat(os.path.join(savefolder, name, name + '.mat'), opdict)
if args.saveVis:
cv2.imwrite(os.path.join(savefolder, name + '_vis.jpg'), deca.visualize(visdict))
if args.render_orig:
cv2.imwrite(os.path.join(savefolder, name + '_vis_original_size.jpg'), deca.visualize(orig_visdict))
if args.saveImages:
for vis_name in ['inputs', 'rendered_images', 'albedo_images', 'shape_images', 'shape_detail_images', 'landmarks2d']:
if vis_name not in visdict.keys():
continue
image = util.tensor2image(visdict[vis_name][0])
cv2.imwrite(os.path.join(savefolder, name, name + '_' + vis_name +'.jpg'), util.tensor2image(visdict[vis_name][0]))
if args.render_orig:
image = util.tensor2image(orig_visdict[vis_name][0])
cv2.imwrite(os.path.join(savefolder, name, 'orig_' + name + '_' + vis_name +'.jpg'), util.tensor2image(orig_visdict[vis_name][0]))
print(f'-- please check the results in {savefolder}')
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='DECA: Detailed Expression Capture and Animation')
parser.add_argument('-i', '--inputpath', default='TestSamples/examples', type=str,
help='path to the test data, can be image folder, image path, image list, video')
parser.add_argument('-s', '--savefolder', default='TestSamples/examples/results', type=str,
help='path to the output directory, where results(obj, txt files) will be stored.')
parser.add_argument('--device', default='cuda', type=str,
help='set device, cpu for using cpu' )
# process test images
parser.add_argument('--iscrop', default=True, type=lambda x: x.lower() in ['true', '1'],
help='whether to crop input image, set false only when the test image are well cropped' )
parser.add_argument('--sample_step', default=10, type=int,
help='sample images from video data for every step' )
parser.add_argument('--detector', default='fan', type=str,
help='detector for cropping face, check decalib/detectors.py for details' )
# rendering option
parser.add_argument('--rasterizer_type', default='standard', type=str,
help='rasterizer type: pytorch3d or standard' )
parser.add_argument('--render_orig', default=True, type=lambda x: x.lower() in ['true', '1'],
help='whether to render results in original image size, currently only works when rasterizer_type=standard')
# save
parser.add_argument('--useTex', default=False, type=lambda x: x.lower() in ['true', '1'],
help='whether to use FLAME texture model to generate uv texture map, \
set it to True only if you downloaded texture model' )
parser.add_argument('--extractTex', default=True, type=lambda x: x.lower() in ['true', '1'],
help='whether to extract texture from input image as the uv texture map, set false if you want albeo map from FLAME mode' )
parser.add_argument('--saveVis', default=True, type=lambda x: x.lower() in ['true', '1'],
help='whether to save visualization of output' )
parser.add_argument('--saveKpt', default=False, type=lambda x: x.lower() in ['true', '1'],
help='whether to save 2D and 3D keypoints' )
parser.add_argument('--saveDepth', default=False, type=lambda x: x.lower() in ['true', '1'],
help='whether to save depth image' )
parser.add_argument('--saveObj', default=False, type=lambda x: x.lower() in ['true', '1'],
help='whether to save outputs as .obj, detail mesh will end with _detail.obj. \
Note that saving objs could be slow' )
parser.add_argument('--saveMat', default=False, type=lambda x: x.lower() in ['true', '1'],
help='whether to save outputs as .mat' )
parser.add_argument('--saveImages', default=False, type=lambda x: x.lower() in ['true', '1'],
help='whether to save visualization output as seperate images' )
main(parser.parse_args())