File size: 7,146 Bytes
45950ff | 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 | import sys
sys.path.append('/mnt/shenzhen2cephfs/capybarali/codes/humanoid')
import torch, yaml, os
from tqdm import tqdm
from smplx import SMPLX
from src.utils.rotation_conversions import quaternion_to_matrix, matrix_to_rotation_6d, matrix_to_axis_angle, rotation_6d_to_matrix, matrix_to_quaternion
import numpy as np
from argparse import ArgumentParser
import joblib
from data.vis import vis_3d_motion
from data.vis_g1 import vis_3d_g1
from copy import deepcopy
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
_debug_save_count = 0
_debug_save_dir = 'data/debug_transform'
parser = ArgumentParser(description="Launch MoCap processing")
parser.add_argument('--save_root', type=str, default="data/final_data")
parser.add_argument('--start_idx', type=int, default=0)
parser.add_argument('--interval', type=int, default=1)
args = parser.parse_args()
os.makedirs(args.save_root, exist_ok=True)
os.makedirs(os.path.join(args.save_root, 'motions'), exist_ok=True)
smplx_model = SMPLX(
model_path='checkpoints/human_model/SMPLX_NEUTRAL.npz',
use_pca=False, num_expression_coeffs=100, num_betas=10, ext='npz'
).to(device)
def extract_g1_component(x):
vel_xy = x[:, :2]
dof = x[:, -29:]
root_rot_mat = rotation_6d_to_matrix(x[:, 2:8])
trans_xyz = torch.cat([torch.cumsum(vel_xy, dim=0), x[:, 10]])
rot_mat = torch.tensor([[1, 0, 0], [0, 0, -1], [0, 1, 0]]).float()
global_orient_mat = torch.from_numpy(root_rot_mat).squeeze(1)
global_orient_mat = torch.einsum('ij,tjk->tik', rot_mat, global_orient_mat)
rot_quat = matrix_to_quaternion(global_orient_mat) # (T, 4) wxyz order
transl = torch.from_numpy(trans_xyz).float()
transl = torch.einsum('ij,tj->ti', rot_mat, transl)
return dof, rot_quat, trans_xyz
def get_smplx_motion(data_path):
smplx_data = joblib.load(data_path)
transl = torch.from_numpy(smplx_data['smplx_transl'])
betas = torch.from_numpy(np.load('checkpoints/humanoid_model/g1/betas.npy'))
global_orient = torch.from_numpy(smplx_data['smplx_global_orient'])
body_pose = torch.from_numpy(smplx_data['smplx_body_pose'])
N = transl.shape[0]
motion_params = dict(
transl=transl,
global_orient=global_orient,
body_pose=body_pose,
betas=betas.unsqueeze(0).repeat(N, 1).float()
)
# 1. process positions
frame_params = {k: v.to(device) for k, v in motion_params.items()}
frame_params['leye_pose'] = torch.zeros((N, 3)).to(device)
frame_params['reye_pose'] = torch.zeros((N, 3)).to(device)
frame_params['left_hand_pose'] = torch.zeros((N, 45)).to(device)
frame_params['right_hand_pose'] = torch.zeros((N, 45)).to(device)
frame_params['jaw_pose'] = torch.zeros((N, 3)).to(device)
frame_params['expression'] = torch.zeros((N, 100)).to(device)
output = smplx_model(**frame_params)
position_data = output.joints.detach().cpu()[:, :22] # T, 22 ,3
position_val_data = position_data[1:] - position_data[:-1]
root_idx = 0
# put on floor and put root on origin for the first frame
ori = deepcopy(position_data[0, root_idx]) # first frame root position
y_min = torch.min(position_data[:, :, 1])
ori[1] = y_min
position_data = position_data - ori
velocities_root = position_data[1:, root_idx, :] - position_data[:-1, root_idx, :]
position_data[:,:,0] -= position_data[:,0:1,0]
position_data[:,:,2] -= position_data[:,0:1,2]
T, njoint, _ = position_data.shape
final_x = torch.zeros((T, 2 + 6 + njoint * 3 + njoint * 3))
final_x[1:, 0] = velocities_root[:, 0]
final_x[1:, 1] = velocities_root[:, 1]
final_x[:, 2:2+6] = matrix_to_rotation_6d(global_orient)
final_x[:, 8:8+njoint*3] = position_data.flatten(1, 2)
final_x[1:, 8+njoint*3:8+njoint*6] = position_val_data.flatten(1, 2) # T, 140
return final_x
def get_g1_motion(data_path):
global _debug_save_count
rotation_matrix = torch.tensor([[1.0, 0, 0], [0, 0, -1], [0, 1, 0]]).inverse()
g1_data = joblib.load(data_path)
dof = g1_data['g1_dof'] # T, 29
global_orient = g1_data['g1_root_ori'] # T, 4, wxyz
joints = g1_data['g1_joints'] # T, 30, 3
# save before-transform data for first 3 motions
if _debug_save_count < 3:
os.makedirs(_debug_save_dir, exist_ok=True)
np.savez(
os.path.join(_debug_save_dir, f'motion_{_debug_save_count:02d}_before.npz'),
g1_trans=joints,
g1_root_rot=global_orient,
g1_dof=dof,
)
global_orient_mat = quaternion_to_matrix(torch.from_numpy(global_orient)).float()
global_orient_mat = torch.einsum('ij,tjk->tik', rotation_matrix, global_orient_mat)
global_orient = matrix_to_axis_angle(global_orient_mat)
position_data = torch.einsum('ij,tkj->tki', rotation_matrix, torch.from_numpy(joints).float())
position_val_data = position_data[1:] - position_data[:-1]
# save after-transform data for first 3 motions
if _debug_save_count < 3:
np.savez(
os.path.join(_debug_save_dir, f'motion_{_debug_save_count:02d}_after.npz'),
g1_trans=position_data,
g1_root_rot=global_orient.numpy(),
g1_dof=dof,
)
root_idx = 0
# put on floor and put root on origin for the first frame
ori = deepcopy(position_data[0, root_idx]) # first frame root position
y_min = torch.min(position_data[:, :, 1])
ori[1] = y_min
position_data = position_data - ori
velocities_root = position_data[1:, root_idx, :] - position_data[:-1, root_idx, :]
position_data[:,:,0] -= position_data[:,0:1,0]
position_data[:,:,2] -= position_data[:,0:1,2]
T, njoint, _ = position_data.shape
final_x = torch.zeros((T, 2 + 6 + njoint * 3 + njoint * 3))
final_x[1:, 0] = velocities_root[:, 0]
final_x[1:, 1] = velocities_root[:, 1]
final_x[:, 2:2+6] = matrix_to_rotation_6d(global_orient)
final_x[:, 8:8+njoint*3] = position_data.flatten(1, 2)
final_x[1:, 8+njoint*3:8+njoint*6] = position_val_data.flatten(1, 2) # T, 140
final_x = torch.concat([final_x, torch.from_numpy(dof)], dim=-1)
return final_x # 217
# python -m data.motionmillion.tools.process --save_root "data/motionmillion/final_data"
if __name__ == '__main__':
with open('data/merge_gmr_retarget_smplx_path.txt', 'r') as f:
paths = f.readlines()
for line in tqdm(paths[args.start_idx::args.interval]):
data_path = line.strip()
smplx_motion = get_smplx_motion(data_path)
g1_motion = get_g1_motion(data_path)
if smplx_motion.shape[0] != g1_motion.shape[0]:
min_len = min(smplx_motion.shape[0], g1_motion.shape[0])
smplx_motion = smplx_motion[:min_len]
g1_motion = g1_motion[:min_len]
data = dict(
g1_motion=g1_motion,
smplx_motion=smplx_motion,
)
data_path = data_path.replace('.npz', '.pkl')
save_path = args.save_root + '/motions/' + '/'.join(data_path.split('/')[2:])
# os.makedirs(os.path.dirname(save_path), exist_ok=True)
# joblib.dump(data, save_path)
|