Spaces:
Running on Zero
Running on Zero
| import torch | |
| from torch import nn | |
| class GaussianLatentDecoder(nn.Module): | |
| def __init__( | |
| self, | |
| latent_channels: int, | |
| hidden_channels: int, | |
| num_gaussian_parameters: int, | |
| init_sh_input_img: bool = True, | |
| ) -> None: | |
| super().__init__() | |
| in_channels = latent_channels + 1 # out_hat + depth | |
| self.decoder = nn.Sequential( | |
| nn.Conv2d(in_channels, hidden_channels, 3, 1, 1, padding_mode="replicate"), | |
| nn.GELU(), | |
| nn.Conv2d( | |
| hidden_channels, | |
| num_gaussian_parameters, | |
| 3, | |
| 1, | |
| 1, | |
| padding_mode="replicate", | |
| ), | |
| nn.GELU(), | |
| nn.Conv2d( | |
| num_gaussian_parameters, | |
| num_gaussian_parameters, | |
| 3, | |
| 1, | |
| 1, | |
| padding_mode="replicate", | |
| ), | |
| ) | |
| self._init_output_layer(init_sh_input_img) | |
| def _init_output_layer(self, init_sh_input_img: bool) -> None: | |
| # Keep the Gaussian head initially close to identity-like residuals. | |
| # Channel order: opacity, offset_xy, scales, quaternion, SH residuals. | |
| output_layer = self.decoder[-1] | |
| if init_sh_input_img: | |
| nn.init.zeros_(output_layer.weight[10:]) | |
| nn.init.zeros_(output_layer.bias[10:]) | |
| nn.init.zeros_(output_layer.weight[3:6]) | |
| nn.init.zeros_(output_layer.bias[3:6]) | |
| def forward( | |
| self, | |
| latent: torch.Tensor, | |
| depth: torch.Tensor, | |
| ) -> torch.Tensor: | |
| x = torch.cat((latent, depth), dim=1) | |
| return self.decoder(x) | |
| class GaussianLatentDepthDecoder(nn.Module): | |
| def __init__( | |
| self, | |
| latent_channels: int, | |
| hidden_channels: int, | |
| ) -> None: | |
| super().__init__() | |
| self.inv_depth_logits = nn.Sequential( | |
| nn.Conv2d(latent_channels, hidden_channels, 3, 1, 1, padding_mode="replicate"), | |
| nn.GELU(), | |
| nn.Conv2d(hidden_channels, hidden_channels, 3, 1, 1, padding_mode="replicate"), | |
| nn.GELU(), | |
| nn.Conv2d(hidden_channels, 1, 3, 1, 1, padding_mode="replicate"), | |
| ) | |
| # Predict the midpoint of the valid inverse-depth interval at initialization. | |
| nn.init.zeros_(self.inv_depth_logits[-1].weight) | |
| nn.init.zeros_(self.inv_depth_logits[-1].bias) | |
| def forward( | |
| self, | |
| latent: torch.Tensor, | |
| near: torch.Tensor, | |
| far: torch.Tensor, | |
| return_inverse: bool = False, | |
| return_normalized_inverse: bool = False, | |
| ) -> torch.Tensor | tuple[torch.Tensor, ...]: | |
| with torch.autocast(device_type=latent.device.type, enabled=False): | |
| inv_depth_logits = self.inv_depth_logits(latent.float()) | |
| near = near.float().clamp_min(torch.finfo(torch.float32).eps) | |
| far = far.float().clamp_min(torch.finfo(torch.float32).eps) | |
| inv_min = far.reciprocal() | |
| inv_max = near.reciprocal() | |
| inv_range = (inv_max - inv_min).clamp_min(torch.finfo(torch.float32).eps) | |
| normalized_inv_depth = torch.sigmoid(inv_depth_logits.float()) | |
| inv_depth = inv_min + normalized_inv_depth * inv_range | |
| depth = inv_depth.clamp_min(torch.finfo(torch.float32).eps).reciprocal() | |
| outputs: list[torch.Tensor] = [depth] | |
| if return_inverse: | |
| outputs.append(inv_depth) | |
| if return_normalized_inverse: | |
| outputs.append(normalized_inv_depth) | |
| if len(outputs) == 1: | |
| return outputs[0] | |
| return tuple(outputs) | |