Maps projected Gaussians to intersecting tiles. Args: means2d: Projected Gaussian means. [..., N, 2] if packed is False, [nnz, 2] if packed is True. radii: Maximum radii of the projected Gaussians. [..., N, 2] if packed is False, [nnz, 2] if packed is True. depths: Z-dep
(
means2d: Tensor, # [..., N, 2] or [nnz, 2]
radii: Tensor, # [..., N, 2] or [nnz, 2]
depths: Tensor, # [..., N] or [nnz]
tile_size: int,
tile_width: int,
tile_height: int,
sort: bool = True,
segmented: bool = False,
packed: bool = False,
n_images: Optional[int] = None,
image_ids: Optional[Tensor] = None,
gaussian_ids: Optional[Tensor] = None,
)
| 441 | |
| 442 | @torch.no_grad() |
| 443 | def isect_tiles( |
| 444 | means2d: Tensor, # [..., N, 2] or [nnz, 2] |
| 445 | radii: Tensor, # [..., N, 2] or [nnz, 2] |
| 446 | depths: Tensor, # [..., N] or [nnz] |
| 447 | tile_size: int, |
| 448 | tile_width: int, |
| 449 | tile_height: int, |
| 450 | sort: bool = True, |
| 451 | segmented: bool = False, |
| 452 | packed: bool = False, |
| 453 | n_images: Optional[int] = None, |
| 454 | image_ids: Optional[Tensor] = None, |
| 455 | gaussian_ids: Optional[Tensor] = None, |
| 456 | ) -> Tuple[Tensor, Tensor, Tensor]: |
| 457 | """Maps projected Gaussians to intersecting tiles. |
| 458 | |
| 459 | Args: |
| 460 | means2d: Projected Gaussian means. [..., N, 2] if packed is False, [nnz, 2] if packed is True. |
| 461 | radii: Maximum radii of the projected Gaussians. [..., N, 2] if packed is False, [nnz, 2] if packed is True. |
| 462 | depths: Z-depth of the projected Gaussians. [..., N] if packed is False, [nnz] if packed is True. |
| 463 | tile_size: Tile size. |
| 464 | tile_width: Tile width. |
| 465 | tile_height: Tile height. |
| 466 | sort: If True, the returned intersections will be sorted by the intersection ids. Default: True. |
| 467 | segmented: If True, segmented radix sort will be used to sort the intersections. Default: False. |
| 468 | packed: If True, the input tensors are packed. Default: False. |
| 469 | n_images: Number of images. Required if packed is True. |
| 470 | image_ids: The image indices of the projected Gaussians. Required if packed is True. |
| 471 | gaussian_ids: The column indices of the projected Gaussians. Required if packed is True. |
| 472 | |
| 473 | Returns: |
| 474 | A tuple: |
| 475 | |
| 476 | - **Tiles per Gaussian**. The number of tiles intersected by each Gaussian. |
| 477 | Int32 [..., N] if packed is False, Int32 [nnz] if packed is True. |
| 478 | - **Intersection ids**. Each id is an 64-bit integer with the following |
| 479 | information: image_id (Xc bits) | tile_id (Xt bits) | depth (32 bits). |
| 480 | Xc and Xt are the maximum number of bits required to represent the image and |
| 481 | tile ids, respectively. Int64 [n_isects] |
| 482 | - **Flatten ids**. The global flatten indices in [I * N] or [nnz] (packed). [n_isects] |
| 483 | """ |
| 484 | if packed: |
| 485 | nnz = means2d.size(0) |
| 486 | assert means2d.shape == (nnz, 2), means2d.shape |
| 487 | assert radii.shape == (nnz, 2), radii.shape |
| 488 | assert depths.shape == (nnz,), depths.shape |
| 489 | assert image_ids is not None, "image_ids is required if packed is True" |
| 490 | assert gaussian_ids is not None, "gaussian_ids is required if packed is True" |
| 491 | assert n_images is not None, "n_images is required if packed is True" |
| 492 | image_ids = image_ids.contiguous() |
| 493 | gaussian_ids = gaussian_ids.contiguous() |
| 494 | I = n_images |
| 495 | |
| 496 | else: |
| 497 | image_dims = means2d.shape[:-2] |
| 498 | I = math.prod(image_dims) |
| 499 | N = means2d.shape[-2] |
| 500 | assert means2d.shape == image_dims + (N, 2), means2d.shape |
searching dependent graphs…