gsplat
gsplat copied to clipboard
nerfstudio-project
`rasterization_2dgs` rendering no longer works when `sh_degree=None` as of v1.5.3
Note: this only appears to affect 2DGS, not 3DGS !
After pip install git+https://github.com/nerfstudio-project/[email protected]
Taking the example from rasterization_2dgs docstring:
from gsplat import rasterization_2dgs
import torch
device = "cuda"
means = torch.randn((100, 3), device=device)
quats = torch.randn((100, 4), device=device)
scales = torch.rand((100, 3), device=device) * 0.1
colors = torch.rand((100, 3), device=device)
opacities = torch.rand((100,), device=device)
# define cameras
viewmats = torch.eye(4, device=device)[None, :, :]
Ks = torch.tensor([
[300., 0., 150.], [0., 300., 100.], [0., 0., 1.]], device=device)[None, :, :]
width, height = 300, 200
# render
colors, alphas, normals, surf_normals, distort, median_depth, meta = rasterization_2dgs(
means, quats, scales, opacities, colors, viewmats, Ks, width, height
)
Results in this error:
---------------------------------------------------------------------------
AssertionError Traceback (most recent call last)
Cell In[1], [line 17](vscode-notebook-cell:?execution_count=1&line=17)
15 width, height = 300, 200
16 # render
---> [17](vscode-notebook-cell:?execution_count=1&line=17) colors, alphas, normals, surf_normals, distort, median_depth, meta = rasterization_2dgs(
18 means, quats, scales, opacities, colors, viewmats, Ks, width, height
19 )
File .venv/lib/python3.10/site-packages/gsplat/rendering.py:1568, in rasterization_2dgs(means, quats, scales, opacities, colors, viewmats, Ks, width, height, near_plane, far_plane, radius_clip, eps2d, sh_degree, packed, tile_size, backgrounds, render_mode, sparse_grad, absgrad, distloss, depth_mode)
1559 else: # RGB
1560 pass
1562 (
1563 render_colors,
1564 render_alphas,
1565 render_normals,
1566 render_distort,
1567 render_median,
-> [1568](.venv/lib/python3.10/site-packages/gsplat/rendering.py:1568) ) = rasterize_to_pixels_2dgs(
1569 means2d,
1570 ray_transforms,
1571 colors,
1572 opacities,
1573 normals,
1574 densify,
1575 width,
1576 height,
1577 tile_size,
1578 isect_offsets,
1579 flatten_ids,
1580 backgrounds=backgrounds,
1581 packed=packed,
1582 absgrad=absgrad,
1583 distloss=distloss,
1584 )
1585 render_normals_from_depth = None
1586 if render_mode in ["ED", "RGB+ED"]:
1587 # normalize the accumulated depth to get the expected depth
File .venv/lib/python3.10/site-packages/gsplat/cuda/_wrapper.py:2283, in rasterize_to_pixels_2dgs(means2d, ray_transforms, colors, opacities, normals, densify, image_width, image_height, tile_size, isect_offsets, flatten_ids, backgrounds, masks, packed, absgrad, distloss)
2281 assert means2d.shape == image_dims + (N, 2), means2d.shape
2282 assert ray_transforms.shape == image_dims + (N, 3, 3), ray_transforms.shape
-> [2283](.venv/lib/python3.10/site-packages/gsplat/cuda/_wrapper.py:2283) assert colors.shape[:-2] == image_dims, colors.shape
2284 assert opacities.shape == image_dims + (N,), opacities.shape
2285 if backgrounds is not None:
AssertionError: torch.Size([100, 3])
Adding a batch dimension does not help:
from gsplat import rasterization_2dgs
import torch
device = "cuda"
means = torch.randn((2, 100, 3), device=device)
quats = torch.randn((2, 100, 4), device=device)
scales = torch.rand((2, 100, 3), device=device) * 0.1
colors = torch.rand((2, 100, 3), device=device)
opacities = torch.rand((2, 100,), device=device)
# define cameras
viewmats = torch.eye(4, device=device)[None, :, :]
Ks = torch.tensor([
[300., 0., 150.], [0., 300., 100.], [0., 0., 1.]], device=device)[None, :, :]
viewmats = viewmats.unsqueeze(1).repeat(2, 1, 1, 1)
Ks = Ks.unsqueeze(1).repeat(2, 1, 1, 1)
width, height = 300, 200
# render
colors, alphas, normals, surf_normals, distort, median_depth, meta = rasterization_2dgs(
means, quats, scales, opacities, colors, viewmats, Ks, width, height
)
---------------------------------------------------------------------------
AssertionError Traceback (most recent call last)
Cell In[2], [line 22](vscode-notebook-cell:?execution_count=2&line=22)
20 width, height = 300, 200
21 # render
---> [22](vscode-notebook-cell:?execution_count=2&line=22) colors, alphas, normals, surf_normals, distort, median_depth, meta = rasterization_2dgs(
23 means, quats, scales, opacities, colors, viewmats, Ks, width, height
24 )
File .venv/lib/python3.10/site-packages/gsplat/rendering.py:1568, in rasterization_2dgs(means, quats, scales, opacities, colors, viewmats, Ks, width, height, near_plane, far_plane, radius_clip, eps2d, sh_degree, packed, tile_size, backgrounds, render_mode, sparse_grad, absgrad, distloss, depth_mode)
1559 else: # RGB
1560 pass
1562 (
1563 render_colors,
1564 render_alphas,
1565 render_normals,
1566 render_distort,
1567 render_median,
-> [1568](.venv/lib/python3.10/site-packages/gsplat/rendering.py:1568) ) = rasterize_to_pixels_2dgs(
1569 means2d,
1570 ray_transforms,
1571 colors,
1572 opacities,
1573 normals,
1574 densify,
1575 width,
1576 height,
1577 tile_size,
1578 isect_offsets,
1579 flatten_ids,
1580 backgrounds=backgrounds,
1581 packed=packed,
1582 absgrad=absgrad,
1583 distloss=distloss,
1584 )
1585 render_normals_from_depth = None
1586 if render_mode in ["ED", "RGB+ED"]:
1587 # normalize the accumulated depth to get the expected depth
File .venv/lib/python3.10/site-packages/gsplat/cuda/_wrapper.py:2283, in rasterize_to_pixels_2dgs(means2d, ray_transforms, colors, opacities, normals, densify, image_width, image_height, tile_size, isect_offsets, flatten_ids, backgrounds, masks, packed, absgrad, distloss)
2281 assert means2d.shape == image_dims + (N, 2), means2d.shape
2282 assert ray_transforms.shape == image_dims + (N, 3, 3), ray_transforms.shape
-> [2283](.venv/lib/python3.10/site-packages/gsplat/cuda/_wrapper.py:2283) assert colors.shape[:-2] == image_dims, colors.shape
2284 assert opacities.shape == image_dims + (N,), opacities.shape
2285 if backgrounds is not None:
AssertionError: torch.Size([2, 100, 3])
BUT if we add an sh dimension to colors and set sh_degree=0, then it works:
from gsplat import rasterization_2dgs
import torch
device = "cuda"
means = torch.randn((2, 100, 3), device=device)
quats = torch.randn((2, 100, 4), device=device)
scales = torch.rand((2, 100, 3), device=device) * 0.1
colors = torch.rand((2, 100, 1, 3), device=device)
opacities = torch.rand((2, 100,), device=device)
# define cameras
viewmats = torch.eye(4, device=device)[None, :, :]
Ks = torch.tensor([
[300., 0., 150.], [0., 300., 100.], [0., 0., 1.]], device=device)[None, :, :]
viewmats = viewmats.unsqueeze(1).repeat(2, 1, 1, 1)
Ks = Ks.unsqueeze(1).repeat(2, 1, 1, 1)
width, height = 300, 200
# render
colors, alphas, normals, surf_normals, distort, median_depth, meta = rasterization_2dgs(
means, quats, scales, opacities, colors, viewmats, Ks, width, height, sh_degree=0
)
<No errors>
There appears to be a bug introduced with the batching introduced v1.5.3, specifically for rasterization_2dgs (great feature btw, thank you!)
also write codes under rasterization_2dgs, and it works
if sh_degree is None: # SH coefficients
# Colors are post-activation values, with shape [..., N, D] or [..., C, N, D]
if packed:
if colors.dim() == num_batch_dims + 2:
# Turn [..., N, D] into [nnz, D]
colors = colors.view(B, N, -1)[batch_ids, gaussian_ids]
else:
# Turn [..., C, N, D] into [nnz, D]
colors = colors.view(B, C, N, -1)[batch_ids, camera_ids, gaussian_ids]
else:
if colors.dim() == num_batch_dims + 2:
# Turn [..., N, D] into [..., C, N, D]
colors = torch.broadcast_to(
colors[..., None, :, :], batch_dims + (C, N, -1)
)
else:
# colors is already [..., C, N, D]
pass
