Camera export (#1191)

* implement the camera exporter

* remove train_dataset from exporter

* fix linter errors

* add docstring

* correct comment

* fix linter errors

* make image_filenames a property of the dataset

nerfstudio/data/datasets/base_dataset.py

@@ -18,7 +18,8 @@ Dataset.
 from __future__ import annotations
 
 from copy import deepcopy
-from typing import Dict
+from pathlib import Path
+from typing import Dict, List
 
 import numpy as np
 import numpy.typing as npt
@@ -115,3 +116,12 @@ class InputDataset(Dataset):
     def __getitem__(self, image_idx: int) -> Dict:
         data = self.get_data(image_idx)
         return data
+
+    @property
+    def image_filenames(self) -> List[Path]:
+        """
+        Returns image filenames for this dataset.
+        The order of filenames is the same as in the Cameras object for easy mapping.
+        """
+
+        return self._dataparser_outputs.image_filenames

nerfstudio/exporter/exporter_utils.py

@@ -22,7 +22,7 @@ from __future__ import annotations
 
 import sys
 from dataclasses import dataclass
-from typing import List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 import numpy as np
 import open3d as o3d
@@ -39,7 +39,8 @@ from rich.progress import (
 from torchtyping import TensorType
 
 from nerfstudio.cameras.cameras import Cameras
-from nerfstudio.pipelines.base_pipeline import Pipeline
+from nerfstudio.data.datasets.base_dataset import InputDataset
+from nerfstudio.pipelines.base_pipeline import Pipeline, VanillaPipeline
 from nerfstudio.utils.rich_utils import ItersPerSecColumn
 
 CONSOLE = Console(width=120)
@@ -255,3 +256,57 @@ def render_trajectory(
             images.append(outputs[rgb_output_name].cpu().numpy())
             depths.append(outputs[depth_output_name].cpu().numpy())
     return images, depths
+
+
+def collect_camera_poses_for_dataset(dataset: Optional[InputDataset]) -> List[Dict[str, Any]]:
+    """Collects rescaled, translated and optimised camera poses for a dataset.
+
+    Args:
+        dataset: Dataset to collect camera poses for.
+
+    Returns:
+        List of dicts containing camera poses.
+    """
+
+    if dataset is None:
+        return []
+
+    cameras = dataset.cameras
+    image_filenames = dataset.image_filenames
+
+    frames: List[Dict[str, Any]] = []
+
+    # 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()
+        frames.append(
+            {
+                "file_path": str(image_filename),
+                "transform": transform,
+            }
+        )
+
+    return frames
+
+
+def collect_camera_poses(pipeline: VanillaPipeline) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
+    """Collects camera poses for train and eval datasets.
+
+    Args:
+        pipeline: Pipeline to evaluate with.
+
+    Returns:
+        List of train camera poses, list of eval camera poses.
+    """
+
+    train_dataset = pipeline.datamanager.train_dataset
+    assert isinstance(train_dataset, InputDataset)
+
+    eval_dataset = pipeline.datamanager.eval_dataset
+    assert isinstance(eval_dataset, InputDataset)
+
+    train_frames = collect_camera_poses_for_dataset(train_dataset)
+    eval_frames = collect_camera_poses_for_dataset(eval_dataset)
+
+    return train_frames, eval_frames

scripts/exporter.py

@@ -4,6 +4,8 @@ Script for exporting NeRF into other formats.
 
 from __future__ import annotations
 
+import json
+import os
 import sys
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -19,10 +21,11 @@ from typing_extensions import Annotated, Literal
 from nerfstudio.cameras.rays import RayBundle
 from nerfstudio.exporter import texture_utils, tsdf_utils
 from nerfstudio.exporter.exporter_utils import (
+    collect_camera_poses,
     generate_point_cloud,
     get_mesh_from_filename,
 )
-from nerfstudio.pipelines.base_pipeline import Pipeline
+from nerfstudio.pipelines.base_pipeline import Pipeline, VanillaPipeline
 from nerfstudio.utils.eval_utils import eval_setup
 
 CONSOLE = Console(width=120)
@@ -314,12 +317,41 @@ class ExportMarchingCubesMesh(Exporter):
         raise NotImplementedError("Marching cubes not implemented yet.")
 
 
+@dataclass
+class ExportCameraPoses(Exporter):
+    """
+    Export camera poses to a .json file.
+    """
+
+    def main(self) -> None:
+        """Export camera poses"""
+        if not self.output_dir.exists():
+            self.output_dir.mkdir(parents=True)
+
+        _, pipeline, _ = eval_setup(self.load_config)
+        assert isinstance(pipeline, VanillaPipeline)
+        train_frames, eval_frames = collect_camera_poses(pipeline)
+
+        for file_name, frames in [("transforms_train.json", train_frames), ("transforms_eval.json", eval_frames)]:
+            if len(frames) == 0:
+                CONSOLE.print(f"[bold yellow]No frames found for {file_name}. Skipping.")
+                continue
+
+            output_file_path = os.path.join(self.output_dir, file_name)
+
+            with open(output_file_path, "w", encoding="UTF-8") as f:
+                json.dump(frames, f, indent=4)
+
+            CONSOLE.print(f"[bold green]:white_check_mark: Saved poses to {output_file_path}")
+
+
 Commands = tyro.conf.FlagConversionOff[
     Union[
         Annotated[ExportPointCloud, tyro.conf.subcommand(name="pointcloud")],
         Annotated[ExportTSDFMesh, tyro.conf.subcommand(name="tsdf")],
         Annotated[ExportPoissonMesh, tyro.conf.subcommand(name="poisson")],
         Annotated[ExportMarchingCubesMesh, tyro.conf.subcommand(name="marching-cubes")],
+        Annotated[ExportCameraPoses, tyro.conf.subcommand(name="cameras")],
     ]
 ]