Account for nerfacto vs splatfacto spherical harmonics differences

Author: brentyi

nerfstudio/scripts/exporter.py

@@ -32,9 +32,6 @@
 import open3d as o3d
 import torch
 import tyro
-from scipy.spatial.transform import Rotation as ScR
-from typing_extensions import Annotated, Literal
-
 from nerfstudio.cameras.rays import RayBundle
 from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager
 from nerfstudio.data.datamanagers.full_images_datamanager import FullImageDatamanager
@@ -50,6 +47,8 @@
 from nerfstudio.utils.eval_utils import eval_setup
 from nerfstudio.utils.rich_utils import CONSOLE
 from nerfstudio.utils.spherical_harmonics import rotate_spherical_harmonics
+from scipy.spatial.transform import Rotation as ScR
+from typing_extensions import Annotated, Literal
 
 
 @dataclass
@@ -636,9 +635,11 @@ def main(self) -> None:
                         dim_sh_all = shs_rest.shape[-1] + 1
                         shs_coeffs_all = torch.zeros((n, 3, dim_sh_all), device=shs_rest.device)
                         shs_coeffs_all[:, :, 1:] = shs_rest
-                        # TODO: check output rotation
-                        output_rotation = pipeline.datamanager.train_dataparser_outputs.dataparser_transform[:3, :3]
-                        shs_rest = rotate_spherical_harmonics(output_rotation, shs_coeffs_all)[:, :, 1:]
+                        shs_rest = rotate_spherical_harmonics(
+                            pipeline.datamanager.train_dataparser_outputs.dataparser_transform[:3, :3].T,
+                            shs_coeffs_all,
+                            component_convention="-y,+z,-x",
+                        )[:, :, 1:]
 
                     shs_rest = shs_rest.cpu().numpy().reshape((n, -1))
 

nerfstudio/utils/spherical_harmonics.py

@@ -15,11 +15,13 @@
 """Sphecal Harmonics utils."""
 
 import math
+from typing import Literal
 
 import torch
 from e3nn.o3 import Irreps
 from jaxtyping import Float
 from torch import Tensor
+from typing_extensions import assert_never
 
 MAX_SH_DEGREE = 4
 
@@ -117,12 +119,16 @@ def SH2RGB(sh):
 def rotate_spherical_harmonics(
     rotation_matrix: Float[Tensor, "3 3"],
     coeffs: Float[Tensor, "*batch dim_sh"],
+    component_convention: Literal["-y,+z,-x", "+y,+z,+x"],
 ) -> Float[Tensor, "*batch dim_sh"]:
     """Rotates real spherical harmonic coefficients using a given 3x3 rotation matrix.
 
     Args:
         rotation_matrix : A 3x3 rotation matrix.
         coeffs : SH coefficients
+        component_convention: Component convention for spherical harmonics.
+            Nerfstudio (nerfacto) uses +y,+z,+x, while gsplat (splatfacto) uses
+            -y,+z,-x.
 
     Returns:
         The rotated SH coefficients
@@ -132,11 +138,31 @@ def rotate_spherical_harmonics(
     sh_degree = int(math.sqrt(dim_sh)) - 1
 
     # e3nn uses the xyz ordering instead of the standard yzx used in ns, equivalent to a change of basis
-    R_yzx_to_xyz = torch.tensor([[0, 0, 1], [1, 0, 0], [0, 1, 0]], dtype=torch.float32)
-    R_total = (R_yzx_to_xyz.T @ rotation_matrix @ R_yzx_to_xyz).cpu()
+    if component_convention == "+y,+z,+x":
+        R_xyz_from_yzx = torch.tensor(
+            [
+                [0, 0, 1],
+                [1, 0, 0],
+                [0, 1, 0],
+            ],
+            dtype=torch.float32,
+        )
+        rotation_matrix = (R_xyz_from_yzx.T @ rotation_matrix @ R_xyz_from_yzx).cpu()
+    elif component_convention == "-y,+z,-x":
+        R_xyz_from_negyznegx = torch.tensor(
+            [
+                [0, 0, -1],
+                [-1, 0, 0],
+                [0, 1, 0],
+            ],
+            dtype=torch.float32,
+        )
+        rotation_matrix = (R_xyz_from_negyznegx.T @ rotation_matrix @ R_xyz_from_negyznegx).cpu()
+    else:
+        assert_never(component_convention)
 
     irreps = Irreps(" + ".join([f"{i}e" for i in range(sh_degree + 1)]))  # Even parity spherical harmonics of degree l
-    D_matrix = irreps.D_from_matrix(R_total).to(coeffs.device)  # Construct Wigner D-matrix
+    D_matrix = irreps.D_from_matrix(rotation_matrix).to(coeffs.device)  # Construct Wigner D-matrix
 
     # Multiply last dimension of coeffs (..., dim_sh) with the Wigner D-matrix (dim_sh, dim_sh)
     rotated_coeffs = coeffs @ D_matrix.T

pyproject.toml

@@ -147,7 +147,7 @@ include = ["nerfstudio*"]
 "*" = ["*.cu", "*.json", "py.typed", "setup.bash", "setup.zsh"]
 
 [tool.pytest.ini_options]
-addopts = "-n=4 --typeguard-packages=nerfstudio --disable-warnings"
+# addopts = "-n=4 --typeguard-packages=nerfstudio --disable-warnings"
 testpaths = [
     "tests",
 ]

tests/utils/test_spherical_harmonics.py

@@ -1,13 +1,16 @@
+from typing import Literal
+
 import numpy as np
 import pytest
 import torch
-from scipy.spatial.transform import Rotation as ScR
-
+from gsplat.cuda._torch_impl import _eval_sh_bases_fast as gsplat_eval_sh_bases
+from gsplat.cuda._torch_impl import _spherical_harmonics as gsplat_spherical_harmonics
 from nerfstudio.utils.spherical_harmonics import (
     components_from_spherical_harmonics,
     num_sh_bases,
     rotate_spherical_harmonics,
 )
+from scipy.spatial.transform import Rotation as ScR
 
 
 @pytest.mark.parametrize("degree", list(range(0, 5)))
@@ -23,7 +26,7 @@ def test_spherical_harmonics_components(degree):
 
 
 @pytest.mark.parametrize("sh_degree", list(range(0, 4)))
-def test_spherical_harmonics_rotation(sh_degree):
+def test_spherical_harmonics_rotation_nerfacto(sh_degree):
     """Test if rotating both the view direction and SH coefficients by the same rotation
     produces the same color output as the original.
 
@@ -43,7 +46,7 @@ def test_spherical_harmonics_rotation(sh_degree):
     color_original = (sh_coeffs * y_lm[..., None, :]).sum(dim=-1)
 
     rot_matrix = torch.tensor(ScR.random().as_matrix(), dtype=torch.float32)
-    sh_coeffs_rotated = rotate_spherical_harmonics(rot_matrix, sh_coeffs)
+    sh_coeffs_rotated = rotate_spherical_harmonics(rot_matrix, sh_coeffs, component_convention="+y,+z,+x")
     dirs_rotated = (rot_matrix @ dirs.T).T
     y_lm_rotated = components_from_spherical_harmonics(sh_degree, dirs_rotated)
     color_rotated = (sh_coeffs_rotated * y_lm_rotated[..., None, :]).sum(dim=-1)
@@ -52,7 +55,42 @@ def test_spherical_harmonics_rotation(sh_degree):
 
 
 @pytest.mark.parametrize("sh_degree", list(range(0, 4)))
-def test_spherical_harmonics_rotation_properties(sh_degree):
+def test_spherical_harmonics_rotation_splatfacto(sh_degree):
+    """Test if rotating both the view direction and SH coefficients by the same rotation
+    produces the same color output as the original.
+
+     In other words, for any rotation R:
+         color(dir, coeffs) = color(R @ dir, rotate_sh(R, coeffs))
+    """
+    torch.manual_seed(0)
+    np.random.seed(0)
+
+    N = 1000
+    num_coeffs = (sh_degree + 1) ** 2
+    sh_coeffs = torch.rand(N, 3, num_coeffs)
+    dirs = torch.rand(N, 3)
+    dirs = dirs / torch.linalg.norm(dirs, dim=-1, keepdim=True)
+
+    assert dirs.shape == (N, 3)
+    y_lm = gsplat_eval_sh_bases(num_coeffs, dirs)
+    color_original = (sh_coeffs * y_lm[..., None, :]).sum(dim=-1)
+
+    rot_matrix = torch.tensor(ScR.random().as_matrix(), dtype=torch.float32)
+    sh_coeffs_rotated = rotate_spherical_harmonics(rot_matrix, sh_coeffs, component_convention="-y,+z,-x")
+    dirs_rotated = (rot_matrix @ dirs.T).T
+    assert dirs_rotated.shape == (N, 3)
+    y_lm_rotated = gsplat_eval_sh_bases(num_coeffs, dirs_rotated)
+    color_rotated = (sh_coeffs_rotated * y_lm_rotated[..., None, :]).sum(dim=-1)
+
+    torch.testing.assert_close(
+        gsplat_spherical_harmonics(sh_degree, coeffs=sh_coeffs.swapaxes(-1, -2), dirs=dirs),
+        gsplat_spherical_harmonics(sh_degree, coeffs=sh_coeffs_rotated.swapaxes(-1, -2), dirs=dirs_rotated),
+    )
+
+
+@pytest.mark.parametrize("sh_degree", list(range(0, 4)))
+@pytest.mark.parametrize("component_convention", ["+y,+z,+x", "-y,+z,-x"])
+def test_spherical_harmonics_rotation_properties(sh_degree: int, component_convention: Literal["+y,+z,+x", "-y,+z,-x"]):
     """Test properties of the SH rotation"""
     torch.manual_seed(0)
     np.random.seed(0)
@@ -61,7 +99,7 @@ def test_spherical_harmonics_rotation_properties(sh_degree):
     num_coeffs = (sh_degree + 1) ** 2
     sh_coeffs = torch.rand(N, 3, num_coeffs)
     rot_matrix = torch.tensor(ScR.random().as_matrix(), dtype=torch.float32)
-    sh_coeffs_rotated = rotate_spherical_harmonics(rot_matrix, sh_coeffs)
+    sh_coeffs_rotated = rotate_spherical_harmonics(rot_matrix, sh_coeffs, component_convention)
 
     # Norm preserving
     norm_original = torch.norm(sh_coeffs, dim=-1)
@@ -73,5 +111,5 @@ def test_spherical_harmonics_rotation_properties(sh_degree):
 
     # Identity rotation
     rot_matrix = torch.eye(3)
-    sh_coeffs_rotated = rotate_spherical_harmonics(rot_matrix, sh_coeffs)
+    sh_coeffs_rotated = rotate_spherical_harmonics(rot_matrix, sh_coeffs, ordering)
     torch.testing.assert_close(sh_coeffs, sh_coeffs_rotated, rtol=0, atol=1e-6)