forked from svip-lab/impersonator
-
Notifications
You must be signed in to change notification settings - Fork 0
/
demo_imitator.py
111 lines (80 loc) · 3.54 KB
/
demo_imitator.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
import numpy as np
from tqdm import tqdm
import os
import glob
from models.imitator import Imitator
from options.test_options import TestOptions
from utils.util import mkdir
import pickle
from utils.video import make_video
from run_imitator import adaptive_personalize
mixamo_root_path = './assets/samples/refs/mixamo'
# MIXAMO_DANCE_ACTION_IDX_LIST = [102]
# MIXAMO_BASE_ACTION_IDX_LIST = [20, 22, 32, 70]
# MIXAMO_ACROBAT_ACTION_IDX_LIST = [7, 31, 83, 131, 145]
MIXAMO_DANCE_ACTION_IDX_LIST = [102]
MIXAMO_BASE_ACTION_IDX_LIST = [20]
MIXAMO_ACROBAT_ACTION_IDX_LIST = [7]
def load_mixamo_smpl(mixamo_idx):
global mixamo_root_path
dir_name = '%.4d' % mixamo_idx
pkl_path = os.path.join(mixamo_root_path, dir_name, 'result.pkl')
with open(pkl_path, 'rb') as f:
result = pickle.load(f)
anim_len = result['anim_len']
pose_array = result['smpl_array'].reshape(anim_len, -1)
cam_array = result['cam_array']
shape_array = np.ones((anim_len, 10))
smpl_array = np.concatenate((cam_array, pose_array, shape_array), axis=1)
return smpl_array
def generate_actor_result(test_opt, src_img_path):
imitator = Imitator(test_opt)
src_img_name = os.path.split(src_img_path)[-1][:-4]
test_opt.src_path = src_img_path
if test_opt.post_tune:
adaptive_personalize(test_opt, imitator, visualizer=None)
else:
imitator.personalize(test_opt.src_path, visualizer=None)
action_list_dict = {'dance': MIXAMO_DANCE_ACTION_IDX_LIST,
'base': MIXAMO_BASE_ACTION_IDX_LIST,
'acrobat': MIXAMO_ACROBAT_ACTION_IDX_LIST}
for action_type in ['dance', 'base', 'acrobat']:
for i in action_list_dict[action_type]:
if test_opt.output_dir:
pred_output_dir = os.path.join(test_opt.output_dir, 'mixamo_preds')
if os.path.exists(pred_output_dir):
os.system("rm -r %s" % pred_output_dir)
mkdir(pred_output_dir)
else:
pred_output_dir = None
print(pred_output_dir)
tgt_smpls = load_mixamo_smpl(i)
imitator.inference_by_smpls(tgt_smpls, cam_strategy='smooth', output_dir=pred_output_dir, visualizer=None)
save_dir = os.path.join(test_opt.output_dir, src_img_name, action_type)
mkdir(save_dir)
output_mp4_path = os.path.join(save_dir, 'mixamo_%.4d_%s.mp4' % (i, src_img_name))
img_path_list = sorted(glob.glob('%s/*.jpg' % pred_output_dir))
make_video(output_mp4_path, img_path_list, save_frames_dir=None, fps=30)
def clean(output_dir):
for item in ['imgs', 'pairs', 'mixamo_preds', 'pairs_meta.pkl']:
filepath = os.path.join(output_dir, item)
if os.path.exists(filepath):
os.system("rm -r %s" % filepath)
def main():
# meta imitator
test_opt = TestOptions().parse()
test_opt.bg_ks = 25
test_opt.front_warp = False
test_opt.post_tune = True
test_opt.output_dir = mkdir('./outputs/results/demos/imitators')
# source images from iPER
iamges_paths = ['./assets/src_imgs/imper_A_Pose/009_5_1_000.jpg',
'./assets/src_imgs/imper_A_Pose/024_8_2_0000.jpg',
'./assets/src_imgs/fashion_woman/Sweaters-id_0000088807_4_full.jpg']
for src_img_path in tqdm(iamges_paths):
generate_actor_result(test_opt, src_img_path)
# clean other files
clean(test_opt.output_dir)
print('Completed! All demo videos are save in {}'.format(test_opt.output_dir))
if __name__ == "__main__":
main()