Create base x and y coordinates for the gaussians in NDC space.
(
depth: torch.Tensor, stride: int, num_layers: int
)
| 254 | |
| 255 | |
| 256 | def _create_base_xy( |
| 257 | depth: torch.Tensor, stride: int, num_layers: int |
| 258 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 259 | """Create base x and y coordinates for the gaussians in NDC space.""" |
| 260 | device = depth.device |
| 261 | batch_size, _, image_height, image_width = depth.shape |
| 262 | xx = torch.arange(0.5 * stride, image_width, stride, device=device) |
| 263 | yy = torch.arange(0.5 * stride, image_height, stride, device=device) |
| 264 | xx = 2 * xx / image_width - 1.0 |
| 265 | yy = 2 * yy / image_height - 1.0 |
| 266 | |
| 267 | xx, yy = torch.meshgrid(xx, yy, indexing="xy") |
| 268 | base_x_ndc = xx[None, None, None].repeat(batch_size, 1, num_layers, 1, 1) |
| 269 | base_y_ndc = yy[None, None, None].repeat(batch_size, 1, num_layers, 1, 1) |
| 270 | |
| 271 | return base_x_ndc, base_y_ndc |
| 272 | |
| 273 | |
| 274 | def _create_base_scale(disparity: torch.Tensor, disparity_scale_factor: float) -> torch.Tensor: |