From 5491df92d9f6a47c86f7e1ecb4e986ae1e1f5632 Mon Sep 17 00:00:00 2001 From: Otto Seiskari Date: Wed, 29 May 2024 17:08:54 +0300 Subject: [PATCH] Export Splatfacto-optimized poses (#3116) * Fix exported poses when Splatfacto pose optimization is enabled Currently works correctly only with --eval-mode=all * Make Splatfacto pose adjustments consistent with Nerfacto Also allow exporting Nerfacto-optimized poses. Fixes #3073. Second revision where the rotation and translation parts are handled separately. --- nerfstudio/cameras/camera_optimizers.py | 21 +++++++++++++++------ nerfstudio/exporter/exporter_utils.py | 23 ++++++++++++++++++++--- nerfstudio/models/splatfacto.py | 6 ++---- 3 files changed, 37 insertions(+), 13 deletions(-) diff --git a/nerfstudio/cameras/camera_optimizers.py b/nerfstudio/cameras/camera_optimizers.py index 9b47b7e53..7b5f4ccbc 100644 --- a/nerfstudio/cameras/camera_optimizers.py +++ b/nerfstudio/cameras/camera_optimizers.py @@ -157,15 +157,24 @@ class CameraOptimizer(nn.Module): if self.config.mode == "off": return camera.camera_to_worlds - assert camera.metadata is not None, "Must provide id of camera in its metadata" - if "cam_idx" not in camera.metadata: - # Evalutaion cams? + if camera.metadata is None or "cam_idx" not in camera.metadata: + # Viser cameras return camera.camera_to_worlds camera_idx = camera.metadata["cam_idx"] - adj = self(torch.tensor([camera_idx], dtype=torch.long, device=camera.device)) # type: ignore - adj = torch.cat([adj, torch.Tensor([0, 0, 0, 1])[None, None].to(adj)], dim=1) - return torch.bmm(camera.camera_to_worlds, adj) + adj = self(torch.tensor([camera_idx], dtype=torch.long)).to(camera.device) # type: ignore + + return torch.cat( + [ + # Apply rotation to directions in world coordinates, without touching the origin. + # Equivalent to: directions -> correction[:3,:3] @ directions + torch.bmm(adj[..., :3, :3], camera.camera_to_worlds[..., :3, :3]), + # Apply translation in world coordinate, independently of rotation. + # Equivalent to: origins -> origins + correction[:3,3] + camera.camera_to_worlds[..., :3, 3:] + adj[..., :3, 3:], + ], + dim=-1, + ) def get_loss_dict(self, loss_dict: dict) -> None: """Add regularization""" diff --git a/nerfstudio/exporter/exporter_utils.py b/nerfstudio/exporter/exporter_utils.py index 672d03968..bad46f517 100644 --- a/nerfstudio/exporter/exporter_utils.py +++ b/nerfstudio/exporter/exporter_utils.py @@ -30,6 +30,7 @@ from jaxtyping import Float from rich.progress import BarColumn, Progress, TaskProgressColumn, TextColumn, TimeRemainingColumn from torch import Tensor +from nerfstudio.cameras.camera_optimizers import CameraOptimizer from nerfstudio.cameras.cameras import Cameras from nerfstudio.cameras.rays import RayBundle from nerfstudio.data.datasets.base_dataset import InputDataset @@ -283,11 +284,14 @@ def render_trajectory( return images, depths -def collect_camera_poses_for_dataset(dataset: Optional[InputDataset]) -> List[Dict[str, Any]]: +def collect_camera_poses_for_dataset( + dataset: Optional[InputDataset], camera_optimizer: Optional[CameraOptimizer] = None +) -> List[Dict[str, Any]]: """Collects rescaled, translated and optimised camera poses for a dataset. Args: dataset: Dataset to collect camera poses for. + camera_optimizer: Camera optimizer that has been used for adjusting the poses Returns: List of dicts containing camera poses. @@ -304,7 +308,15 @@ def collect_camera_poses_for_dataset(dataset: Optional[InputDataset]) -> List[Di # new cameras are in cameras, whereas image paths are stored in a private member of the dataset for idx in range(len(cameras)): image_filename = image_filenames[idx] - transform = cameras.camera_to_worlds[idx].tolist() + if camera_optimizer is None: + transform = cameras.camera_to_worlds[idx].tolist() + else: + # print('exporting optimized camera pose for camera %d' % idx) + camera = cameras[idx : idx + 1] + assert camera.metadata is not None + camera.metadata["cam_idx"] = idx + transform = camera_optimizer.apply_to_camera(camera).tolist()[0] + frames.append( { "file_path": str(image_filename), @@ -331,7 +343,12 @@ def collect_camera_poses(pipeline: VanillaPipeline) -> Tuple[List[Dict[str, Any] eval_dataset = pipeline.datamanager.eval_dataset assert isinstance(eval_dataset, InputDataset) - train_frames = collect_camera_poses_for_dataset(train_dataset) + camera_optimizer = None + if hasattr(pipeline.model, "camera_optimizer"): + camera_optimizer = pipeline.model.camera_optimizer + + train_frames = collect_camera_poses_for_dataset(train_dataset, camera_optimizer) + # Note: returning original poses, even if --eval-mode=all eval_frames = collect_camera_poses_for_dataset(eval_dataset) return train_frames, eval_frames diff --git a/nerfstudio/models/splatfacto.py b/nerfstudio/models/splatfacto.py index e996dd640..9b29eca62 100644 --- a/nerfstudio/models/splatfacto.py +++ b/nerfstudio/models/splatfacto.py @@ -672,10 +672,10 @@ class SplatfactoModel(Model): return {} assert camera.shape[0] == 1, "Only one camera at a time" + optimized_camera_to_world = self.camera_optimizer.apply_to_camera(camera)[0, ...] + # get the background color if self.training: - optimized_camera_to_world = self.camera_optimizer.apply_to_camera(camera)[0, ...] - if self.config.background_color == "random": background = torch.rand(3, device=self.device) elif self.config.background_color == "white": @@ -685,8 +685,6 @@ class SplatfactoModel(Model): else: background = self.background_color.to(self.device) else: - optimized_camera_to_world = camera.camera_to_worlds[0, ...] - if renderers.BACKGROUND_COLOR_OVERRIDE is not None: background = renderers.BACKGROUND_COLOR_OVERRIDE.to(self.device) else: