import copy import os import joblib import numpy as np from scipy.spatial.transform import Slerp, Rotation import torch from hawor.utils.process import run_mano, run_mano_left from hawor.utils.rotation import angle_axis_to_quaternion, angle_axis_to_rotation_matrix, quaternion_to_rotation_matrix, rotation_matrix_to_angle_axis from lib.utils.geometry import rotmat_to_rot6d from lib.utils.geometry import rot6d_to_rotmat def slerp_interpolation_aa(pos, valid): B, T, N, _ = pos.shape # B: 批次大小, T: 时间步长, N: 关节数, 4: 四元数维度 pos_interp = pos.copy() # 创建副本以存储插值结果 for b in range(B): for n in range(N): quat_b_n = pos[b, :, n, :] valid_b_n = valid[b, :] invalid_idxs = np.where(~valid_b_n)[0] valid_idxs = np.where(valid_b_n)[0] if len(invalid_idxs) == 0: continue if len(valid_idxs) > 1: valid_times = valid_idxs # 有效时间步 valid_rots = Rotation.from_rotvec(quat_b_n[valid_idxs]) # 有效四元数 slerp = Slerp(valid_times, valid_rots) for idx in invalid_idxs: if idx < valid_idxs[0]: # 时间步小于第一个有效时间步，进行外推 pos_interp[b, idx, n, :] = quat_b_n[valid_idxs[0]] # 复制第一个有效四元数 elif idx > valid_idxs[-1]: # 时间步大于最后一个有效时间步，进行外推 pos_interp[b, idx, n, :] = quat_b_n[valid_idxs[-1]] # 复制最后一个有效四元数 else: interp_rot = slerp([idx]) pos_interp[b, idx, n, :] = interp_rot.as_rotvec()[0] # print("#######") # if N > 1: # print(pos[1,0,11]) # print(pos_interp[1,0,11]) return pos_interp def slerp_interpolation_quat(pos, valid): # wxyz to xyzw pos = pos[:, :, :, [1, 2, 3, 0]] B, T, N, _ = pos.shape # B: 批次大小, T: 时间步长, N: 关节数, 4: 四元数维度 pos_interp = pos.copy() # 创建副本以存储插值结果 for b in range(B): for n in range(N): quat_b_n = pos[b, :, n, :] valid_b_n = valid[b, :] invalid_idxs = np.where(~valid_b_n)[0] valid_idxs = np.where(valid_b_n)[0] if len(invalid_idxs) == 0: continue if len(valid_idxs) > 1: valid_times = valid_idxs # 有效时间步 valid_rots = Rotation.from_quat(quat_b_n[valid_idxs]) # 有效四元数 slerp = Slerp(valid_times, valid_rots) for idx in invalid_idxs: if idx < valid_idxs[0]: # 时间步小于第一个有效时间步，进行外推 pos_interp[b, idx, n, :] = quat_b_n[valid_idxs[0]] # 复制第一个有效四元数 elif idx > valid_idxs[-1]: # 时间步大于最后一个有效时间步，进行外推 pos_interp[b, idx, n, :] = quat_b_n[valid_idxs[-1]] # 复制最后一个有效四元数 else: interp_rot = slerp([idx]) pos_interp[b, idx, n, :] = interp_rot.as_quat()[0] # xyzw to wxyz pos_interp = pos_interp[:, :, :, [3, 0, 1, 2]] return pos_interp def linear_interpolation_nd(pos, valid): B, T = pos.shape[:2] # 取出批次大小B和时间步长T feature_dim = pos.shape[2] # ** 代表的任意维度 pos_interp = pos.copy() # 创建一个副本，用来保存插值结果 for b in range(B): for idx in range(feature_dim): # 针对任意维度 pos_b_idx = pos[b, :, idx] # 取出第b批次对应的**维度下的一个时间序列 valid_b = valid[b, :] # 当前批次的有效标志 # 找到无效的索引（False） invalid_idxs = np.where(~valid_b)[0] valid_idxs = np.where(valid_b)[0] if len(invalid_idxs) == 0: continue # 对无效部分进行线性插值 if len(valid_idxs) > 1: # 确保有足够的有效点用于插值 pos_b_idx[invalid_idxs] = np.interp(invalid_idxs, valid_idxs, pos_b_idx[valid_idxs]) pos_interp[b, :, idx] = pos_b_idx # 保存插值结果 return pos_interp def world2canonical_convert(R_c2w_sla, t_c2w_sla, data_out, handedness): init_rot_mat = copy.deepcopy(data_out["init_root_orient"]) init_rot_mat = torch.einsum("tij,btjk->btik", R_c2w_sla, init_rot_mat) init_rot = rotation_matrix_to_angle_axis(init_rot_mat) init_rot_quat = angle_axis_to_quaternion(init_rot) # data_out["init_root_orient"] = rotation_matrix_to_angle_axis(data_out["init_root_orient"]) # data_out["init_hand_pose"] = rotation_matrix_to_angle_axis(data_out["init_hand_pose"]) data_out_init_root_orient = rotation_matrix_to_angle_axis(data_out["init_root_orient"]) data_out_init_hand_pose = rotation_matrix_to_angle_axis(data_out["init_hand_pose"]) init_trans = data_out["init_trans"] # (B, T, 3) if handedness == "left": outputs = run_mano_left(data_out["init_trans"], data_out_init_root_orient, data_out_init_hand_pose, betas=data_out["init_betas"]) elif handedness == "right": outputs = run_mano(data_out["init_trans"], data_out_init_root_orient, data_out_init_hand_pose, betas=data_out["init_betas"]) root_loc = outputs["joints"][..., 0, :].cpu() # (B, T, 3) offset = init_trans - root_loc # It is a constant, no matter what the rotation is. init_trans = ( torch.einsum("tij,btj->bti", R_c2w_sla, root_loc) + t_c2w_sla[None, :] + offset ) data_world = { "init_root_orient": init_rot, # (B, T, 3) "init_hand_pose": data_out_init_hand_pose, # (B, T, 15, 3) "init_trans": init_trans, # (B, T, 3) "init_betas": data_out["init_betas"] # (B, T, 10) } return data_world def filling_preprocess(item): num_joints = 15 global_trans = item['trans'] # (2, seq_len, 3) global_rot = item['rot'] #(2, seq_len, 3) hand_pose = item['hand_pose'] # (2, seq_len, 45) betas = item['betas'] # (2, seq_len, 10) valid = item['valid'] # (2, seq_len) N, T, _ = global_trans.shape R_canonical2world_left_aa = torch.from_numpy(global_rot[0, 0]) R_canonical2world_right_aa = torch.from_numpy(global_rot[1, 0]) R_world2canonical_left = angle_axis_to_rotation_matrix(R_canonical2world_left_aa).t() R_world2canonical_right = angle_axis_to_rotation_matrix(R_canonical2world_right_aa).t() # transform left hand to canonical hand_pose = hand_pose.reshape(N, T, num_joints, 3) data_world_left = { "init_trans": torch.from_numpy(global_trans[0:1]), "init_root_orient": angle_axis_to_rotation_matrix(torch.from_numpy(global_rot[0:1])), "init_hand_pose": angle_axis_to_rotation_matrix(torch.from_numpy(hand_pose[0:1])), "init_betas": torch.from_numpy(betas[0:1]), } data_left_init_root_orient = rotation_matrix_to_angle_axis(data_world_left["init_root_orient"]) data_left_init_hand_pose = rotation_matrix_to_angle_axis(data_world_left["init_hand_pose"]) outputs = run_mano_left(data_world_left["init_trans"], data_left_init_root_orient, data_left_init_hand_pose, betas=data_world_left["init_betas"]) init_trans = data_world_left["init_trans"][0, 0] # (3,) root_loc = outputs["joints"][0, 0, 0, :].cpu() # (3,) offset = init_trans - root_loc # It is a constant, no matter what the rotation is. t_world2canonical_left = -torch.einsum("ij,j->i", R_world2canonical_left, root_loc) - offset R_world2canonical_left = R_world2canonical_left.repeat(T, 1, 1) t_world2canonical_left = t_world2canonical_left.repeat(T, 1) data_canonical_left = world2canonical_convert(R_world2canonical_left, t_world2canonical_left, data_world_left, "left") # transform right hand to canonical data_world_right = { "init_trans": torch.from_numpy(global_trans[1:2]), "init_root_orient": angle_axis_to_rotation_matrix(torch.from_numpy(global_rot[1:2])), "init_hand_pose": angle_axis_to_rotation_matrix(torch.from_numpy(hand_pose[1:2])), "init_betas": torch.from_numpy(betas[1:2]), } data_right_init_root_orient = rotation_matrix_to_angle_axis(data_world_right["init_root_orient"]) data_right_init_hand_pose = rotation_matrix_to_angle_axis(data_world_right["init_hand_pose"]) outputs = run_mano(data_world_right["init_trans"], data_right_init_root_orient, data_right_init_hand_pose, betas=data_world_right["init_betas"]) init_trans = data_world_right["init_trans"][0, 0] # (3,) root_loc = outputs["joints"][0, 0, 0, :].cpu() # (3,) offset = init_trans - root_loc # It is a constant, no matter what the rotation is. t_world2canonical_right = -torch.einsum("ij,j->i", R_world2canonical_right, root_loc) - offset R_world2canonical_right = R_world2canonical_right.repeat(T, 1, 1) t_world2canonical_right = t_world2canonical_right.repeat(T, 1) data_canonical_right = world2canonical_convert(R_world2canonical_right, t_world2canonical_right, data_world_right, "right") # merge left and right canonical data global_rot = torch.cat((data_canonical_left['init_root_orient'], data_canonical_right['init_root_orient'])) global_trans = torch.cat((data_canonical_left['init_trans'], data_canonical_right['init_trans'])).numpy() # global_rot = angle_axis_to_quaternion(global_rot).numpy().reshape(N, T, 1, 4) global_rot = global_rot.reshape(N, T, 1, 3).numpy() hand_pose = hand_pose.reshape(N, T, 15, 3) # hand_pose = angle_axis_to_quaternion(torch.from_numpy(hand_pose)).numpy() # lerp and slerp global_trans_lerped = linear_interpolation_nd(global_trans, valid) betas_lerped = linear_interpolation_nd(betas, valid) global_rot_slerped = slerp_interpolation_aa(global_rot, valid) hand_pose_slerped = slerp_interpolation_aa(hand_pose, valid) # convert to rot6d global_rot_slerped_mat = angle_axis_to_rotation_matrix(torch.from_numpy(global_rot_slerped.reshape(N*T, -1))) # global_rot_slerped_mat = quaternion_to_rotation_matrix(torch.from_numpy(global_rot_slerped.reshape(N*T, -1))) global_rot_slerped_rot6d = rotmat_to_rot6d(global_rot_slerped_mat).reshape(N, T, -1).numpy() hand_pose_slerped_mat = angle_axis_to_rotation_matrix(torch.from_numpy(hand_pose_slerped.reshape(N*T*num_joints, -1))) # hand_pose_slerped_mat = quaternion_to_rotation_matrix(torch.from_numpy(hand_pose_slerped.reshape(N*T*num_joints, -1))) hand_pose_slerped_rot6d = rotmat_to_rot6d(hand_pose_slerped_mat).reshape(N, T, -1).numpy() # concat to (T, concat_dim) global_pose_vec_input = np.concatenate((global_trans_lerped, betas_lerped, global_rot_slerped_rot6d, hand_pose_slerped_rot6d), axis=-1).transpose(1, 0, 2).reshape(T, -1) R_canon2w_left = R_world2canonical_left.transpose(-1, -2) t_canon2w_left = -torch.einsum("tij,tj->ti", R_canon2w_left, t_world2canonical_left) R_canon2w_right = R_world2canonical_right.transpose(-1, -2) t_canon2w_right = -torch.einsum("tij,tj->ti", R_canon2w_right, t_world2canonical_right) transform_w_canon = { "R_w2canon_left": R_world2canonical_left, "t_w2canon_left": t_world2canonical_left, "R_canon2w_left": R_canon2w_left, "t_canon2w_left": t_canon2w_left, "R_w2canon_right": R_world2canonical_right, "t_w2canon_right": t_world2canonical_right, "R_canon2w_right": R_canon2w_right, "t_canon2w_right": t_canon2w_right, } return global_pose_vec_input, transform_w_canon def custom_rot6d_to_rotmat(rot6d): original_shape = rot6d.shape[:-1] rot6d = rot6d.reshape(-1, 6) mat = rot6d_to_rotmat(rot6d) mat = mat.reshape(*original_shape, 3, 3) return mat def filling_postprocess(output, transform_w_canon): # output = output.numpy() output = output.permute(1, 0, 2) # (2, T, -1) N, T, _ = output.shape canon_trans = output[:, :, :3] betas = output[:, :, 3:13] canon_rot_rot6d = output[:, :, 13:19] hand_pose_rot6d = output[:, :, 19:109].reshape(N, T, 15, 6) canon_rot_mat = custom_rot6d_to_rotmat(canon_rot_rot6d) hand_pose_mat = custom_rot6d_to_rotmat(hand_pose_rot6d) data_canonical_left = { "init_trans": canon_trans[[0], :, :], "init_root_orient": canon_rot_mat[[0], :, :, :], "init_hand_pose": hand_pose_mat[[0], :, :, :, :], "init_betas": betas[[0], :, :] } data_canonical_right = { "init_trans": canon_trans[[1], :, :], "init_root_orient": canon_rot_mat[[1], :, :, :], "init_hand_pose": hand_pose_mat[[1], :, :, :, :], "init_betas": betas[[1], :, :] } R_canon2w_left = transform_w_canon['R_canon2w_left'] t_canon2w_left = transform_w_canon['t_canon2w_left'] R_canon2w_right = transform_w_canon['R_canon2w_right'] t_canon2w_right = transform_w_canon['t_canon2w_right'] world_left = world2canonical_convert(R_canon2w_left, t_canon2w_left, data_canonical_left, "left") world_right = world2canonical_convert(R_canon2w_right, t_canon2w_right, data_canonical_right, "right") global_rot = torch.cat((world_left['init_root_orient'], world_right['init_root_orient'])).numpy() global_trans = torch.cat((world_left['init_trans'], world_right['init_trans'])).numpy() pred_data = { "trans": global_trans, # (2, T, 3) "rot": global_rot, # (2, T, 3) "hand_pose": rotation_matrix_to_angle_axis(hand_pose_mat).flatten(-2).numpy(), # (2, T, 45) "betas": betas.numpy(), # (2, T, 10) } return pred_data