from utils.skeleton import Skeleton from utils.quaternion import * from utils.motion_process import t2m_kinematic_chain, t2m_raw_offsets import torch import os n_raw_offsets = torch.from_numpy(t2m_raw_offsets) kinematic_chain = t2m_kinematic_chain l_idx1, l_idx2 = 5, 8 face_joint_indx = [2, 1, 17, 16] # Lazy loading of tgt_offsets - only needed when is_mesh=True _tgt_offsets_cache = None def _get_tgt_offsets(): """Lazily load target offsets for mesh processing. Only called when is_mesh=True.""" global _tgt_offsets_cache if _tgt_offsets_cache is None: example_data_path = os.path.join("datasets/HumanML3D/new_joints", "000021" + '.npy') if not os.path.exists(example_data_path): raise FileNotFoundError( f"Example data file not found: {example_data_path}\n" "This file is only needed for mesh-level motion generation (is_mesh=True).\n" "For regular text-to-motion generation, use is_mesh=False." ) example_data = np.load(example_data_path) example_data = example_data.reshape(len(example_data), -1, 3) example_data = torch.from_numpy(example_data) tgt_skel = Skeleton(n_raw_offsets, kinematic_chain, 'cpu') _tgt_offsets_cache = tgt_skel.get_offsets_joints(example_data[0]) return _tgt_offsets_cache l_idx1, l_idx2 = 5, 8 fid_r, fid_l = [8, 11], [7, 10] r_hip, l_hip = 2, 1 joints_num = 22 def uniform_skeleton(positions, target_offset): src_skel = Skeleton(n_raw_offsets, kinematic_chain, 'cpu') src_offset = src_skel.get_offsets_joints(torch.from_numpy(positions[0])) src_offset = src_offset.numpy() tgt_offset = target_offset.numpy() # print(src_offset) # print(tgt_offset) '''Calculate Scale Ratio as the ratio of legs''' src_leg_len = np.abs(src_offset[l_idx1]).max() + np.abs(src_offset[l_idx2]).max() tgt_leg_len = np.abs(tgt_offset[l_idx1]).max() + np.abs(tgt_offset[l_idx2]).max() scale_rt = tgt_leg_len / src_leg_len # print(scale_rt) src_root_pos = positions[:, 0] tgt_root_pos = src_root_pos * scale_rt '''Inverse Kinematics''' quat_params = src_skel.inverse_kinematics_np(positions, face_joint_indx) # print(quat_params.shape) '''Forward Kinematics''' src_skel.set_offset(target_offset) new_joints = src_skel.forward_kinematics_np(quat_params, tgt_root_pos) return new_joints def process_file(positions, feet_thre, is_mesh=False): # (seq_len, joints_num, 3) # '''Down Sample''' # positions = positions[::ds_num] if is_mesh: '''Uniform Skeleton''' tgt_offsets = _get_tgt_offsets() # Lazy load only when needed positions = uniform_skeleton(positions, tgt_offsets) '''Put on Floor''' floor_height = positions.min(axis=0).min(axis=0)[1] positions[:, :, 1] -= floor_height # print(floor_height) # plot_3d_motion("./positions_1.mp4", kinematic_chain, positions, 'title', fps=20) '''XZ at origin''' root_pos_init = positions[0] root_pose_init_xz = root_pos_init[0] * np.array([1, 0, 1]) positions = positions - root_pose_init_xz # '''Move the first pose to origin ''' # root_pos_init = positions[0] # positions = positions - root_pos_init[0] '''All initially face Z+''' r_hip, l_hip, sdr_r, sdr_l = face_joint_indx across1 = root_pos_init[r_hip] - root_pos_init[l_hip] across2 = root_pos_init[sdr_r] - root_pos_init[sdr_l] across = across1 + across2 across = across / np.sqrt((across ** 2).sum(axis=-1))[..., np.newaxis] # forward (3,), rotate around y-axis forward_init = np.cross(np.array([[0, 1, 0]]), across, axis=-1) # forward (3,) forward_init = forward_init / np.sqrt((forward_init ** 2).sum(axis=-1))[..., np.newaxis] # print(forward_init) target = np.array([[0, 0, 1]]) root_quat_init = qbetween_np(forward_init, target) root_quat_init = np.ones(positions.shape[:-1] + (4,)) * root_quat_init positions_b = positions.copy() positions = qrot_np(root_quat_init, positions) # plot_3d_motion("./positions_2.mp4", kinematic_chain, positions, 'title', fps=20) '''New ground truth positions''' global_positions = positions.copy() # plt.plot(positions_b[:, 0, 0], positions_b[:, 0, 2], marker='*') # plt.plot(positions[:, 0, 0], positions[:, 0, 2], marker='o', color='r') # plt.xlabel('x') # plt.ylabel('z') # plt.axis('equal') # plt.show() """ Get Foot Contacts """ def foot_detect(positions, thres): velfactor, heightfactor = np.array([thres, thres]), np.array([3.0, 2.0]) feet_l_x = (positions[1:, fid_l, 0] - positions[:-1, fid_l, 0]) ** 2 feet_l_y = (positions[1:, fid_l, 1] - positions[:-1, fid_l, 1]) ** 2 feet_l_z = (positions[1:, fid_l, 2] - positions[:-1, fid_l, 2]) ** 2 # feet_l_h = positions[:-1,fid_l,1] # feet_l = (((feet_l_x + feet_l_y + feet_l_z) < velfactor) & (feet_l_h < heightfactor)).astype(np.float) feet_l = ((feet_l_x + feet_l_y + feet_l_z) < velfactor).astype(np.float32) feet_r_x = (positions[1:, fid_r, 0] - positions[:-1, fid_r, 0]) ** 2 feet_r_y = (positions[1:, fid_r, 1] - positions[:-1, fid_r, 1]) ** 2 feet_r_z = (positions[1:, fid_r, 2] - positions[:-1, fid_r, 2]) ** 2 # feet_r_h = positions[:-1,fid_r,1] # feet_r = (((feet_r_x + feet_r_y + feet_r_z) < velfactor) & (feet_r_h < heightfactor)).astype(np.float) feet_r = (((feet_r_x + feet_r_y + feet_r_z) < velfactor)).astype(np.float32) return feet_l, feet_r # feet_l, feet_r = foot_detect(positions, feet_thre) # feet_l, feet_r = foot_detect(positions, 0.002) '''Quaternion and Cartesian representation''' r_rot = None def get_rifke(positions): '''Local pose''' positions[..., 0] -= positions[:, 0:1, 0] positions[..., 2] -= positions[:, 0:1, 2] '''All pose face Z+''' positions = qrot_np(np.repeat(r_rot[:, None], positions.shape[1], axis=1), positions) return positions def get_quaternion(positions): skel = Skeleton(n_raw_offsets, kinematic_chain, "cpu") # (seq_len, joints_num, 4) quat_params = skel.inverse_kinematics_np(positions, face_joint_indx, smooth_forward=False) '''Fix Quaternion Discontinuity''' quat_params = qfix(quat_params) # (seq_len, 4) r_rot = quat_params[:, 0].copy() # print(r_rot[0]) '''Root Linear Velocity''' # (seq_len - 1, 3) velocity = (positions[1:, 0] - positions[:-1, 0]).copy() # print(r_rot.shape, velocity.shape) velocity = qrot_np(r_rot[1:], velocity) '''Root Angular Velocity''' # (seq_len - 1, 4) r_velocity = qmul_np(r_rot[1:], qinv_np(r_rot[:-1])) quat_params[1:, 0] = r_velocity # (seq_len, joints_num, 4) return quat_params, r_velocity, velocity, r_rot def get_cont6d_params(positions): skel = Skeleton(n_raw_offsets, kinematic_chain, "cpu") # (seq_len, joints_num, 4) quat_params = skel.inverse_kinematics_np(positions, face_joint_indx, smooth_forward=True) '''Quaternion to continuous 6D''' cont_6d_params = quaternion_to_cont6d_np(quat_params) # (seq_len, 4) r_rot = quat_params[:, 0].copy() # print(r_rot[0]) '''Root Linear Velocity''' # (seq_len - 1, 3) velocity = (positions[1:, 0] - positions[:-1, 0]).copy() # print(r_rot.shape, velocity.shape) velocity = qrot_np(r_rot[1:], velocity) '''Root Angular Velocity''' # (seq_len - 1, 4) r_velocity = qmul_np(r_rot[1:], qinv_np(r_rot[:-1])) # (seq_len, joints_num, 4) return cont_6d_params, r_velocity, velocity, r_rot cont_6d_params, r_velocity, velocity, r_rot = get_cont6d_params(positions) positions = get_rifke(positions) # trejec = np.cumsum(np.concatenate([np.array([[0, 0, 0]]), velocity], axis=0), axis=0) # r_rotations, r_pos = recover_ric_glo_np(r_velocity, velocity[:, [0, 2]]) # plt.plot(positions_b[:, 0, 0], positions_b[:, 0, 2], marker='*') # plt.plot(ground_positions[:, 0, 0], ground_positions[:, 0, 2], marker='o', color='r') # plt.plot(trejec[:, 0], trejec[:, 2], marker='^', color='g') # plt.plot(r_pos[:, 0], r_pos[:, 2], marker='s', color='y') # plt.xlabel('x') # plt.ylabel('z') # plt.axis('equal') # plt.show() '''Root height''' root_y = positions[:, 0, 1:2] '''Root rotation and linear velocity''' # (seq_len-1, 1) rotation velocity along y-axis # (seq_len-1, 2) linear velovity on xz plane r_velocity = np.arcsin(r_velocity[:, 2:3]) l_velocity = velocity[:, [0, 2]] # print(r_velocity.shape, l_velocity.shape, root_y.shape) root_data = np.concatenate([r_velocity, l_velocity, root_y[:-1]], axis=-1) '''Get Joint Rotation Representation''' # (seq_len, (joints_num-1) *6) quaternion for skeleton joints rot_data = cont_6d_params[:, 1:].reshape(len(cont_6d_params), -1) '''Get Joint Rotation Invariant Position Represention''' # (seq_len, (joints_num-1)*3) local joint position ric_data = positions[:, 1:].reshape(len(positions), -1) '''Get Joint Velocity Representation''' # (seq_len-1, joints_num*3) local_vel = qrot_np(np.repeat(r_rot[:-1, None], global_positions.shape[1], axis=1), global_positions[1:] - global_positions[:-1]) local_vel = local_vel.reshape(len(local_vel), -1) data = root_data data = np.concatenate([data, ric_data[:-1]], axis=-1) data = np.concatenate([data, rot_data[:-1]], axis=-1) # print(data.shape, local_vel.shape) data = np.concatenate([data, local_vel], axis=-1) data = np.concatenate([data, feet_l, feet_r], axis=-1) return data, global_positions, positions, l_velocity def back_process(data, is_mesh=False): data, ground_positions, positions, l_velocity = process_file(data, 0.002, is_mesh=is_mesh) return data[:, :67]