| import torch |
| import numpy as np |
| from .rotation_conversions import rotation_6d_to_matrix, matrix_to_axis_angle |
|
|
| def accumulate_rotations(relative_rotations): |
| R_total = [relative_rotations[0]] |
| for R_rel in relative_rotations[1:]: |
| R_total.append(np.matmul(R_rel, R_total[-1])) |
| return np.array(R_total) |
|
|
| def accumulate_rotations_batch(R_rel): |
| B, T, _, _ = R_rel.shape |
|
|
| R_list = [] |
| R_prev = R_rel[:, 0] |
| R_list.append(R_prev) |
|
|
| for t in range(1, T): |
| R_curr = R_rel[:, t] @ R_prev |
| R_list.append(R_curr) |
| R_prev = R_curr |
|
|
| R_out = torch.stack(R_list, dim=1) |
| return R_out |
|
|
| def recover_from_local_position_fast(final_x: torch.Tensor, njoint=22): |
| """ |
| final_x: (B, T, D) |
| return: (B, T, J, 3) |
| """ |
| B, T, _ = final_x.shape |
| device = final_x.device |
|
|
| |
| positions_no_heading = final_x[:, :, 8:8+3*njoint].reshape(B, T, njoint, 3) |
| velocities_root_xy = final_x[:, :, :2] |
| global_heading_diff_rot = final_x[:, :, 2:8] |
|
|
| |
| R_rel = rotation_6d_to_matrix(global_heading_diff_rot) |
| R_global = accumulate_rotations_batch(R_rel) |
| R_inv = R_global.transpose(-1, -2) |
|
|
| |
| pos = torch.matmul( |
| R_inv[:, :, None, :, :], |
| positions_no_heading[..., None] |
| ).squeeze(-1) |
|
|
| |
| vel_root = torch.zeros((B, T, 3), device=device) |
| vel_root[:, :, 0] = velocities_root_xy[:, :, 0] |
| vel_root[:, :, 2] = velocities_root_xy[:, :, 1] |
|
|
| vel_root[:, 1:] = torch.matmul( |
| R_inv[:, :-1], |
| vel_root[:, 1:, :, None] |
| ).squeeze(-1) |
|
|
| |
| root_translation = torch.cumsum(vel_root, dim=1) |
|
|
| pos[:, :, :, 0] += root_translation[:, :, 0:1] |
| pos[:, :, :, 2] += root_translation[:, :, 2:3] |
|
|
| return pos |
|
|
| def recover_from_local_position(final_x, njoint): |
|
|
| if final_x.ndim == 3: |
| bs, nfrm, _ = final_x.shape |
| is_batched = True |
| else: |
| nfrm, _ = final_x.shape |
| bs = 1 |
| is_batched = False |
| final_x = final_x.reshape(1, *final_x.shape) |
|
|
| |
| positions_no_heading = final_x[:,:,8:8+3*njoint].reshape(bs, nfrm, njoint, 3) |
| velocities_root_xy_no_heading = final_x[:,:,:2] |
| global_heading_diff_rot = final_x[:,:,2:8] |
| |
| positions_with_heading = [] |
| for b in range(bs): |
| |
| global_heading_rot = accumulate_rotations(rotation_6d_to_matrix(torch.from_numpy(global_heading_diff_rot[b])).numpy()) |
| inv_global_heading_rot = np.transpose(global_heading_rot, (0, 2, 1)) |
| |
| curr_pos_with_heading = np.matmul(np.repeat(inv_global_heading_rot[:, None,:, :], njoint, axis=1), |
| positions_no_heading[b][...,None]).squeeze(-1) |
| |
| velocities_root_xyz_no_heading = np.zeros((velocities_root_xy_no_heading[b].shape[0], 3)) |
| velocities_root_xyz_no_heading[:, 0] = velocities_root_xy_no_heading[b, :, 0] |
| velocities_root_xyz_no_heading[:, 2] = velocities_root_xy_no_heading[b, :, 1] |
| velocities_root_xyz_no_heading[1:, :] = np.matmul(inv_global_heading_rot[:-1], |
| velocities_root_xyz_no_heading[1:, :,None]).squeeze(-1) |
|
|
| root_translation = np.cumsum(velocities_root_xyz_no_heading, axis=0) |
| |
| curr_pos_with_heading[:, :, 0] += root_translation[:, 0:1] |
| curr_pos_with_heading[:, :, 2] += root_translation[:, 2:] |
| |
| positions_with_heading.append(curr_pos_with_heading) |
|
|
| positions_with_heading = np.stack(positions_with_heading, axis=0) |
|
|
| if not is_batched: |
| positions_with_heading = positions_with_heading.squeeze(0) |
|
|
| return positions_with_heading |
|
|
| def recover_from_local_rotation(final_x, njoint): |
| nfrm, _ = final_x.shape |
| rotations_matrix = rotation_6d_to_matrix(torch.from_numpy(final_x[:,8+6*njoint:8+12*njoint]).reshape(nfrm, -1, 6)).numpy() |
| global_heading_diff_rot = final_x[:,2:8] |
| velocities_root_xy_no_heading = final_x[:,:2] |
| positions_no_heading = final_x[:, 8:8+3*njoint].reshape(nfrm, -1, 3) |
| height = positions_no_heading[:, 0, 1] |
|
|
| global_heading_rot = accumulate_rotations(rotation_6d_to_matrix(torch.from_numpy(global_heading_diff_rot)).numpy()) |
| inv_global_heading_rot = np.transpose(global_heading_rot, (0, 2, 1)) |
| |
| rotations_matrix[:,0,...] = np.matmul(inv_global_heading_rot, rotations_matrix[:,0,...]) |
| velocities_root_xyz_no_heading = np.zeros((velocities_root_xy_no_heading.shape[0], 3)) |
| velocities_root_xyz_no_heading[:, 0] = velocities_root_xy_no_heading[:, 0] |
| velocities_root_xyz_no_heading[:, 2] = velocities_root_xy_no_heading[:, 1] |
| velocities_root_xyz_no_heading[1:, :] = np.matmul(inv_global_heading_rot[:-1], velocities_root_xyz_no_heading[1:, :,None]).squeeze(-1) |
| root_translation = np.cumsum(velocities_root_xyz_no_heading, axis=0) |
| root_translation[:, 1] = height |
|
|
| nfrm, njoint, _, _ = rotations_matrix.shape |
| axis_angle = matrix_to_axis_angle(torch.from_numpy(rotations_matrix)).numpy().reshape(nfrm, -1) |
| global_orient = axis_angle[:, :3] |
|
|
| |
| |
| |
| |
| |
| |
| |
| return dict( |
| global_orient = global_orient, |
| body_pose = axis_angle[:, 3:], |
| transl = root_translation, |
| betas = torch.zeros((nfrm, 10)) |
| ) |
|
|
