A version of rasterization() that utilies on PyTorch's autograd. .. note:: This function still relies on gsplat's CUDA backend for some computation, but the entire differentiable graph is on of PyTorch (and nerfacc) so could use Pytorch's autograd for backpropagation.
(
means: Tensor, # [N, 3]
quats: Tensor, # [N, 4]
scales: Tensor, # [N, 3]
opacities: Tensor, # [N]
colors: Tensor, # [(C,) N, D] or [(C,) N, K, 3]
viewmats: Tensor, # [C, 4, 4]
Ks: Tensor, # [C, 3, 3]
width: int,
height: int,
near_plane: float = 0.01,
far_plane: float = 1e10,
eps2d: float = 0.3,
sh_degree: Optional[int] = None,
tile_size: int = 16,
backgrounds: Optional[Tensor] = None,
render_mode: Literal["RGB", "D", "ED", "RGB+D", "RGB+ED"] = "RGB",
rasterize_mode: Literal["classic", "antialiased"] = "classic",
channel_chunk: int = 32,
batch_per_iter: int = 100,
)
| 583 | |
| 584 | |
| 585 | def _rasterization( |
| 586 | means: Tensor, # [N, 3] |
| 587 | quats: Tensor, # [N, 4] |
| 588 | scales: Tensor, # [N, 3] |
| 589 | opacities: Tensor, # [N] |
| 590 | colors: Tensor, # [(C,) N, D] or [(C,) N, K, 3] |
| 591 | viewmats: Tensor, # [C, 4, 4] |
| 592 | Ks: Tensor, # [C, 3, 3] |
| 593 | width: int, |
| 594 | height: int, |
| 595 | near_plane: float = 0.01, |
| 596 | far_plane: float = 1e10, |
| 597 | eps2d: float = 0.3, |
| 598 | sh_degree: Optional[int] = None, |
| 599 | tile_size: int = 16, |
| 600 | backgrounds: Optional[Tensor] = None, |
| 601 | render_mode: Literal["RGB", "D", "ED", "RGB+D", "RGB+ED"] = "RGB", |
| 602 | rasterize_mode: Literal["classic", "antialiased"] = "classic", |
| 603 | channel_chunk: int = 32, |
| 604 | batch_per_iter: int = 100, |
| 605 | ) -> Tuple[Tensor, Tensor, Dict]: |
| 606 | """A version of rasterization() that utilies on PyTorch's autograd. |
| 607 | |
| 608 | .. note:: |
| 609 | This function still relies on gsplat's CUDA backend for some computation, but the |
| 610 | entire differentiable graph is on of PyTorch (and nerfacc) so could use Pytorch's |
| 611 | autograd for backpropagation. |
| 612 | |
| 613 | .. note:: |
| 614 | This function relies on installing latest nerfacc, via: |
| 615 | pip install git+https://github.com/nerfstudio-project/nerfacc |
| 616 | |
| 617 | .. note:: |
| 618 | Compared to rasterization(), this function does not support some arguments such as |
| 619 | `packed`, `sparse_grad` and `absgrad`. |
| 620 | """ |
| 621 | from gsplat.cuda._torch_impl import ( |
| 622 | _fully_fused_projection, |
| 623 | _quat_scale_to_covar_preci, |
| 624 | _rasterize_to_pixels, |
| 625 | ) |
| 626 | |
| 627 | N = means.shape[0] |
| 628 | C = viewmats.shape[0] |
| 629 | assert means.shape == (N, 3), means.shape |
| 630 | assert quats.shape == (N, 4), quats.shape |
| 631 | assert scales.shape == (N, 3), scales.shape |
| 632 | assert opacities.shape == (N,), opacities.shape |
| 633 | assert viewmats.shape == (C, 4, 4), viewmats.shape |
| 634 | assert Ks.shape == (C, 3, 3), Ks.shape |
| 635 | assert render_mode in ["RGB", "D", "ED", "RGB+D", "RGB+ED"], render_mode |
| 636 | |
| 637 | if sh_degree is None: |
| 638 | # treat colors as post-activation values, should be in shape [N, D] or [C, N, D] |
| 639 | assert (colors.dim() == 2 and colors.shape[0] == N) or ( |
| 640 | colors.dim() == 3 and colors.shape[:2] == (C, N) |
| 641 | ), colors.shape |
| 642 | else: |