-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathprepare_data.py
158 lines (112 loc) · 7.28 KB
/
prepare_data.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
import torch
import numpy as np
import os
import time
from torch.utils.data import Dataset, DataLoader
from human_body_prior.body_model.body_model import BodyModel
from human_body_prior.tools.omni_tools import copy2cpu as c2c
from human_body_prior.tools.rotation_tools import aa2matrot,matrot2aa,local2global_pose
from utils import utils_transform
import glob
# from IPython import embed
import time
import copy
import pickle
dataroot_amass ="amass" # root of amass dataset
for dataroot_subset in ["MPI_HDM05", "BioMotionLab_NTroje", "CMU"]:
print(dataroot_subset)
for phase in ["test"]:
print(phase)
savedir = os.path.join("./data_fps60", dataroot_subset, phase)
if not os.path.exists(savedir):
os.makedirs(savedir)
split_file = os.path.join("./data_split", dataroot_subset, phase+"_split.txt")
with open(split_file, 'r') as f:
filepaths = [line.rstrip('\n') for line in f]
rotation_local_full_gt_list = []
hmd_position_global_full_gt_list = []
body_parms_list = []
head_global_trans_list = []
support_dir = 'support_data/'
bm_fname_male = os.path.join(support_dir, 'body_models/smplh/{}/model.npz'.format('male'))
dmpl_fname_male = os.path.join(support_dir, 'body_models/dmpls/{}/model.npz'.format('male'))
bm_fname_female = os.path.join(support_dir, 'body_models/smplh/{}/model.npz'.format('female'))
dmpl_fname_female = os.path.join(support_dir, 'body_models/dmpls/{}/model.npz'.format('female'))
num_betas = 16 # number of body parameters
num_dmpls = 8 # number of DMPL parameters
bm_male = BodyModel(bm_fname=bm_fname_male, num_betas=num_betas, num_dmpls=num_dmpls, dmpl_fname=dmpl_fname_male)#.to(comp_device)
bm_female = BodyModel(bm_fname=bm_fname_female, num_betas=num_betas, num_dmpls=num_dmpls, dmpl_fname=dmpl_fname_female)
idx = 0
for filepath in filepaths:
data = dict()
bdata = np.load(os.path.join('AMASS',filepath),allow_pickle=True)
# print(list(bdata.keys())) ### check keys of body data: ['trans', 'gender', 'mocap_framerate', 'betas', 'dmpls', 'poses']
try:
framerate = bdata["mocap_framerate"]
print("framerate is {}".format(framerate))
except:
print(filepath)
print(list(bdata.keys()))
continue # skip shape.npz
# pass
# else:
idx+=1
print(idx)
if framerate == 120:
stride = 2
elif framerate == 60:
stride = 1
bdata_poses = bdata["poses"][::stride,...]
bdata_trans = bdata["trans"][::stride,...]
subject_gender = bdata["gender"]
bm = bm_male# if subject_gender == 'male' else bm_female
body_parms = {
'root_orient': torch.Tensor(bdata_poses[:, :3]),#.to(comp_device), # controls the global root orientation
'pose_body': torch.Tensor(bdata_poses[:, 3:66]),#.to(comp_device), # controls the body
'trans': torch.Tensor(bdata_trans),#.to(comp_device), # controls the global body position
}
body_parms_list = body_parms
body_pose_world = bm(**{k:v for k,v in body_parms.items() if k in ['pose_body','root_orient','trans']})
# self.rotation_local_full_gt_list.append(body_parms['pose_body'])
# self.rotation_local_full_gt_list.append(torch.Tensor(bdata['poses'][:, :66]))
output_aa = torch.Tensor(bdata_poses[:, :66]).reshape(-1,3)
output_6d = utils_transform.aa2sixd(output_aa).reshape(bdata_poses.shape[0],-1)
rotation_local_full_gt_list = output_6d[1:] # GT
rotation_local_matrot = aa2matrot(torch.tensor(bdata_poses).reshape(-1,3)).reshape(bdata_poses.shape[0],-1,9)
rotation_global_matrot = local2global_pose(rotation_local_matrot, bm.kintree_table[0].long()) # rotation of joints relative to the origin ALL
head_rotation_global_matrot = rotation_global_matrot[:,[15],:,:] # Head
hands_rotation_global_matrot = rotation_global_matrot[:,[20,21],:,:] # Hands
rotation_global_6d = utils_transform.matrot2sixd(rotation_global_matrot.reshape(-1,3,3)).reshape(rotation_global_matrot.shape[0],-1,6)
input_rotation_global_6d = rotation_global_6d[1:,[15,20,21],:] # Head and hands
rotation_velocity_global_matrot = torch.matmul(torch.inverse(rotation_global_matrot[:-1]),rotation_global_matrot[1:])
rotation_velocity_global_6d = utils_transform.matrot2sixd(rotation_velocity_global_matrot.reshape(-1,3,3)).reshape(rotation_velocity_global_matrot.shape[0],-1,6)
input_rotation_velocity_global_6d = rotation_velocity_global_6d[:,[15,20,21],:]
position_global_full_gt_world = body_pose_world.Jtr[:,:22,:] # position of joints relative to the world origin
position_head_world = position_global_full_gt_world[:,15,:] # world position of head
position_hands_world = position_global_full_gt_world[:,[20,21],:] # world position of hands
head_global_trans = torch.eye(4).repeat(position_head_world.shape[0],1,1)
head_global_trans[:,:3,:3] = head_rotation_global_matrot.squeeze()
head_global_trans[:,:3,3] = position_global_full_gt_world[:,15,:]
hands_global_trans = torch.eye(4).repeat(position_hands_world.shape[0],2,1,1)
hands_global_trans[:,:,:3,:3] = hands_rotation_global_matrot.squeeze()
hands_global_trans[:,:,:3,3] = position_global_full_gt_world[:,[20,21],:]
head_global_trans_list = head_global_trans[1:]
hands_global_trans_list = hands_global_trans[1:]
num_frames = position_global_full_gt_world.shape[0]-1
hmd_position_global_full_gt_list = torch.cat([
input_rotation_global_6d.reshape(num_frames,-1), # Rotation
input_rotation_velocity_global_6d.reshape(num_frames,-1), # Angular Velocity
position_global_full_gt_world[1:, [15,20,21], :].reshape(num_frames,-1), # Position
position_global_full_gt_world[1:, [15,20,21], :].reshape(num_frames,-1)-position_global_full_gt_world[:-1, [15,20,21], :].reshape(num_frames,-1)], dim=-1) # Linear velocity
data_count = len(hmd_position_global_full_gt_list)
print(str(idx)+'/'+str(len(filepaths)))
data['rotation_local_full_gt_list'] = rotation_local_full_gt_list
data['hmd_position_global_full_gt_list'] = hmd_position_global_full_gt_list
data['body_parms_list'] = body_parms_list
data['head_global_trans_list'] = head_global_trans_list
data['hands_global_trans_list'] = hands_global_trans_list
data['framerate'] = 60
data['gender'] = subject_gender
data['filepath'] = filepath
with open(os.path.join(savedir,'{}.pkl'.format(idx)), 'wb') as f:
pickle.dump(data, f)