Convert 6d rotation representation to 3d vector of axis-angle rotation. Shape: - Input: :Torch:`(N, 6)` - Output: :Torch:`(N, 3)`
(rot6d)
| 35 | return rot6d |
| 36 | |
| 37 | def rot6d_to_axis_angle(rot6d): |
| 38 | """Convert 6d rotation representation to 3d vector of axis-angle rotation. |
| 39 | Shape: |
| 40 | - Input: :Torch:`(N, 6)` |
| 41 | - Output: :Torch:`(N, 3)` |
| 42 | """ |
| 43 | batch_size = rot6d.shape[0] |
| 44 | |
| 45 | rot6d = rot6d.view(batch_size, 3, 2) |
| 46 | a1 = rot6d[:, :, 0] |
| 47 | a2 = rot6d[:, :, 1] |
| 48 | b1 = F.normalize(a1) |
| 49 | b2 = F.normalize(a2 - torch.einsum('bi,bi->b', b1, a2).unsqueeze(-1) * b1) |
| 50 | b3 = torch.cross(b1, b2, dim=-1) |
| 51 | rot_mat = torch.stack((b1, b2, b3), dim=-1) # 3x3 rotation matrix |
| 52 | |
| 53 | rot_mat = torch.cat([rot_mat, torch.zeros((batch_size, 3, 1), device=rot_mat.device).float()], |
| 54 | 2) # 3x4 rotation matrix |
| 55 | axis_angle = tgm.rotation_matrix_to_angle_axis(rot_mat).reshape(-1, 3) # axis-angle |
| 56 | axis_angle[torch.isnan(axis_angle)] = 0.0 |
| 57 | return axis_angle |
| 58 | |
| 59 | def axis_angle_to_mat3x3(angle_axis): |
| 60 | """ |
no outgoing calls
no test coverage detected