From 0c3a4cdcef8f79e8f7adf2054d594a63b6ca6955 Mon Sep 17 00:00:00 2001 From: Aayush Date: Sat, 18 Jan 2025 20:43:11 -0800 Subject: [PATCH] Complete merge --- nerfstudio/scripts/exporter.py | 32 ++++++++------------- nerfstudio/utils/spherical_harmonics.py | 27 ++++++++++++++++++ tests/utils/test_spherical_harmonics.py | 38 ++++++++++++++++++++++++- 3 files changed, 76 insertions(+), 21 deletions(-) diff --git a/nerfstudio/scripts/exporter.py b/nerfstudio/scripts/exporter.py index 4a62428f2..b6985b2c9 100644 --- a/nerfstudio/scripts/exporter.py +++ b/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] diff --git a/nerfstudio/utils/spherical_harmonics.py b/nerfstudio/utils/spherical_harmonics.py index 07936281c..a4b97479b 100644 --- a/nerfstudio/utils/spherical_harmonics.py +++ b/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 diff --git a/tests/utils/test_spherical_harmonics.py b/tests/utils/test_spherical_harmonics.py index a8949891a..b2e98ea1c 100644 --- a/tests/utils/test_spherical_harmonics.py +++ b/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)