""" Joint to HumanML3D Representation Converter This module provides a class to convert joint positions (shape: [T, 22, 3]) to HumanML3D representation format (shape: [T, 263]). """ import numpy as np import torch from os.path import join as pjoin from .humanml_skeleton import Skeleton from .quaternion import * from .humanml_paramUtil import t2m_raw_offsets, t2m_kinematic_chain class JointToHumanML3D: """ Convert joint positions to HumanML3D representation format. The output format is a 263-dimensional feature vector containing: - Root rotation velocity (1D): rotation velocity along y-axis - Root linear velocity (2D): linear velocity on xz plane - Root height (1D): y-coordinate of root joint - Joint rotation invariant positions (63D): local positions of 21 joints - Joint rotations (126D): continuous 6D rotation representation for 21 joints - Joint velocities (66D): local velocities of 22 joints - Foot contacts (4D): contact labels for left/right foot Total: 1 + 2 + 1 + 63 + 126 + 66 + 4 = 263 """ def __init__(self, example_id="000021", data_dir='./joints/'): """ Initialize the converter with pre-loaded target skeleton. Args: example_id (str): The example motion ID to extract target skeleton offsets data_dir (str): Directory containing the example motion data """ # Lower legs indices for scale calculation self.l_idx1, self.l_idx2 = 5, 8 # Right/Left foot indices for foot contact detection self.fid_r, self.fid_l = [8, 11], [7, 10] # Face direction joints: r_hip, l_hip, sdr_r, sdr_l self.face_joint_indx = [2, 1, 17, 16] # Hip indices self.r_hip, self.l_hip = 2, 1 # Number of joints self.joints_num = 22 # Load kinematic chain and raw offsets self.n_raw_offsets = torch.from_numpy(t2m_raw_offsets) self.kinematic_chain = t2m_kinematic_chain # Get target skeleton offsets from example example_data = np.load(pjoin(data_dir, example_id + '.npy')) example_data = example_data.reshape(len(example_data), -1, 3) example_data = torch.from_numpy(example_data) tgt_skel = Skeleton(self.n_raw_offsets, self.kinematic_chain, 'cpu') self.tgt_offsets = tgt_skel.get_offsets_joints(example_data[0]) print(f"JointToHumanML3D initialized with target skeleton from {example_id}") def uniform_skeleton(self, positions, target_offset): """ Normalize skeleton to target proportions using leg length scaling. Args: positions (np.ndarray): Joint positions, shape [T, joints_num, 3] target_offset (torch.Tensor): Target skeleton offsets Returns: np.ndarray: Normalized joint positions """ src_skel = Skeleton(self.n_raw_offsets, self.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() # Calculate scale ratio based on leg lengths src_leg_len = np.abs(src_offset[self.l_idx1]).max() + np.abs(src_offset[self.l_idx2]).max() tgt_leg_len = np.abs(tgt_offset[self.l_idx1]).max() + np.abs(tgt_offset[self.l_idx2]).max() scale_rt = tgt_leg_len / src_leg_len src_root_pos = positions[:, 0] tgt_root_pos = src_root_pos * scale_rt # Inverse Kinematics quat_params = src_skel.inverse_kinematics_np(positions, self.face_joint_indx) # Forward Kinematics with target skeleton src_skel.set_offset(target_offset) new_joints = src_skel.forward_kinematics_np(quat_params, tgt_root_pos) return new_joints def process_joints(self, positions, feet_thre=0.002): """ Process joint positions and convert to HumanML3D representation. Args: positions (np.ndarray): Joint positions, shape [T, joints_num, 3] feet_thre (float): Threshold for foot contact detection Returns: np.ndarray: HumanML3D representation, shape [T-1, 263] """ # Uniform skeleton normalization positions = self.uniform_skeleton(positions, self.tgt_offsets) # Put on floor floor_height = positions.min(axis=0).min(axis=0)[1] positions[:, :, 1] -= floor_height # Center 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 # Align all poses to initially face Z+ r_hip, l_hip, sdr_r, sdr_l = self.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] # Calculate forward direction (rotate around y-axis) forward_init = np.cross(np.array([[0, 1, 0]]), across, axis=-1) forward_init = forward_init / np.sqrt((forward_init ** 2).sum(axis=-1))[..., np.newaxis] # Rotate to face Z+ 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 = qrot_np(root_quat_init, positions) # Store global positions global_positions = positions.copy() # Detect foot contacts feet_l, feet_r = self._foot_detect(positions, feet_thre) # Get continuous 6D representation cont_6d_params, r_velocity, velocity, r_rot = self._get_cont6d_params(positions) # Get rotation invariant position representation positions = self._get_rifke(positions, r_rot) # Root height root_y = positions[:, 0, 1:2] # Root rotation and linear velocity r_velocity = np.arcsin(r_velocity[:, 2:3]) # (T-1, 1) l_velocity = velocity[:, [0, 2]] # (T-1, 2) root_data = np.concatenate([r_velocity, l_velocity, root_y[:-1]], axis=-1) # (T-1, 4) # Joint rotation representation rot_data = cont_6d_params[:, 1:].reshape(len(cont_6d_params), -1) # (T, 126) # Joint rotation invariant position representation ric_data = positions[:, 1:].reshape(len(positions), -1) # (T, 63) # Joint velocity representation 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) # (T-1, 66) # Concatenate all features data = root_data # (T-1, 4) data = np.concatenate([data, ric_data[:-1]], axis=-1) # (T-1, 4+63) data = np.concatenate([data, rot_data[:-1]], axis=-1) # (T-1, 4+63+126) data = np.concatenate([data, local_vel], axis=-1) # (T-1, 4+63+126+66) data = np.concatenate([data, feet_l, feet_r], axis=-1) # (T-1, 263) return data def _foot_detect(self, positions, thres): """ Detect foot contacts based on velocity threshold. Args: positions (np.ndarray): Joint positions thres (float): Velocity threshold Returns: tuple: (feet_l, feet_r) contact labels for left and right feet """ velfactor = np.array([thres, thres]) # Left foot feet_l_x = (positions[1:, self.fid_l, 0] - positions[:-1, self.fid_l, 0]) ** 2 feet_l_y = (positions[1:, self.fid_l, 1] - positions[:-1, self.fid_l, 1]) ** 2 feet_l_z = (positions[1:, self.fid_l, 2] - positions[:-1, self.fid_l, 2]) ** 2 feet_l = ((feet_l_x + feet_l_y + feet_l_z) < velfactor).astype(np.float32) # Right foot feet_r_x = (positions[1:, self.fid_r, 0] - positions[:-1, self.fid_r, 0]) ** 2 feet_r_y = (positions[1:, self.fid_r, 1] - positions[:-1, self.fid_r, 1]) ** 2 feet_r_z = (positions[1:, self.fid_r, 2] - positions[:-1, self.fid_r, 2]) ** 2 feet_r = ((feet_r_x + feet_r_y + feet_r_z) < velfactor).astype(np.float32) return feet_l, feet_r def _get_cont6d_params(self, positions): """ Get continuous 6D rotation parameters. Args: positions (np.ndarray): Joint positions Returns: tuple: (cont_6d_params, r_velocity, velocity, r_rot) """ skel = Skeleton(self.n_raw_offsets, self.kinematic_chain, "cpu") quat_params = skel.inverse_kinematics_np(positions, self.face_joint_indx, smooth_forward=True) # Quaternion to continuous 6D cont_6d_params = quaternion_to_cont6d_np(quat_params) # Root rotation r_rot = quat_params[:, 0].copy() # Root linear velocity velocity = (positions[1:, 0] - positions[:-1, 0]).copy() velocity = qrot_np(r_rot[1:], velocity) # Root angular velocity r_velocity = qmul_np(r_rot[1:], qinv_np(r_rot[:-1])) return cont_6d_params, r_velocity, velocity, r_rot def _get_rifke(self, positions, r_rot): """ Get rotation invariant position representation. Args: positions (np.ndarray): Joint positions r_rot (np.ndarray): Root rotations Returns: np.ndarray: Rotation invariant positions """ # Local pose (relative to root XZ) positions[..., 0] -= positions[:, 0:1, 0] positions[..., 2] -= positions[:, 0:1, 2] # All poses face Z+ positions = qrot_np(np.repeat(r_rot[:, None], positions.shape[1], axis=1), positions) return positions def convert(self, joint_positions, valid_frames): """ Convert joint positions to HumanML3D representation. Args: joint_positions (np.ndarray or torch.Tensor): Joint positions with shape [T, 22, 3] valid_frames (int): Number of valid frames to use. If provided and less than T, only the first valid_frames will be processed, and the output will be padded to match the expected length [T-1, 263]. If None, all frames are processed. Returns: np.ndarray: HumanML3D representation with shape [T-1, 263] """ # Convert to numpy if input is torch tensor if isinstance(joint_positions, torch.Tensor): joint_positions = joint_positions.cpu().numpy() # Validate input shape if len(joint_positions.shape) != 3: raise ValueError(f"Expected 3D input with shape [T, 22, 3], got shape {joint_positions.shape}") if joint_positions.shape[1] != 22 or joint_positions.shape[2] != 3: raise ValueError(f"Expected shape [T, 22, 3], got {joint_positions.shape}") original_length = joint_positions.shape[0] # Truncate to valid frames if specified if valid_frames is not None and valid_frames < original_length: joint_positions = joint_positions[:valid_frames, ...] else: valid_frames = original_length # Process joints (this will reduce length by 1 due to velocity calculation) humanml3d_repr = self.process_joints(joint_positions, feet_thre=0.002) # humanml3d_repr shape: [valid_frames-1, 263] # Padding to match expected output length [original_length-1, 263] # if frames were truncated if valid_frames < original_length: target_length = original_length - 1 current_length = humanml3d_repr.shape[0] # valid_frames - 1 if current_length < target_length: # Pad with the last valid frame pad_length = target_length - current_length last_frame = humanml3d_repr[-1:, :] # [1, 263] padding = np.repeat(last_frame, pad_length, axis=0) # [pad_length, 263] humanml3d_repr = np.concatenate([humanml3d_repr, padding], axis=0) return humanml3d_repr if __name__ == "__main__": # Example usage converter = JointToHumanML3D(example_id="000021", data_dir='./joints/') # Load example joint data example_joints = np.load('./joints/000021.npy')[:, :22, :] example_joints = example_joints.reshape(-1, 22, 3) print(f"Input shape: {example_joints.shape}") # Test 1: Convert without truncation humanml3d_data = converter.convert(example_joints, valid_frames=None) print(f"Output shape (no truncation): {humanml3d_data.shape}") print(f"Expected output shape: [{example_joints.shape[0]-1}, 263]") # Test 2: Convert with truncation and padding if example_joints.shape[0] > 50: valid_frames = 50 humanml3d_data_truncated = converter.convert(example_joints, valid_frames=valid_frames) print(f"\nOutput shape (truncated to {valid_frames} frames): {humanml3d_data_truncated.shape}") print(f"Expected shape: [{example_joints.shape[0]-1}, 263]") print(f"Frames processed: {valid_frames-1}, Frames padded: {example_joints.shape[0]-1-(valid_frames-1)}")