-
Notifications
You must be signed in to change notification settings - Fork 15
/
Copy pathdemo_voca.py
executable file
·246 lines (202 loc) · 10.6 KB
/
demo_voca.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
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
import numpy as np
import librosa
import os, argparse, pickle
from SelfTalk import SelfTalk
from transformers import Wav2Vec2Processor
import torch
import time
import cv2
import tempfile
from subprocess import call
import pyrender
from psbody.mesh import Mesh
import trimesh
os.environ['PYOPENGL_PLATFORM'] = 'egl' # egl
@torch.no_grad()
def test_model(args):
if not os.path.exists(args.result_path):
os.makedirs(args.result_path)
# build model
model = SelfTalk(args)
model.load_state_dict(torch.load(os.path.join(args.dataset, '{}.pth'.format(args.model_name)),
map_location=torch.device(args.device)))
model = model.to(torch.device(args.device))
model.eval()
template_file = os.path.join(args.dataset, args.template_path)
with open(template_file, 'rb') as fin:
templates = pickle.load(fin, encoding='latin1')
temp = templates[args.subject]
template = temp.reshape((-1))
template = np.reshape(template, (-1, template.shape[0]))
template = torch.FloatTensor(template).to(device=args.device)
wav_path = args.wav_path
test_name = os.path.basename(wav_path).split(".")[0]
speech_array, sampling_rate = librosa.load(os.path.join(wav_path), sr=16000)
processor = Wav2Vec2Processor.from_pretrained("jonatasgrosman/wav2vec2-large-xlsr-53-english")
audio_feature = np.squeeze(processor(speech_array, sampling_rate=16000).input_values)
audio_feature = np.reshape(audio_feature, (-1, audio_feature.shape[0]))
audio_feature = torch.FloatTensor(audio_feature).to(device=args.device)
start = time.time()
prediction, lip_features, logits = model.predict(audio_feature, template)
end = time.time()
print("Model predict time: ", end - start)
prediction = prediction.squeeze()
np.save(os.path.join(args.result_path, test_name), prediction.detach().cpu().numpy())
def get_unit_factor(unit):
if unit == 'mm':
return 1000.0
elif unit == 'cm':
return 100.0
elif unit == 'm':
return 1.0
else:
raise ValueError('Unit not supported')
# The implementation of rendering is borrowed from VOCA: https://github.com/TimoBolkart/voca/blob/master/utils/rendering.py
def render_mesh_helper(mesh, t_center, rot=np.zeros(3), tex_img=None, v_colors=None,
errors=None, error_unit='m', min_dist_in_mm=0.0, max_dist_in_mm=3.0, z_offset=0, xmag=0.5,
y=0.7, z=1, camera='o', r=None):
camera_params = {'c': np.array([400, 400]),
'k': np.array([-0.19816071, 0.92822711, 0, 0, 0]),
'f': np.array([4754.97941935 / 2, 4754.97941935 / 2])}
frustum = {'near': 0.01, 'far': 3.0, 'height': 800, 'width': 800}
mesh_copy = Mesh(mesh.v, mesh.f)
mesh_copy.v[:] = cv2.Rodrigues(rot)[0].dot((mesh_copy.v - t_center).T).T + t_center
intensity = 2.0
rgb_per_v = None
primitive_material = pyrender.material.MetallicRoughnessMaterial(
alphaMode='BLEND',
baseColorFactor=[0.3, 0.3, 0.3, 1.0],
metallicFactor=0.8,
roughnessFactor=0.8
)
color = np.array([0.3, 0.5, 0.55])
tri_mesh = trimesh.Trimesh(vertices=mesh_copy.v, faces=mesh_copy.f, vertex_colors=rgb_per_v)
render_mesh = pyrender.Mesh.from_trimesh(tri_mesh, material=primitive_material, smooth=True)
if 1 == 1:
scene = pyrender.Scene(ambient_light=[.2, .2, .2], bg_color=[0, 0, 0])
else:
scene = pyrender.Scene(ambient_light=[.2, .2, .2], bg_color=[255, 255, 255])
camera = pyrender.IntrinsicsCamera(fx=camera_params['f'][0],
fy=camera_params['f'][1],
cx=camera_params['c'][0],
cy=camera_params['c'][1],
znear=frustum['near'],
zfar=frustum['far'])
scene.add(render_mesh, pose=np.eye(4))
camera_pose = np.eye(4)
camera_pose[:3, 3] = np.array([0, 0, 1.0 - z_offset])
scene.add(camera, pose=[[1, 0, 0, 0],
[0, 1, 0, 0],
[0, 0, 1, 1],
[0, 0, 0, 1]])
angle = np.pi / 6.0
pos = camera_pose[:3, 3]
light_color = np.array([1., 1., 1.])
light = pyrender.DirectionalLight(color=light_color, intensity=intensity)
light_pose = np.eye(4)
light_pose[:3, 3] = pos
scene.add(light, pose=light_pose.copy())
light_pose[:3, 3] = cv2.Rodrigues(np.array([angle, 0, 0]))[0].dot(pos)
scene.add(light, pose=light_pose.copy())
light_pose[:3, 3] = cv2.Rodrigues(np.array([-angle, 0, 0]))[0].dot(pos)
scene.add(light, pose=light_pose.copy())
light_pose[:3, 3] = cv2.Rodrigues(np.array([0, -angle, 0]))[0].dot(pos)
scene.add(light, pose=light_pose.copy())
light_pose[:3, 3] = cv2.Rodrigues(np.array([0, angle, 0]))[0].dot(pos)
scene.add(light, pose=light_pose.copy())
flags = pyrender.RenderFlags.SKIP_CULL_FACES
# try:
# egl
r = pyrender.OffscreenRenderer(viewport_width=frustum['width'], viewport_height=frustum['height'])
color, _ = r.render(scene, flags=flags)
# except:
# print('pyrender: Failed rendering frame')
# color = np.zeros((frustum['height'], frustum['width'], 3), dtype='uint8')
return color[..., ::-1]
def render_sequence_meshes(audio_fname, sequence_vertices, template, out_path, uv_template_fname='',
texture_img_fname=''):
if not os.path.exists(out_path):
os.makedirs(out_path)
tmp_video_file = tempfile.NamedTemporaryFile('w', suffix='.mp4', dir=out_path)
if int(cv2.__version__[0]) < 3:
writer = cv2.VideoWriter(tmp_video_file.name, cv2.cv.CV_FOURCC(*'mp4v'), 30, (800, 800), True)
else:
writer = cv2.VideoWriter(tmp_video_file.name, cv2.VideoWriter_fourcc(*'mp4v'), 30, (800, 800), True)
if os.path.exists(uv_template_fname) and os.path.exists(texture_img_fname):
uv_template = Mesh(filename=uv_template_fname)
vt, ft = uv_template.vt, uv_template.ft
tex_img = cv2.imread(texture_img_fname)[:, :, ::-1]
else:
vt, ft = None, None
tex_img = None
num_frames = sequence_vertices.shape[0]
center = np.mean(sequence_vertices[0], axis=0)
for i_frame in range(num_frames):
render_mesh = Mesh(sequence_vertices[i_frame], template.f)
if vt is not None and ft is not None:
render_mesh.vt, render_mesh.ft = vt, ft
img = render_mesh_helper(render_mesh, center)
writer.write(img)
writer.release()
video_fname = os.path.join(out_path, 'video.mp4')
cmd = ('ffmpeg' + ' -i {0} -i {1} -pix_fmt yuv420p -qscale 0 {2} -y'.format(
audio_fname, tmp_video_file.name, video_fname)).split()
call(cmd)
def output_sequence_meshes(sequence_vertices, template, out_path, uv_template_fname='', texture_img_fname=''):
mesh_out_path = os.path.join(out_path, 'meshes')
if not os.path.exists(mesh_out_path):
os.makedirs(mesh_out_path)
if os.path.exists(uv_template_fname):
uv_template = Mesh(filename=uv_template_fname)
vt, ft = uv_template.vt, uv_template.ft
else:
vt, ft = None, None
num_frames = sequence_vertices.shape[0]
for i_frame in range(num_frames):
out_fname = os.path.join(mesh_out_path, '%05d.obj' % i_frame)
out_mesh = Mesh(sequence_vertices[i_frame], template.f)
if vt is not None and ft is not None:
out_mesh.vt, out_mesh.ft = vt, ft
if os.path.exists(texture_img_fname):
out_mesh.set_texture_image(texture_img_fname)
out_mesh.write_obj(out_fname)
def main():
parser = argparse.ArgumentParser(
description='SelfTalk: A Self-Supervised Commutative Training Diagram to Comprehend 3D Talking Faces')
parser.add_argument("--model_name", type=str, default="vocaset", help='vocaset or BIWI')
parser.add_argument("--dataset", type=str, default="vocaset", help='vocaset or BIWI')
parser.add_argument("--fps", type=float, default=30, help='frame rate - 30 for vocaset; 25 for BIWI')
parser.add_argument("--feature_dim", type=int, default=512, help='512 for vocaset; 1024 for BIWI')
parser.add_argument("--period", type=int, default=30, help='period in PPE - 30 for vocaset; 25 for BIWI')
parser.add_argument("--vertice_dim", type=int, default=5023 * 3,
help='number of vertices - 5023*3 for vocaset; 23370*3 for BIWI')
parser.add_argument("--device", type=str, default="cuda", help='cuda or cpu')
parser.add_argument("--train_subjects", type=str, default="FaceTalk_170728_03272_TA"
" FaceTalk_170904_00128_TA FaceTalk_170725_00137_TA FaceTalk_170915_00223_TA"
" FaceTalk_170811_03274_TA FaceTalk_170913_03279_TA"
" FaceTalk_170904_03276_TA FaceTalk_170912_03278_TA")
parser.add_argument("--test_subjects", type=str, default="FaceTalk_170809_00138_TA"
" FaceTalk_170731_00024_TA")
parser.add_argument("--output_path", type=str, default="demo/output", help='path of the rendered video sequence')
parser.add_argument("--wav_path", type=str, default="demo/wav/test.wav", help='path of the input audio signal')
parser.add_argument("--result_path", type=str, default="demo/result", help='path of the predictions')
parser.add_argument("--subject", type=str, default="FaceTalk_170908_03277_TA",
help='select a subject from test_subjects or train_subjects')
parser.add_argument("--background_black", type=bool, default=True, help='whether to use black background')
parser.add_argument("--template_path", type=str, default="templates.pkl", help='path of the personalized templates')
parser.add_argument("--render_template_path", type=str, default="templates",
help='path of the mesh in BIWI/FLAME topology')
args = parser.parse_args()
test_model(args)
fa_path = args.result_path + "/" + args.wav_path.split("/")[-1].split(".")[0] + ".npy"
temp = "./vocaset/templates/FLAME_sample.ply"
out_path = fa_path.split(".")[0]
audio_fname = args.wav_path
template = Mesh(filename=temp)
predicted_vertices_out = np.load(fa_path).reshape(-1, 5023, 3)
print("Start rendering...")
output_sequence_meshes(predicted_vertices_out, template, out_path)
render_sequence_meshes(audio_fname, predicted_vertices_out, template, out_path, uv_template_fname='',
texture_img_fname='')
if __name__ == "__main__":
main()