Camera pose optimization module.
| 12 | |
| 13 | |
| 14 | class CameraOptModule(torch.nn.Module): |
| 15 | """Camera pose optimization module.""" |
| 16 | |
| 17 | def __init__(self, n: int): |
| 18 | super().__init__() |
| 19 | # Delta positions (3D) + Delta rotations (6D) |
| 20 | self.embeds = torch.nn.Embedding(n, 9) |
| 21 | # Identity rotation in 6D representation |
| 22 | self.register_buffer("identity", torch.tensor([1.0, 0.0, 0.0, 0.0, 1.0, 0.0])) |
| 23 | |
| 24 | def zero_init(self): |
| 25 | torch.nn.init.zeros_(self.embeds.weight) |
| 26 | |
| 27 | def random_init(self, std: float): |
| 28 | torch.nn.init.normal_(self.embeds.weight, std=std) |
| 29 | |
| 30 | def forward(self, camtoworlds: Tensor, embed_ids: Tensor) -> Tensor: |
| 31 | """Adjust camera pose based on deltas. |
| 32 | |
| 33 | Args: |
| 34 | camtoworlds: (..., 4, 4) |
| 35 | embed_ids: (...,) |
| 36 | |
| 37 | Returns: |
| 38 | updated camtoworlds: (..., 4, 4) |
| 39 | """ |
| 40 | assert camtoworlds.shape[:-2] == embed_ids.shape |
| 41 | batch_shape = camtoworlds.shape[:-2] |
| 42 | pose_deltas = self.embeds(embed_ids) # (..., 9) |
| 43 | dx, drot = pose_deltas[..., :3], pose_deltas[..., 3:] |
| 44 | rot = rotation_6d_to_matrix( |
| 45 | drot + self.identity.expand(*batch_shape, -1) |
| 46 | ) # (..., 3, 3) |
| 47 | transform = torch.eye(4, device=pose_deltas.device).repeat((*batch_shape, 1, 1)) |
| 48 | transform[..., :3, :3] = rot |
| 49 | transform[..., :3, 3] = dx |
| 50 | return torch.matmul(camtoworlds, transform) |
| 51 | |
| 52 | |
| 53 | class AppearanceOptModule(torch.nn.Module): |