From a24faf66bf53ecd4f6ca4be0d8073c19a7758098 Mon Sep 17 00:00:00 2001 From: Brent Yi Date: Thu, 22 Feb 2024 23:44:51 -0800 Subject: [PATCH] Export splats in "world frame" by default --- nerfstudio/cameras/camera_utils.py | 4 +-- nerfstudio/scripts/exporter.py | 54 ++++++++++++++++++++++++++---- nerfstudio/viewer/export_panel.py | 11 +++++- 3 files changed, 60 insertions(+), 9 deletions(-) diff --git a/nerfstudio/cameras/camera_utils.py b/nerfstudio/cameras/camera_utils.py index 9c1ee0220..3f66bb045 100644 --- a/nerfstudio/cameras/camera_utils.py +++ b/nerfstudio/cameras/camera_utils.py @@ -311,8 +311,8 @@ def viewmatrix(lookat: torch.Tensor, up: torch.Tensor, pos: torch.Tensor) -> Flo """ vec2 = normalize(lookat) vec1_avg = normalize(up) - vec0 = normalize(torch.cross(vec1_avg, vec2)) - vec1 = normalize(torch.cross(vec2, vec0)) + vec0 = normalize(torch.linalg.cross(vec1_avg, vec2)) + vec1 = normalize(torch.linalg.cross(vec2, vec0)) m = torch.stack([vec0, vec1, vec2, pos], 1) return m diff --git a/nerfstudio/scripts/exporter.py b/nerfstudio/scripts/exporter.py index e52626889..fdfa1a4ec 100644 --- a/nerfstudio/scripts/exporter.py +++ b/nerfstudio/scripts/exporter.py @@ -30,10 +30,10 @@ import numpy as np import open3d as o3d import torch import tyro -from typing_extensions import Annotated, Literal - +from nerfstudio.cameras.camera_utils import quaternion_from_matrix from nerfstudio.cameras.rays import RayBundle from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager +from nerfstudio.data.datamanagers.full_images_datamanager import FullImageDatamanager from nerfstudio.data.datamanagers.parallel_datamanager import ParallelDataManager from nerfstudio.data.scene_box import OrientedBox from nerfstudio.exporter import texture_utils, tsdf_utils @@ -44,6 +44,7 @@ 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.rich_utils import CONSOLE +from typing_extensions import Annotated, Literal @dataclass @@ -121,7 +122,7 @@ class ExportPointCloud(Exporter): """Number of rays to evaluate per batch. Decrease if you run out of memory.""" std_ratio: float = 10.0 """Threshold based on STD of the average distances across the point cloud to remove outliers.""" - save_world_frame: bool = False + save_world_frame: bool = True """If set, saves the point cloud in the same frame as the original dataset. Otherwise, uses the scaled and reoriented coordinate space expected by the NeRF models.""" @@ -482,6 +483,11 @@ class ExportGaussianSplat(Exporter): Export 3D Gaussian Splatting model to a .ply """ + save_world_frame: bool = True + """If set, saves the splat in the same frame as the original dataset. + Otherwise, uses the scaled and reoriented coordinate space produced + internally by Nerfstudio.""" + def main(self) -> None: if not self.output_dir.exists(): self.output_dir.mkdir(parents=True) @@ -497,7 +503,26 @@ class ExportGaussianSplat(Exporter): map_to_tensors = {} with torch.no_grad(): - positions = model.means.cpu().numpy() + if self.save_world_frame: + assert isinstance(pipeline.datamanager, FullImageDatamanager) + dataparser_outputs = pipeline.datamanager.train_dataparser_outputs + dataparser_scale = dataparser_outputs.dataparser_scale + dataparser_transform = dataparser_outputs.dataparser_transform.numpy(force=True) + + 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]) + + positions = ( + np.einsum("ij,bj->bi", output_transform[:3, :3], model.means.cpu().numpy() * output_scale) + + output_transform[None, :3, 3] + ) n = positions.shape[0] map_to_tensors["positions"] = positions map_to_tensors["normals"] = np.zeros_like(positions, dtype=np.float32) @@ -518,11 +543,28 @@ class ExportGaussianSplat(Exporter): map_to_tensors["opacity"] = model.opacities.data.cpu().numpy() - scales = model.scales.data.cpu().numpy() + # Note that scales are in log space! + scales = model.scales.data.cpu().numpy() + np.log(output_scale) for i in range(3): map_to_tensors[f"scale_{i}"] = scales[:, i, None] - quats = model.quats.data.cpu().numpy() + # Invert and aapply the dataparser transform. + 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/viewer/export_panel.py b/nerfstudio/viewer/export_panel.py index 76808a272..d52494c11 100644 --- a/nerfstudio/viewer/export_panel.py +++ b/nerfstudio/viewer/export_panel.py @@ -97,7 +97,7 @@ def populate_point_cloud_tab( num_points = server.add_gui_number("# Points", initial_value=1_000_000, min=1, max=None, step=1) world_frame = server.add_gui_checkbox( "Save in world frame", - False, + True, hint=( "If checked, saves the point cloud in the same frame as the original dataset. Otherwise, uses the " "scaled and reoriented coordinate space expected by the NeRF models." @@ -195,6 +195,14 @@ def populate_splat_tab( output_directory = server.add_gui_text("Output Directory", initial_value="exports/splat/") generate_command = server.add_gui_button("Generate Command", icon=viser.Icon.TERMINAL_2) + world_frame = server.add_gui_checkbox( + "Save in world frame", + True, + hint=( + "If checked, saves the splat file in the same frame as the original dataset. Otherwise, uses the " + "scaled and reoriented coordinate space generated by Nerfstudio." + ), + ) @generate_command.on_click def _(event: viser.GuiEvent) -> None: @@ -204,6 +212,7 @@ def populate_splat_tab( "ns-export gaussian-splat", f"--load-config {config_path}", f"--output-dir {output_directory.value}", + f"--save-world-frame {world_frame.value}", ] ) show_command_modal(event.client, "splat", command)