Appearance optimization module.
| 51 | |
| 52 | |
| 53 | class AppearanceOptModule(torch.nn.Module): |
| 54 | """Appearance optimization module.""" |
| 55 | |
| 56 | def __init__( |
| 57 | self, |
| 58 | n: int, |
| 59 | feature_dim: int, |
| 60 | embed_dim: int = 16, |
| 61 | sh_degree: int = 3, |
| 62 | mlp_width: int = 64, |
| 63 | mlp_depth: int = 2, |
| 64 | ): |
| 65 | super().__init__() |
| 66 | self.embed_dim = embed_dim |
| 67 | self.sh_degree = sh_degree |
| 68 | self.embeds = torch.nn.Embedding(n, embed_dim) |
| 69 | layers = [] |
| 70 | layers.append( |
| 71 | torch.nn.Linear(embed_dim + feature_dim + (sh_degree + 1) ** 2, mlp_width) |
| 72 | ) |
| 73 | layers.append(torch.nn.ReLU(inplace=True)) |
| 74 | for _ in range(mlp_depth - 1): |
| 75 | layers.append(torch.nn.Linear(mlp_width, mlp_width)) |
| 76 | layers.append(torch.nn.ReLU(inplace=True)) |
| 77 | layers.append(torch.nn.Linear(mlp_width, 3)) |
| 78 | self.color_head = torch.nn.Sequential(*layers) |
| 79 | |
| 80 | def forward( |
| 81 | self, features: Tensor, embed_ids: Tensor, dirs: Tensor, sh_degree: int |
| 82 | ) -> Tensor: |
| 83 | """Adjust appearance based on embeddings. |
| 84 | |
| 85 | Args: |
| 86 | features: (N, feature_dim) |
| 87 | embed_ids: (C,) |
| 88 | dirs: (C, N, 3) |
| 89 | |
| 90 | Returns: |
| 91 | colors: (C, N, 3) |
| 92 | """ |
| 93 | from gsplat.cuda._torch_impl import _eval_sh_bases_fast |
| 94 | |
| 95 | C, N = dirs.shape[:2] |
| 96 | # Camera embeddings |
| 97 | if embed_ids is None: |
| 98 | embeds = torch.zeros(C, self.embed_dim, device=features.device) |
| 99 | else: |
| 100 | embeds = self.embeds(embed_ids) # [C, D2] |
| 101 | embeds = embeds[:, None, :].expand(-1, N, -1) # [C, N, D2] |
| 102 | # GS features |
| 103 | features = features[None, :, :].expand(C, -1, -1) # [C, N, D1] |
| 104 | # View directions |
| 105 | dirs = F.normalize(dirs, dim=-1) # [C, N, 3] |
| 106 | num_bases_to_use = (sh_degree + 1) ** 2 |
| 107 | num_bases = (self.sh_degree + 1) ** 2 |
| 108 | sh_bases = torch.zeros(C, N, num_bases, device=features.device) # [C, N, K] |
| 109 | sh_bases[:, :, :num_bases_to_use] = _eval_sh_bases_fast(num_bases_to_use, dirs) |
| 110 | # Get colors |