(
self,
context: dict,
global_step: int = 0,
visualization_dump: Optional[dict] = None,
)
| 124 | return head(decout, img_shape, ray_embedding=ray_embedding) |
| 125 | |
| 126 | def forward( |
| 127 | self, |
| 128 | context: dict, |
| 129 | global_step: int = 0, |
| 130 | visualization_dump: Optional[dict] = None, |
| 131 | ) -> Gaussians: |
| 132 | device = context["image"].device |
| 133 | b, v, _, h, w = context["image"].shape |
| 134 | |
| 135 | # Encode the context images. |
| 136 | dec_feat, shape, images = self.backbone(context) |
| 137 | with torch.cuda.amp.autocast(enabled=False): |
| 138 | all_mean_res = [] |
| 139 | all_other_params = [] |
| 140 | res1 = self._downstream_head(1, [tok[:, 0].float() for tok in dec_feat], shape[:, 0]) |
| 141 | all_mean_res.append(res1) |
| 142 | for i in range(1, v): |
| 143 | res2 = self._downstream_head(2, [tok[:, i].float() for tok in dec_feat], shape[:, i]) |
| 144 | all_mean_res.append(res2) |
| 145 | |
| 146 | # for the 3DGS heads |
| 147 | if self.gs_params_head_type == 'dpt_gs': |
| 148 | GS_res1 = self.gaussian_param_head([tok[:, 0].float() for tok in dec_feat], all_mean_res[0]['pts3d'].permute(0, 3, 1, 2), images[:, 0, :3], shape[0, 0].cpu().tolist()) |
| 149 | GS_res1 = rearrange(GS_res1, "b d h w -> b (h w) d") |
| 150 | all_other_params.append(GS_res1) |
| 151 | for i in range(1, v): |
| 152 | GS_res2 = self.gaussian_param_head2([tok[:, i].float() for tok in dec_feat], all_mean_res[i]['pts3d'].permute(0, 3, 1, 2), images[:, i, :3], shape[0, i].cpu().tolist()) |
| 153 | GS_res2 = rearrange(GS_res2, "b d h w -> b (h w) d") |
| 154 | all_other_params.append(GS_res2) |
| 155 | else: |
| 156 | raise NotImplementedError(f"unexpected {self.gs_params_head_type=}") |
| 157 | |
| 158 | pts_all = [all_mean_res_i['pts3d'] for all_mean_res_i in all_mean_res] |
| 159 | pts_all = torch.stack(pts_all, dim=1) |
| 160 | pts_all = rearrange(pts_all, "b v h w xyz -> b v (h w) xyz") |
| 161 | pts_all = pts_all.unsqueeze(-2) # for cfg.num_surfaces |
| 162 | |
| 163 | depths = pts_all[..., -1].unsqueeze(-1) |
| 164 | |
| 165 | gaussians = torch.stack(all_other_params, dim=1) |
| 166 | gaussians = rearrange(gaussians, "... (srf c) -> ... srf c", srf=self.cfg.num_surfaces) |
| 167 | densities = gaussians[..., 0].sigmoid().unsqueeze(-1) |
| 168 | |
| 169 | # Convert the features and depths into Gaussians. |
| 170 | if self.pose_free: |
| 171 | gaussians = self.gaussian_adapter.forward( |
| 172 | pts_all.unsqueeze(-2), |
| 173 | depths, |
| 174 | self.map_pdf_to_opacity(densities, global_step), |
| 175 | rearrange(gaussians[..., 1:], "b v r srf c -> b v r srf () c"), |
| 176 | ) |
| 177 | else: |
| 178 | xy_ray, _ = sample_image_grid((h, w), device) |
| 179 | xy_ray = rearrange(xy_ray, "h w xy -> (h w) () xy") |
| 180 | xy_ray = xy_ray[None, None, ...].expand(b, v, -1, -1, -1) |
| 181 | |
| 182 | gaussians = self.gaussian_adapter.forward( |
| 183 | rearrange(context["extrinsics"], "b v i j -> b v () () () i j"), |
nothing calls this directly
no test coverage detected