Adjust appearance based on embeddings. Args: features: (N, feature_dim) embed_ids: (C,) dirs: (C, N, 3) Returns: colors: (C, N, 3)
(
self, features: Tensor, embed_ids: Tensor, dirs: Tensor, sh_degree: int
)
| 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 |
| 111 | if self.embed_dim > 0: |
| 112 | h = torch.cat([embeds, features, sh_bases], dim=-1) # [C, N, D1 + D2 + K] |
| 113 | else: |
| 114 | h = torch.cat([features, sh_bases], dim=-1) |
| 115 | colors = self.color_head(h) |
| 116 | return colors |
| 117 | |
| 118 | |
| 119 | def rotation_6d_to_matrix(d6: Tensor) -> Tensor: |
nothing calls this directly
no test coverage detected