Complete merge

nerfstudio/scripts/exporter.py

@@ -48,8 +48,8 @@ from nerfstudio.fields.sdf_field import SDFField  # noqa
 from nerfstudio.models.splatfacto import SplatfactoModel
 from nerfstudio.pipelines.base_pipeline import Pipeline, VanillaPipeline
 from nerfstudio.utils.eval_utils import eval_setup
-from nerfstudio.utils.math import rotate_spherical_harmonics
 from nerfstudio.utils.rich_utils import CONSOLE
+from nerfstudio.utils.spherical_harmonics import rotate_spherical_harmonics
 
 
 @dataclass
@@ -494,7 +494,6 @@ class ExportGaussianSplat(Exporter):
     Export 3D Gaussian Splatting model to a .ply
     """
 
-
     output_filename: str = "splat.ply"
     """Name of the output file."""
     save_world_frame: bool = True
@@ -582,10 +581,9 @@ class ExportGaussianSplat(Exporter):
         map_to_tensors = OrderedDict()
 
         with torch.no_grad():
-
             positions = model.means.cpu().numpy()
             quats = model.quats.data.cpu().numpy()
-            
+
             if self.save_world_frame:
                 assert isinstance(pipeline.datamanager, FullImageDatamanager)
                 # ns gplat uses quaternions in [w,x,y,z] format while scipy uses [x, y, z, w]
@@ -598,17 +596,10 @@ class ExportGaussianSplat(Exporter):
                 poses = pipeline.datamanager.train_dataparser_outputs.transform_poses_to_original_space(
                     torch.from_numpy(poses)
                 )
-                
-                rot_world = ScR.from_matrix(poses[:, :3, :3].numpy())
-                quats = rot_world.as_quat()[:, [3, 0, 1, 2]] # convert back to [w,x,y,z] format
-                positions = poses[:, :3, 3].cpu().numpy()
-                
-                dataparser_outputs = pipeline.datamanager.train_dataparser_outputs
-                dataparser_scale = dataparser_outputs.dataparser_scale
-                dataparser_transform = dataparser_outputs.dataparser_transform
 
-                output_scale = 1 / dataparser_scale
-                output_rotation = dataparser_transform[:3, :3].T
+                rot_world = ScR.from_matrix(poses[:, :3, :3].numpy())
+                quats = rot_world.as_quat()[:, [3, 0, 1, 2]]  # convert back to [w,x,y,z] format
+                positions = poses[:, :3, 3].cpu().numpy()
 
             count = positions.shape[0]
             n = count
@@ -637,16 +628,17 @@ class ExportGaussianSplat(Exporter):
                     )
                 elif self.ply_color_mode == "sh_coeffs":
                     # transpose(1, 2) was needed to match the sh order in Inria version
-                    shs_rest = model.shs_rest.transpose(1, 2).contiguous() # (n, 3, dim_sh)
+                    shs_rest = model.shs_rest.transpose(1, 2).contiguous()  # (n, 3, dim_sh)
                     if self.save_world_frame:
-                        # The 0th order coefficients (l=0) are constant and invariant to rotation. 
+                        # The 0th order coefficients (l=0) are constant and invariant to rotation.
                         # They don't affect the rotation of the rest of the SH.
-                        shs_coeffs_all = torch.zeros((n, 3, shs_rest.shape[-1]+1), device=shs_rest.device)
+                        shs_coeffs_all = torch.zeros((n, 3, shs_rest.shape[-1] + 1), device=shs_rest.device)
                         shs_coeffs_all[:, :, 1:] = shs_rest
+                        output_rotation = pipeline.datamanager.train_dataparser_outputs.dataparser_transform[:3, :3].T
                         shs_rest = rotate_spherical_harmonics(shs_coeffs_all, output_rotation)[:, :, 1:]
 
                     shs_rest = shs_rest.cpu().numpy().reshape((n, -1))
-                    
+
                     for i in range(shs_rest.shape[-1]):
                         map_to_tensors[f"f_rest_{i}"] = shs_rest[:, i, None]
 
@@ -655,8 +647,8 @@ class ExportGaussianSplat(Exporter):
             # Note that scales are in log space!
             scales = model.scales.data.cpu().numpy()
             if self.save_world_frame:
-                scales += np.log(output_scale)
-                
+                scales += np.log(pipeline.datamanager.train_dataparser_outputs.dataparser_scale)
+
             for i in range(3):
                 map_to_tensors[f"scale_{i}"] = scales[:, i, None]
 

nerfstudio/utils/spherical_harmonics.py

@@ -14,7 +14,10 @@
 
 """Sphecal Harmonics utils."""
 
+import math
+
 import torch
+from e3nn.o3 import Irreps
 from jaxtyping import Float
 from torch import Tensor
 
@@ -109,3 +112,27 @@ def SH2RGB(sh):
     """
     C0 = 0.28209479177387814
     return sh * C0 + 0.5
+
+
+def rotate_spherical_harmonics(
+    coeffs: Float[Tensor, "*batch dim_sh"],
+    rotation_matrix: Float[Tensor, "3 3"],
+) -> Float[Tensor, "*batch dim_sh"]:
+    """Rotates real spherical harmonic coefficients of degree l using given 3x3 rotation matrix.
+
+    Args:
+        coeffs : SH coefficients
+        rotation_matrix : A 3x3 rotation matrix.
+
+    Returns:
+        The rotated SH coefficients
+    """
+    dim_sh = coeffs.shape[-1]
+    assert math.isqrt(dim_sh) ** 2 == dim_sh, "dim_sh must be a perfect square (l+1)^2"
+    sh_degree = int(math.sqrt(dim_sh)) - 1
+
+    rotation_matrix = rotation_matrix.cpu()
+    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(rotation_matrix).to(coeffs.device)  # Construct Wigner D-matrix
+    rotated_coeffs = torch.einsum("ij,...j->...i", D_matrix, coeffs)
+    return rotated_coeffs

tests/utils/test_spherical_harmonics.py

@@ -1,7 +1,13 @@
+import numpy as np
 import pytest
 import torch
+from scipy.spatial.transform import Rotation as ScR
 
-from nerfstudio.utils.spherical_harmonics import components_from_spherical_harmonics, num_sh_bases
+from nerfstudio.utils.spherical_harmonics import (
+    components_from_spherical_harmonics,
+    num_sh_bases,
+    rotate_spherical_harmonics,
+)
 
 
 @pytest.mark.parametrize("degree", list(range(0, 5)))
@@ -14,3 +20,33 @@ def test_spherical_harmonics(degree):
     sh = components_from_spherical_harmonics(degree, dx)
     matrix = (sh.T @ sh) / N * 4 * torch.pi
     torch.testing.assert_close(matrix, torch.eye(num_sh_bases(degree)), rtol=0, atol=1.5e-2)
+
+
+@pytest.mark.parametrize("sh_degree", list(range(0, 4)))
+def test_rotate_spherical_harmonics(sh_degree):
+    torch.manual_seed(0)
+    np.random.seed(0)
+
+    num_coeffs = (sh_degree + 1) ** 2
+    coeffs = torch.rand(1000, 3, num_coeffs)
+
+    rot_matrix = torch.tensor(ScR.random().as_matrix(), dtype=torch.float32)
+    rotated_coeffs = rotate_spherical_harmonics(coeffs, rot_matrix)
+
+    # Norm Preservation
+    norm_original = torch.norm(coeffs, dim=-1)
+    norm_rotated = torch.norm(rotated_coeffs, dim=-1)
+    torch.testing.assert_close(norm_original, norm_rotated, rtol=0, atol=1e-5)
+
+    # 0th order coeffs are invariant to rotation
+    torch.testing.assert_close(coeffs[..., 0], rotated_coeffs[..., 0], rtol=0, atol=1e-6)
+
+    # 1st order coeffs undergo simple rotation
+    if sh_degree > 0:
+        act_1st_order = torch.einsum("ij,...j->...i", rot_matrix, coeffs[..., 1:4])
+        torch.testing.assert_close(act_1st_order, rotated_coeffs[..., 1:4], rtol=0, atol=1e-5)
+
+    # Identity Rotation
+    rot_matrix = torch.eye(3)
+    rotated_coeffs = rotate_spherical_harmonics(coeffs, rot_matrix)
+    torch.testing.assert_close(coeffs, rotated_coeffs, rtol=0, atol=1e-6)