From a5eea5f14dabba87e80135b73a8cbf41b988106f Mon Sep 17 00:00:00 2001 From: Martin Piala <25415255+mackopes@users.noreply.github.com> Date: Sun, 1 Jan 2023 06:38:43 +0100 Subject: [PATCH] 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 | 12 ++++- nerfstudio/exporter/exporter_utils.py | 59 +++++++++++++++++++++++- scripts/exporter.py | 34 +++++++++++++- 3 files changed, 101 insertions(+), 4 deletions(-) diff --git a/nerfstudio/data/datasets/base_dataset.py b/nerfstudio/data/datasets/base_dataset.py index 44793fbb0..c01d26739 100644 --- a/nerfstudio/data/datasets/base_dataset.py +++ b/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 diff --git a/nerfstudio/exporter/exporter_utils.py b/nerfstudio/exporter/exporter_utils.py index 3e5f4425e..eb6c5f9b6 100644 --- a/nerfstudio/exporter/exporter_utils.py +++ b/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 diff --git a/scripts/exporter.py b/scripts/exporter.py index 175cee187..a555de98c 100644 --- a/scripts/exporter.py +++ b/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")], ] ]