Convert quaternion coefficients to rotation matrix. Args: quat: size = [B, 4] 4 <===>(w, x, y, z) Returns: Rotation matrix corresponding to the quaternion -- size = [B, 3, 3]
(quat)
| 22 | |
| 23 | |
| 24 | def quat_to_rotmat(quat): |
| 25 | """Convert quaternion coefficients to rotation matrix. |
| 26 | |
| 27 | Args: |
| 28 | quat: size = [B, 4] 4 <===>(w, x, y, z) |
| 29 | Returns: |
| 30 | Rotation matrix corresponding to the quaternion -- size = [B, 3, 3] |
| 31 | """ |
| 32 | norm_quat = quat |
| 33 | norm_quat = norm_quat / norm_quat.norm(p=2, dim=1, keepdim=True) |
| 34 | w = norm_quat[:, 0] |
| 35 | x = norm_quat[:, 1] |
| 36 | y = norm_quat[:, 2] |
| 37 | z = norm_quat[:, 3] |
| 38 | B = quat.size(0) |
| 39 | |
| 40 | w2, x2, y2, z2 = w.pow(2), x.pow(2), y.pow(2), z.pow(2) |
| 41 | wx, wy, wz = w * x, w * y, w * z |
| 42 | xy, xz, yz = x * y, x * z, y * z |
| 43 | |
| 44 | rotMat = torch.stack([ |
| 45 | w2 + x2 - y2 - z2, 2 * xy - 2 * wz, 2 * wy + 2 * xz, 2 * wz + 2 * xy, |
| 46 | w2 - x2 + y2 - z2, 2 * yz - 2 * wx, 2 * xz - 2 * wy, 2 * wx + 2 * yz, |
| 47 | w2 - x2 - y2 + z2 |
| 48 | ], |
| 49 | dim=1).view(B, 3, 3) |
| 50 | return rotMat |
| 51 | |
| 52 | |
| 53 | def rot6d_to_rotmat(x): |