(self, cont6d_params, root_pos, skel_joints=None, do_root_R=True)
| 171 | return joints |
| 172 | |
| 173 | def forward_kinematics_cont6d(self, cont6d_params, root_pos, skel_joints=None, do_root_R=True): |
| 174 | # cont6d_params (batch_size, joints_num, 6) |
| 175 | # joints (batch_size, joints_num, 3) |
| 176 | # root_pos (batch_size, 3) |
| 177 | if skel_joints is not None: |
| 178 | # skel_joints = torch.from_numpy(skel_joints) |
| 179 | offsets = self.get_offsets_joints_batch(skel_joints) |
| 180 | if len(self._offset.shape) == 2: |
| 181 | offsets = self._offset.expand(cont6d_params.shape[0], -1, -1) |
| 182 | joints = torch.zeros(cont6d_params.shape[:-1] + (3,)).to(cont6d_params.device) |
| 183 | joints[..., 0, :] = root_pos |
| 184 | for chain in self._kinematic_tree: |
| 185 | if do_root_R: |
| 186 | matR = cont6d_to_matrix(cont6d_params[:, 0]) |
| 187 | else: |
| 188 | matR = torch.eye(3).expand((len(cont6d_params), -1, -1)).detach().to(cont6d_params.device) |
| 189 | for i in range(1, len(chain)): |
| 190 | matR = torch.matmul(matR, cont6d_to_matrix(cont6d_params[:, chain[i]])) |
| 191 | offset_vec = offsets[:, chain[i]].unsqueeze(-1) |
| 192 | # print(matR.shape, offset_vec.shape) |
| 193 | joints[:, chain[i]] = torch.matmul(matR, offset_vec).squeeze(-1) + joints[:, chain[i-1]] |
| 194 | return joints |
| 195 | |
| 196 | |
| 197 |
no test coverage detected