Construct Gaussian base values and prepare feature input. Args: image: The image to process. depth: The corresponding depth map from the monodepth network. Returns: The base value for Gaussians.
(self, image: torch.Tensor, depth: torch.Tensor)
| 125 | return features_in |
| 126 | |
| 127 | def forward(self, image: torch.Tensor, depth: torch.Tensor) -> InitializerOutput: |
| 128 | """Construct Gaussian base values and prepare feature input. |
| 129 | |
| 130 | Args: |
| 131 | image: The image to process. |
| 132 | depth: The corresponding depth map from the monodepth network. |
| 133 | |
| 134 | Returns: |
| 135 | The base value for Gaussians. |
| 136 | """ |
| 137 | image = image.contiguous() |
| 138 | depth = depth.contiguous() |
| 139 | device = depth.device |
| 140 | batch_size, _, image_height, image_width = depth.shape |
| 141 | base_height, base_width = ( |
| 142 | image_height // self.stride, |
| 143 | image_width // self.stride, |
| 144 | ) |
| 145 | # global_scale is the inverse of the depth_factor, which is used to rescale |
| 146 | # the depth such that it is numerically stable for training. |
| 147 | global_scale: torch.Tensor | None = None |
| 148 | if self.normalize_depth: |
| 149 | depth, depth_factor = _rescale_depth(depth) |
| 150 | global_scale = 1.0 / depth_factor |
| 151 | |
| 152 | def _create_disparity_layers(num_layers: int = 1) -> torch.Tensor: |
| 153 | """Create multiple disparity layers.""" |
| 154 | disparity = torch.linspace(1.0 / self.base_depth, 0.0, num_layers + 1, device=device) |
| 155 | return disparity[None, None, :-1, None, None].repeat( |
| 156 | batch_size, 1, 1, base_height, base_width |
| 157 | ) |
| 158 | |
| 159 | def _create_surface_layer( |
| 160 | depth: torch.Tensor, |
| 161 | depth_pooling_mode: str, |
| 162 | ) -> torch.Tensor: |
| 163 | """Create multiple surface layers.""" |
| 164 | disparity = 1.0 / depth |
| 165 | if depth_pooling_mode == "min": |
| 166 | disparity = torch.max_pool2d(disparity, self.stride, self.stride) |
| 167 | elif depth_pooling_mode == "max": |
| 168 | disparity = -torch.max_pool2d(-disparity, self.stride, self.stride) |
| 169 | else: |
| 170 | raise ValueError(f"Invalid depth pooling mode {depth_pooling_mode}.") |
| 171 | |
| 172 | return disparity[:, :, None, :, :] |
| 173 | |
| 174 | # Input disparity dimensions: |
| 175 | # (batch_size, num_channels in (1, 2), height, width) |
| 176 | |
| 177 | # Output disparity dimensions: |
| 178 | # (batch_size, num_channels=1, num_layers in (1, 2), height, width) |
| 179 | if self.first_layer_depth_option == "surface_min": |
| 180 | first_disparity = _create_surface_layer(depth[:, 0:1], "min") |
| 181 | elif self.first_layer_depth_option == "surface_max": |
| 182 | first_disparity = _create_surface_layer(depth[:, 0:1], "max") |
| 183 | elif self.first_layer_depth_option in ("base_depth", "linear_disparity"): |
| 184 | first_disparity = _create_disparity_layers() |
nothing calls this directly
no test coverage detected