File size: 4,288 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 | 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, axis_angle_to_6d
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")
parser = ArgumentParser(description="Launch MoCap processing")
parser.add_argument('--save_root', type=str, default="data/smplx_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 get_smplx_motion(data_path):
data = dict(np.load(data_path))
import ipdb; ipdb.set_trace()
smplx_data = dict(
smplx_transl=data["transl"],
smplx_global_orient=data["global_orient"],
smplx_body_pose=data["body_pose"],
)
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
# import ipdb; ipdb.set_trace()
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, :]
# import ipdb; ipdb.set_trace()
position_data_cp = deepcopy(position_data)
position_data[:,:,0] -= position_data_cp[:,0:1,0]
position_data[:,:,2] -= position_data_cp[:,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] = axis_angle_to_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
# import ipdb; ipdb.set_trace()
# vis_3d_motion([position_data.numpy()], None, ['video.mp4'], fps=30)
return final_x
# python -m data.motionmillion.tools.process --save_root "data/motionmillion/final_data"
if __name__ == '__main__':
with open('data/smplx_path.txt', 'r') as f:
paths = f.readlines()
for line in tqdm(paths[args.start_idx::args.interval]):
data_path = line.strip()
save_path = data_path.replace('.npz', '.npy').replace('/mnt/shenzhen2cephfs/capybarali/codes/neobot/data/motionmillion/1_filtered_smplx/', '')
save_path = args.save_root + '/' + save_path
if os.path.exists(save_path):
continue
smplx_motion = get_smplx_motion(data_path)
os.makedirs(os.path.dirname(save_path), exist_ok=True)
np.save(save_path, smplx_motion)
# joblib.dump(data, save_path)
|