Spaces:
Running
Running
| from common.quaternion import * | |
| import scipy.ndimage.filters as filters | |
| class Skeleton(object): | |
| def __init__(self, offset, kinematic_tree, device): | |
| self.device = device | |
| self._raw_offset_np = offset.numpy() | |
| self._raw_offset = offset.clone().detach().to(device).float() | |
| self._kinematic_tree = kinematic_tree | |
| self._offset = None | |
| self._parents = [0] * len(self._raw_offset) | |
| self._parents[0] = -1 | |
| for chain in self._kinematic_tree: | |
| for j in range(1, len(chain)): | |
| self._parents[chain[j]] = chain[j-1] | |
| def njoints(self): | |
| return len(self._raw_offset) | |
| def offset(self): | |
| return self._offset | |
| def set_offset(self, offsets): | |
| self._offset = offsets.clone().detach().to(self.device).float() | |
| def kinematic_tree(self): | |
| return self._kinematic_tree | |
| def parents(self): | |
| return self._parents | |
| # joints (batch_size, joints_num, 3) | |
| def get_offsets_joints_batch(self, joints): | |
| assert len(joints.shape) == 3 | |
| _offsets = self._raw_offset.expand(joints.shape[0], -1, -1).clone() | |
| for i in range(1, self._raw_offset.shape[0]): | |
| _offsets[:, i] = torch.norm(joints[:, i] - joints[:, self._parents[i]], p=2, dim=1)[:, None] * _offsets[:, i] | |
| self._offset = _offsets.detach() | |
| return _offsets | |
| # joints (joints_num, 3) | |
| def get_offsets_joints(self, joints): | |
| assert len(joints.shape) == 2 | |
| _offsets = self._raw_offset.clone() | |
| for i in range(1, self._raw_offset.shape[0]): | |
| # print(joints.shape) | |
| _offsets[i] = torch.norm(joints[i] - joints[self._parents[i]], p=2, dim=0) * _offsets[i] | |
| self._offset = _offsets.detach() | |
| return _offsets | |
| # face_joint_idx should follow the order of right hip, left hip, right shoulder, left shoulder | |
| # joints (batch_size, joints_num, 3) | |
| def inverse_kinematics_np(self, joints, face_joint_idx, smooth_forward=False): | |
| assert len(face_joint_idx) == 4 | |
| '''Get Forward Direction''' | |
| l_hip, r_hip, sdr_r, sdr_l = face_joint_idx | |
| across1 = joints[:, r_hip] - joints[:, l_hip] | |
| across2 = joints[:, sdr_r] - joints[:, sdr_l] | |
| across = across1 + across2 | |
| across = across / np.sqrt((across**2).sum(axis=-1))[:, np.newaxis] | |
| # print(across1.shape, across2.shape) | |
| # forward (batch_size, 3) | |
| forward = np.cross(np.array([[0, 1, 0]]), across, axis=-1) | |
| if smooth_forward: | |
| forward = filters.gaussian_filter1d(forward, 20, axis=0, mode='nearest') | |
| # forward (batch_size, 3) | |
| forward = forward / np.sqrt((forward**2).sum(axis=-1))[..., np.newaxis] | |
| '''Get Root Rotation''' | |
| target = np.array([[0,0,1]]).repeat(len(forward), axis=0) | |
| root_quat = qbetween_np(forward, target) | |
| '''Inverse Kinematics''' | |
| # quat_params (batch_size, joints_num, 4) | |
| # print(joints.shape[:-1]) | |
| quat_params = np.zeros(joints.shape[:-1] + (4,)) | |
| # print(quat_params.shape) | |
| root_quat[0] = np.array([[1.0, 0.0, 0.0, 0.0]]) | |
| quat_params[:, 0] = root_quat | |
| # quat_params[0, 0] = np.array([[1.0, 0.0, 0.0, 0.0]]) | |
| for chain in self._kinematic_tree: | |
| R = root_quat | |
| for j in range(len(chain) - 1): | |
| # (batch, 3) | |
| u = self._raw_offset_np[chain[j+1]][np.newaxis,...].repeat(len(joints), axis=0) | |
| # print(u.shape) | |
| # (batch, 3) | |
| v = joints[:, chain[j+1]] - joints[:, chain[j]] | |
| v = v / np.sqrt((v**2).sum(axis=-1))[:, np.newaxis] | |
| # print(u.shape, v.shape) | |
| rot_u_v = qbetween_np(u, v) | |
| R_loc = qmul_np(qinv_np(R), rot_u_v) | |
| quat_params[:,chain[j + 1], :] = R_loc | |
| R = qmul_np(R, R_loc) | |
| return quat_params | |
| # Be sure root joint is at the beginning of kinematic chains | |
| def forward_kinematics(self, quat_params, root_pos, skel_joints=None, do_root_R=True): | |
| # quat_params (batch_size, joints_num, 4) | |
| # joints (batch_size, joints_num, 3) | |
| # root_pos (batch_size, 3) | |
| if skel_joints is not None: | |
| offsets = self.get_offsets_joints_batch(skel_joints) | |
| if len(self._offset.shape) == 2: | |
| offsets = self._offset.expand(quat_params.shape[0], -1, -1) | |
| joints = torch.zeros(quat_params.shape[:-1] + (3,)).to(self.device) | |
| joints[:, 0] = root_pos | |
| for chain in self._kinematic_tree: | |
| if do_root_R: | |
| R = quat_params[:, 0] | |
| else: | |
| R = torch.tensor([[1.0, 0.0, 0.0, 0.0]]).expand(len(quat_params), -1).detach().to(self.device) | |
| for i in range(1, len(chain)): | |
| R = qmul(R, quat_params[:, chain[i]]) | |
| offset_vec = offsets[:, chain[i]] | |
| joints[:, chain[i]] = qrot(R, offset_vec) + joints[:, chain[i-1]] | |
| return joints | |
| # Be sure root joint is at the beginning of kinematic chains | |
| def forward_kinematics_np(self, quat_params, root_pos, skel_joints=None, do_root_R=True): | |
| # quat_params (batch_size, joints_num, 4) | |
| # joints (batch_size, joints_num, 3) | |
| # root_pos (batch_size, 3) | |
| if skel_joints is not None: | |
| skel_joints = torch.from_numpy(skel_joints) | |
| offsets = self.get_offsets_joints_batch(skel_joints) | |
| if len(self._offset.shape) == 2: | |
| offsets = self._offset.expand(quat_params.shape[0], -1, -1) | |
| offsets = offsets.numpy() | |
| joints = np.zeros(quat_params.shape[:-1] + (3,)) | |
| joints[:, 0] = root_pos | |
| for chain in self._kinematic_tree: | |
| if do_root_R: | |
| R = quat_params[:, 0] | |
| else: | |
| R = np.array([[1.0, 0.0, 0.0, 0.0]]).repeat(len(quat_params), axis=0) | |
| for i in range(1, len(chain)): | |
| R = qmul_np(R, quat_params[:, chain[i]]) | |
| offset_vec = offsets[:, chain[i]] | |
| joints[:, chain[i]] = qrot_np(R, offset_vec) + joints[:, chain[i - 1]] | |
| return joints | |
| def forward_kinematics_cont6d_np(self, cont6d_params, root_pos, skel_joints=None, do_root_R=True): | |
| # cont6d_params (batch_size, joints_num, 6) | |
| # joints (batch_size, joints_num, 3) | |
| # root_pos (batch_size, 3) | |
| if skel_joints is not None: | |
| skel_joints = torch.from_numpy(skel_joints) | |
| offsets = self.get_offsets_joints_batch(skel_joints) | |
| if len(self._offset.shape) == 2: | |
| offsets = self._offset.expand(cont6d_params.shape[0], -1, -1) | |
| offsets = offsets.numpy() | |
| joints = np.zeros(cont6d_params.shape[:-1] + (3,)) | |
| joints[:, 0] = root_pos | |
| for chain in self._kinematic_tree: | |
| if do_root_R: | |
| matR = cont6d_to_matrix_np(cont6d_params[:, 0]) | |
| else: | |
| matR = np.eye(3)[np.newaxis, :].repeat(len(cont6d_params), axis=0) | |
| for i in range(1, len(chain)): | |
| matR = np.matmul(matR, cont6d_to_matrix_np(cont6d_params[:, chain[i]])) | |
| offset_vec = offsets[:, chain[i]][..., np.newaxis] | |
| # print(matR.shape, offset_vec.shape) | |
| joints[:, chain[i]] = np.matmul(matR, offset_vec).squeeze(-1) + joints[:, chain[i-1]] | |
| return joints | |
| def forward_kinematics_cont6d(self, cont6d_params, root_pos, skel_joints=None, do_root_R=True): | |
| # cont6d_params (batch_size, joints_num, 6) | |
| # joints (batch_size, joints_num, 3) | |
| # root_pos (batch_size, 3) | |
| if skel_joints is not None: | |
| # skel_joints = torch.from_numpy(skel_joints) | |
| offsets = self.get_offsets_joints_batch(skel_joints) | |
| if len(self._offset.shape) == 2: | |
| offsets = self._offset.expand(cont6d_params.shape[0], -1, -1) | |
| joints = torch.zeros(cont6d_params.shape[:-1] + (3,)).to(cont6d_params.device) | |
| joints[..., 0, :] = root_pos | |
| for chain in self._kinematic_tree: | |
| if do_root_R: | |
| matR = cont6d_to_matrix(cont6d_params[:, 0]) | |
| else: | |
| matR = torch.eye(3).expand((len(cont6d_params), -1, -1)).detach().to(cont6d_params.device) | |
| for i in range(1, len(chain)): | |
| matR = torch.matmul(matR, cont6d_to_matrix(cont6d_params[:, chain[i]])) | |
| offset_vec = offsets[:, chain[i]].unsqueeze(-1) | |
| # print(matR.shape, offset_vec.shape) | |
| joints[:, chain[i]] = torch.matmul(matR, offset_vec).squeeze(-1) + joints[:, chain[i-1]] | |
| return joints | |