From cd0ff1dfb40facf91e2869789cbd5e39237955bf Mon Sep 17 00:00:00 2001 From: Aayush Date: Sat, 18 Jan 2025 01:48:53 -0800 Subject: [PATCH] SH rotation, cleanup --- nerfstudio/scripts/exporter.py | 73 +++++++++++++++++-------------- nerfstudio/utils/math.py | 25 +++++++++++ nerfstudio/viewer/export_panel.py | 6 +-- tests/utils/test_math.py | 31 ++++++++++++- 4 files changed, 97 insertions(+), 38 deletions(-) diff --git a/nerfstudio/scripts/exporter.py b/nerfstudio/scripts/exporter.py index c58cd5c5a..778f56d2c 100644 --- a/nerfstudio/scripts/exporter.py +++ b/nerfstudio/scripts/exporter.py @@ -31,7 +31,9 @@ import numpy as np import open3d as o3d import torch import tyro -from nerfstudio.cameras.camera_utils import quaternion_from_matrix +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 @@ -45,8 +47,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 typing_extensions import Annotated, Literal @dataclass @@ -567,26 +569,34 @@ 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] + rot_ns = ScR.from_quat(quats[:, [1, 2, 3, 0]]).as_matrix() + + # apply the inverse dataparser transform to the splat + poses = np.eye(4, dtype=np.float32)[None, ...].repeat(positions.shape[0], axis=0)[:, :3, :] + poses[:, :3, :3] = rot_ns + poses[:, :3, 3] = positions + 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.numpy(force=True) + dataparser_transform = dataparser_outputs.dataparser_transform output_scale = 1 / dataparser_scale - output_transform = np.zeros((3, 4)) - output_transform[:3, :3] = dataparser_transform[:3, :3].T - output_transform[:3, 3] = -dataparser_transform[:3, :3].T @ dataparser_transform[:3, 3] - else: - output_scale = 1 - output_transform = np.zeros((3, 4)) - output_transform[:3, :3] = np.eye(3) - inv_dataparser_quat = quaternion_from_matrix(output_transform[:3, :3]) + output_rotation = dataparser_transform[:3, :3].T - positions = ( - np.einsum("ij,bj->bi", output_transform[:3, :3], model.means.cpu().numpy() * output_scale) - + output_transform[None, :3, 3] - ) count = positions.shape[0] n = count map_to_tensors["x"] = positions[:, 0] @@ -614,34 +624,29 @@ 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().cpu().numpy() - shs_rest = shs_rest.reshape((n, -1)) + 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. + # 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[:, :, 1:] = shs_rest + 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] map_to_tensors["opacity"] = model.opacities.data.cpu().numpy() # Note that scales are in log space! - scales = model.scales.data.cpu().numpy() + np.log(output_scale) + scales = model.scales.data.cpu().numpy() + if self.save_world_frame: + scales += np.log(output_scale) + for i in range(3): map_to_tensors[f"scale_{i}"] = scales[:, i, None] - def quaternion_multiply(wxyz0: np.ndarray, wxyz1: np.ndarray) -> np.ndarray: - assert wxyz0.shape[-1] == 4 - assert wxyz1.shape[-1] == 4 - w0, x0, y0, z0 = np.moveaxis(wxyz0, -1, 0) - w1, x1, y1, z1 = np.moveaxis(wxyz1, -1, 0) - return np.stack( - [ - -x0 * x1 - y0 * y1 - z0 * z1 + w0 * w1, - x0 * w1 + y0 * z1 - z0 * y1 + w0 * x1, - -x0 * z1 + y0 * w1 + z0 * x1 + w0 * y1, - x0 * y1 - y0 * x1 + z0 * w1 + w0 * z1, - ], - axis=-1, - ) - - quats = quaternion_multiply(inv_dataparser_quat, model.quats.data.cpu().numpy()) for i in range(4): map_to_tensors[f"rot_{i}"] = quats[:, i, None] diff --git a/nerfstudio/utils/math.py b/nerfstudio/utils/math.py index d71907bee..3f6ba18f2 100644 --- a/nerfstudio/utils/math.py +++ b/nerfstudio/utils/math.py @@ -20,6 +20,7 @@ from dataclasses import dataclass from typing import Literal, Tuple import torch +from e3nn.o3 import Irreps from jaxtyping import Bool, Float from torch import Tensor @@ -92,6 +93,30 @@ def components_from_spherical_harmonics( return components +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 + + @dataclass class Gaussians: """Stores Gaussians diff --git a/nerfstudio/viewer/export_panel.py b/nerfstudio/viewer/export_panel.py index 4c24e2c7f..db4bcd292 100644 --- a/nerfstudio/viewer/export_panel.py +++ b/nerfstudio/viewer/export_panel.py @@ -16,14 +16,14 @@ from __future__ import annotations from pathlib import Path +import viser import viser.transforms as vtf +from typing_extensions import Literal + from nerfstudio.data.scene_box import OrientedBox from nerfstudio.models.base_model import Model from nerfstudio.models.splatfacto import SplatfactoModel from nerfstudio.viewer.control_panel import ControlPanel -from typing_extensions import Literal - -import viser def populate_export_tab( diff --git a/tests/utils/test_math.py b/tests/utils/test_math.py index 952c6a1be..821d19084 100644 --- a/tests/utils/test_math.py +++ b/tests/utils/test_math.py @@ -1,7 +1,9 @@ +import numpy as np import pytest import torch +from scipy.spatial.transform import Rotation as ScR -from nerfstudio.utils.math import components_from_spherical_harmonics +from nerfstudio.utils.math import components_from_spherical_harmonics, rotate_spherical_harmonics @pytest.mark.parametrize("components", list(range(1, 5 + 1))) @@ -14,3 +16,30 @@ def test_spherical_harmonics(components): sh = components_from_spherical_harmonics(components, dx) matrix = (sh.T @ sh) / N * 4 * torch.pi torch.testing.assert_close(matrix, torch.eye(components**2), 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) + + # Norm Preservation + rot_matrix = ScR.random().as_matrix() + rotated_coeffs = rotate_spherical_harmonics(coeffs, rot_matrix) + 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-6) + + # 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: + torch.testing.assert_close(rot_matrix @ coeffs[..., 1:4], rotated_coeffs[..., 1:4], rtol=0, atol=1e-6) + + # Identity Rotation + rot_matrix = np.eye(3) + rotated_coeffs = rotate_spherical_harmonics(coeffs, rot_matrix) + torch.testing.assert_close(coeffs, rotated_coeffs, rtol=0, atol=1e-6) \ No newline at end of file