Applies a batch of rigid transformations to the joints. Parameters ---------- rot_mats : torch.tensor BxNx3x3 Tensor of rotation matrices joints : torch.tensor BxNx3 Locations of joints parents : torch.tensor BxN The kinematic tree of each object dtyp
(rot_mats: Tensor,
joints: Tensor,
parents: Tensor,
dtype=torch.float32)
| 334 | |
| 335 | |
| 336 | def batch_rigid_transform(rot_mats: Tensor, |
| 337 | joints: Tensor, |
| 338 | parents: Tensor, |
| 339 | dtype=torch.float32) -> Tensor: |
| 340 | """Applies a batch of rigid transformations to the joints. |
| 341 | |
| 342 | Parameters |
| 343 | ---------- |
| 344 | rot_mats : torch.tensor BxNx3x3 |
| 345 | Tensor of rotation matrices |
| 346 | joints : torch.tensor BxNx3 |
| 347 | Locations of joints |
| 348 | parents : torch.tensor BxN |
| 349 | The kinematic tree of each object |
| 350 | dtype : torch.dtype, optional: |
| 351 | The data type of the created tensors, the default is torch.float32 |
| 352 | |
| 353 | Returns |
| 354 | ------- |
| 355 | posed_joints : torch.tensor BxNx3 |
| 356 | The locations of the joints after applying the pose rotations |
| 357 | rel_transforms : torch.tensor BxNx4x4 |
| 358 | The relative (with respect to the root joint) rigid transformations |
| 359 | for all the joints |
| 360 | """ |
| 361 | |
| 362 | joints = torch.unsqueeze(joints, dim=-1) |
| 363 | |
| 364 | rel_joints = joints.clone() |
| 365 | rel_joints[:, 1:] -= joints[:, parents[1:]] |
| 366 | |
| 367 | transforms_mat = transform_mat(rot_mats.reshape(-1, 3, 3), |
| 368 | rel_joints.reshape(-1, 3, 1)).reshape( |
| 369 | -1, joints.shape[1], 4, 4) |
| 370 | |
| 371 | transform_chain = [transforms_mat[:, 0]] |
| 372 | for i in range(1, parents.shape[0]): |
| 373 | # Subtract the joint location at the rest pose |
| 374 | # No need for rotation, since it's identity when at rest |
| 375 | curr_res = torch.matmul(transform_chain[parents[i]], transforms_mat[:, |
| 376 | i]) |
| 377 | transform_chain.append(curr_res) |
| 378 | |
| 379 | transforms = torch.stack(transform_chain, dim=1) |
| 380 | |
| 381 | # The last column of the transformations contains the posed joints |
| 382 | posed_joints = transforms[:, :, :3, 3] |
| 383 | |
| 384 | # The last column of the transformations contains the posed joints |
| 385 | posed_joints = transforms[:, :, :3, 3] |
| 386 | |
| 387 | joints_homogen = F.pad(joints, [0, 0, 0, 1]) |
| 388 | |
| 389 | rel_transforms = transforms - F.pad( |
| 390 | torch.matmul(transforms, joints_homogen), [3, 0, 0, 0, 0, 0, 0, 0]) |
| 391 | |
| 392 | return posed_joints, rel_transforms |
no test coverage detected