| 6 | |
| 7 | |
| 8 | class OptimizablePose(nn.Module): |
| 9 | def __init__(self, init_pose): |
| 10 | super().__init__() |
| 11 | assert (isinstance(init_pose, torch.FloatTensor)) |
| 12 | self.register_parameter('data', nn.Parameter(init_pose)) |
| 13 | self.data.required_grad_ = True |
| 14 | |
| 15 | def copy_from(self, pose): |
| 16 | self.data = deepcopy(pose.data) |
| 17 | |
| 18 | def matrix(self): |
| 19 | Rt = torch.eye(4) |
| 20 | Rt[:3, :3] = self.rotation() |
| 21 | Rt[:3, 3] = self.translation() |
| 22 | return Rt |
| 23 | |
| 24 | def rotation(self): |
| 25 | w = self.data[3:] |
| 26 | wx = self.skew_symmetric(w) |
| 27 | theta = w.norm(dim=-1)[..., None, None] |
| 28 | I = torch.eye(3, device=w.device, dtype=torch.float32) |
| 29 | A = self.taylor_A(theta) |
| 30 | B = self.taylor_B(theta) |
| 31 | R = I + A * wx + B * wx @ wx |
| 32 | return R |
| 33 | |
| 34 | def translation(self, ): |
| 35 | return self.data[:3] |
| 36 | |
| 37 | @classmethod |
| 38 | def log(cls, R, eps=1e-7): # [...,3,3] |
| 39 | """ |
| 40 | compute so(3) from the SO(3) |
| 41 | Args: |
| 42 | R (tensor): SO(3),rotation matrix. |
| 43 | Returns: |
| 44 | w (tensor,R^{3}): so(3). |
| 45 | """ |
| 46 | trace = R[..., 0, 0] + R[..., 1, 1] + R[..., 2, 2] |
| 47 | # ln(R) will explode if theta==pi |
| 48 | theta = ((trace - 1) / 2).clamp(-1 + eps, 1 - |
| 49 | eps).acos_()[..., None, None] % np.pi |
| 50 | lnR = 1 / (2 * cls.taylor_A(theta) + 1e-8) * \ |
| 51 | (R - R.transpose(-2, -1)) # FIXME: wei-chiu finds it weird |
| 52 | w0, w1, w2 = lnR[..., 2, 1], lnR[..., 0, 2], lnR[..., 1, 0] |
| 53 | w = torch.stack([w0, w1, w2], dim=-1) |
| 54 | return w |
| 55 | |
| 56 | @classmethod |
| 57 | def from_matrix(cls, Rt, eps=1e-8): # [...,3,4] |
| 58 | """ |
| 59 | Conversion of Lie Groups to Lie Algebras(SE(3)->se(3)), and add grad using torch.nn.Module. |
| 60 | Args: |
| 61 | cls (class): self.class |
| 62 | Rt (tensor,4*4): SE(3)T, transformation matrix |
| 63 | Returns: |
| 64 | OptimizablePose (class): add 'data'(se(3),R^{6}) parameters incluing grad |
| 65 | """ |