import numpy as np import torch from ..builder import PIPELINES from ..quaternion import qbetween_np, qinv_np, qmul_np, qrot_np face_joint_indx = [2, 1, 17, 16] fid_l = [7, 10] fid_r = [8, 11] trans_matrix = torch.Tensor([[1.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, -1.0, 0.0]]) def rigid_transform(relative, data): global_positions = data[..., :22 * 3].reshape(data.shape[:-1] + (22, 3)) global_vel = data[..., 22 * 3:22 * 6].reshape(data.shape[:-1] + (22, 3)) relative_rot = relative[0] relative_t = relative[1:3] relative_r_rot_quat = np.zeros(global_positions.shape[:-1] + (4, )) relative_r_rot_quat[..., 0] = np.cos(relative_rot) relative_r_rot_quat[..., 2] = np.sin(relative_rot) global_positions = qrot_np(qinv_np(relative_r_rot_quat), global_positions) global_positions[..., [0, 2]] += relative_t data[..., :22 * 3] = global_positions.reshape(data.shape[:-1] + (-1, )) global_vel = qrot_np(qinv_np(relative_r_rot_quat), global_vel) data[..., 22 * 3:22 * 6] = global_vel.reshape(data.shape[:-1] + (-1, )) return data @PIPELINES.register_module() class SwapSiameseMotion(object): r"""Swap motion sequences. Args: prob (float): The probability of swapping siamese motions """ def __init__(self, prob=0.5): self.prob = prob assert prob >= 0 and prob <= 1.0 def __call__(self, results): if np.random.rand() <= self.prob: motion1 = results['motion1'] motion2 = results['motion2'] results['motion1'] = motion2 results['motion2'] = motion1 return results def __repr__(self): repr_str = self.__class__.__name__ + f'(prob={self.prob})' return repr_str @PIPELINES.register_module() class ProcessSiameseMotion(object): r"""Process siamese motion sequences. The code is borrowed from https://github.com/tr3e/InterGen/blob/master/utils/utils.py """ def __init__(self, feet_threshold, prev_frames, n_joints, prob): self.feet_threshold = feet_threshold self.prev_frames = prev_frames self.n_joints = n_joints self.prob = prob def process_single_motion(self, motion): feet_thre = self.feet_threshold prev_frames = self.prev_frames n_joints = self.n_joints '''Uniform Skeleton''' # positions = uniform_skeleton(positions, tgt_offsets) positions = motion[:, :n_joints * 3].reshape(-1, n_joints, 3) rotations = motion[:, n_joints * 3:] positions = np.einsum("mn, tjn->tjm", trans_matrix, positions) '''Put on Floor''' floor_height = positions.min(axis=0).min(axis=0)[1] positions[:, :, 1] -= floor_height '''XZ at origin''' root_pos_init = positions[prev_frames] root_pose_init_xz = root_pos_init[0] * np.array([1, 0, 1]) positions = positions - root_pose_init_xz '''All initially face Z+''' r_hip, l_hip, sdr_r, sdr_l = face_joint_indx across = root_pos_init[r_hip] - root_pos_init[l_hip] 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)) forward_init = forward_init[..., np.newaxis] target = np.array([[0, 0, 1]]) root_quat_init = qbetween_np(forward_init, target) root_quat_init_for_all = \ np.ones(positions.shape[:-1] + (4,)) * root_quat_init positions = qrot_np(root_quat_init_for_all, positions) """ Get Foot Contacts """ def foot_detect(positions, thres): velfactor, heightfactor = \ np.array([thres, thres]), np.array([0.12, 0.05]) 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_sum = feet_l_x + feet_l_y + feet_l_z feet_l = ((feet_l_sum < velfactor) & (feet_l_h < heightfactor)) feet_l = feet_l.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_sum = feet_r_x + feet_r_y + feet_r_z feet_r = ((feet_r_sum < velfactor) & (feet_r_h < heightfactor)) feet_r = feet_r.astype(np.float32) return feet_l, feet_r feet_l, feet_r = foot_detect(positions, feet_thre) '''Get Joint Rotation Representation''' rot_data = rotations '''Get Joint Rotation Invariant Position Represention''' joint_positions = positions.reshape(len(positions), -1) joint_vels = positions[1:] - positions[:-1] joint_vels = joint_vels.reshape(len(joint_vels), -1) data = joint_positions[:-1] data = np.concatenate([data, joint_vels], axis=-1) data = np.concatenate([data, rot_data[:-1]], axis=-1) data = np.concatenate([data, feet_l, feet_r], axis=-1) return data, root_quat_init, root_pose_init_xz[None] def __call__(self, results): motion1, root_quat_init1, root_pos_init1 = \ self.process_single_motion(results['motion1']) motion2, root_quat_init2, root_pos_init2 = \ self.process_single_motion(results['motion2']) r_relative = qmul_np(root_quat_init2, qinv_np(root_quat_init1)) angle = np.arctan2(r_relative[:, 2:3], r_relative[:, 0:1]) xz = qrot_np(root_quat_init1, root_pos_init2 - root_pos_init1)[:, [0, 2]] relative = np.concatenate([angle, xz], axis=-1)[0] motion2 = rigid_transform(relative, motion2) if np.random.rand() <= self.prob: motion2, motion1 = motion1, motion2 motion = np.concatenate((motion1, motion2), axis=-1) results['motion'] = motion return results def __repr__(self): repr_str = self.__class__.__name__ repr_str += f'(feet_threshold={self.feet_threshold})' repr_str += f'(feet_threshold={self.feet_threshold})' repr_str += f'(n_joints={self.n_joints})' repr_str += f'(prob={self.prob})' return repr_str