Support exporting point clouds (#924)

* Support exporting point clouds

* Fix typing

* lint

* Update requirements

nerfstudio/model_components/exporters.py

@@ -0,0 +1,128 @@
+# Copyright 2022 The Nerfstudio Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+"""Algorithms for exporting data."""
+
+from __future__ import annotations
+
+import sys
+from typing import Tuple
+
+import open3d as o3d
+import torch
+from rich.console import Console
+from rich.progress import (
+    BarColumn,
+    Progress,
+    TaskProgressColumn,
+    TextColumn,
+    TimeRemainingColumn,
+)
+
+from nerfstudio.configs.base_config import Config  # pylint: disable=unused-import
+from nerfstudio.pipelines.base_pipeline import Pipeline
+
+CONSOLE = Console(width=120)
+
+
+def generate_point_cloud(
+    pipeline: Pipeline,
+    num_points: int = 1000000,
+    remove_outliers: bool = True,
+    estimate_normals: bool = False,
+    rgb_output_name: str = "rgb",
+    depth_output_name: str = "depth",
+    use_bounding_box: bool = True,
+    bounding_box_min: Tuple[float, float, float] = (-1.0, -1.0, -1.0),
+    bounding_box_max: Tuple[float, float, float] = (1.0, 1.0, 1.0),
+) -> o3d.geometry.PointCloud:
+    """Generate a point cloud from a nerf.
+
+    Args:
+        pipeline: Pipeline to evaluate with.
+        num_points: Number of points to generate. May result in less if outlier removal is used.
+        remove_outliers: Whether to remove outliers.
+        estimate_normals: Whether to estimate normals.
+        rgb_output_name: Name of the RGB output.
+        depth_output_name: Name of the depth output.
+        use_bounding_box: Whether to use a bounding box to sample points.
+        bounding_box_min: Minimum of the bounding box.
+        bounding_box_max: Maximum of the bounding box.
+
+    Returns:
+        Point cloud.
+    """
+
+    progress = Progress(
+        TextColumn(":cloud: Computing Point Cloud :cloud:"),
+        BarColumn(),
+        TaskProgressColumn(show_speed=True),
+        TimeRemainingColumn(elapsed_when_finished=True, compact=True),
+    )
+    points = []
+    rgbs = []
+    with progress as progress_bar:
+        task = progress_bar.add_task("Generating Point Cloud", total=num_points)
+        while not progress_bar.finished:
+            with torch.no_grad():
+                ray_bundle, _ = pipeline.datamanager.next_train(0)
+                outputs = pipeline.model(ray_bundle)
+            if rgb_output_name not in outputs:
+                CONSOLE.rule("Error", style="red")
+                CONSOLE.print(f"Could not find {rgb_output_name} in the model outputs", justify="center")
+                CONSOLE.print(f"Please set --rgb_output_name to one of: {outputs.keys()}", justify="center")
+                sys.exit(1)
+            if depth_output_name not in outputs:
+                CONSOLE.rule("Error", style="red")
+                CONSOLE.print(f"Could not find {depth_output_name} in the model outputs", justify="center")
+                CONSOLE.print(f"Please set --depth_output_name to one of: {outputs.keys()}", justify="center")
+                sys.exit(1)
+            rgb = outputs[rgb_output_name]
+            depth = outputs[depth_output_name]
+            point = ray_bundle.origins + ray_bundle.directions * depth
+
+            if use_bounding_box:
+                comp_l = torch.tensor(bounding_box_min, device=point.device)
+                comp_m = torch.tensor(bounding_box_max, device=point.device)
+                assert torch.all(
+                    comp_l < comp_m
+                ), f"Bounding box min {bounding_box_min} must be smaller than max {bounding_box_max}"
+                mask = torch.all(torch.concat([point > comp_l, point < comp_m], dim=-1), dim=-1)
+                point = point[mask]
+                rgb = rgb[mask]
+
+            points.append(point)
+            rgbs.append(rgb)
+            progress.advance(task, point.shape[0])
+    points = torch.cat(points, dim=0)
+    rgbs = torch.cat(rgbs, dim=0)
+
+    pcd = o3d.geometry.PointCloud()
+    pcd.points = o3d.utility.Vector3dVector(points.float().cpu().numpy())
+    pcd.colors = o3d.utility.Vector3dVector(rgbs.float().cpu().numpy())
+
+    if remove_outliers:
+        CONSOLE.print("Cleaning Point Cloud")
+        pcd, _ = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=1.0)
+        print("\033[A\033[A")
+        CONSOLE.print("[bold green]:white_check_mark: Cleaning Point Cloud")
+
+    if estimate_normals:
+        CONSOLE.print("Estimating Point Cloud Normals")
+        pcd.estimate_normals()
+        print("\033[A\033[A")
+        CONSOLE.print("[bold green]:white_check_mark: Estimating Point Cloud Normals")
+
+    return pcd

pyproject.toml

@@ -31,6 +31,7 @@ dependencies = [
     "msgpack==1.0.4",
     "msgpack_numpy==0.4.8",
     "nerfacc==0.2.1",
+    "open3d>=0.16.0",
     "opencv-python==4.6.0.66",
     "plotly==5.7.0",
     "protobuf==3.20.0",
@@ -66,7 +67,7 @@ docs = [
     "furo==2022.4.7",
     "readthedocs-sphinx-search==0.1.2",
     "myst-nb==0.16.0",
-    "nbformat==5.4.0",
+    "nbformat==5.5.0",
     "sphinx==4.5.0",
     "sphinxemoji==0.2.0",
     "sphinx-argparse==0.3.1",
@@ -83,6 +84,7 @@ ns-download-data = "scripts.downloads.download_data:entrypoint"
 ns-train = "scripts.train:entrypoint"
 ns-eval = "scripts.eval:entrypoint"
 ns-render = "scripts.render:entrypoint"
+ns-export = "scripts.exporter:entrypoint"
 ns-dev-test = "scripts.github.run_actions:entrypoint"
 ns-bridge-server = "nerfstudio.viewer.server.server:entrypoint"
 

scripts/exporter.py

@@ -0,0 +1,127 @@
+"""
+Script for exporting NeRF into other formats.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Tuple, Union
+
+import open3d as o3d
+import tyro
+from rich.console import Console
+from typing_extensions import Annotated
+
+from nerfstudio.model_components.exporters import generate_point_cloud
+from nerfstudio.utils.eval_utils import eval_setup
+
+CONSOLE = Console(width=120)
+
+
+@dataclass
+class Exporter:
+    """Export the mesh from a YML config to a folder."""
+
+    load_config: Path
+    """Path to the config YAML file."""
+    output_dir: Path
+    """Path to the output directory."""
+
+
+@dataclass
+class ExportPointCloud(Exporter):
+    """Class to export meshes with the poisson algorithm."""
+
+    num_points: int = 1000000
+    """Number of points to generate. May result in less if outlier removal is used."""
+    remove_outliers: bool = True
+    """Remove outliers from the point cloud."""
+    estimate_normals: bool = False
+    """Estimate normals for the point cloud."""
+    depth_output_name: str = "depth"
+    """Name of the depth output."""
+    rgb_output_name: str = "rgb"
+    """Name of the RGB output."""
+    use_bounding_box: bool = True
+    """Only query points within the bounding box"""
+    bounding_box_min: Tuple[float, float, float] = (-1, -1, -1)
+    """Minimum of the bounding box, used if use_bounding_box is True."""
+    bounding_box_max: Tuple[float, float, float] = (1, 1, 1)
+    """Minimum of the bounding box, used if use_bounding_box is True."""
+    num_rays_per_batch: int = 65536
+    """Number of rays to evaluate per batch. Decrease if you run out of memory."""
+
+    def main(self) -> None:
+        """Export point cloud."""
+
+        if not self.output_dir.exists():
+            self.output_dir.mkdir(parents=True)
+
+        _, pipeline, _ = eval_setup(self.load_config)
+
+        # Increase the batchsize to speed up the evaluation.
+        pipeline.datamanager.train_pixel_sampler.num_rays_per_batch = self.num_rays_per_batch
+
+        pcd = generate_point_cloud(
+            pipeline=pipeline,
+            num_points=self.num_points,
+            remove_outliers=self.remove_outliers,
+            estimate_normals=self.estimate_normals,
+            rgb_output_name=self.rgb_output_name,
+            depth_output_name=self.depth_output_name,
+            use_bounding_box=self.use_bounding_box,
+            bounding_box_min=self.bounding_box_min,
+            bounding_box_max=self.bounding_box_max,
+        )
+
+        CONSOLE.print(f"[bold green]:white_check_mark: Generated {pcd}")
+        CONSOLE.print("Saving Point Cloud")
+        o3d.io.write_point_cloud(str(self.output_dir / "point_cloud.ply"), pcd)
+        print("\033[A\033[A")
+        CONSOLE.print("[bold green]:white_check_mark: Saving Point Cloud")
+
+
+@dataclass
+class ExportTSDFMesh(Exporter):
+    """Class to export meshes with the poisson algorithm."""
+
+    def main(self) -> None:
+        """Export mesh"""
+        raise NotImplementedError("TSDF mesh reconstruction not implemented yet.")
+
+
+@dataclass
+class ExportPoissonMesh(Exporter):
+    """Exports a mesh using poisson mesh reconstruction."""
+
+    def main(self) -> None:
+        """Export mesh"""
+        raise NotImplementedError("Poisson mesh reconstruction not implemented yet.")
+
+
+@dataclass
+class ExportMarchingCubesMesh(Exporter):
+    """Exports a mesh using marching cubes."""
+
+    def main(self) -> None:
+        """Export mesh"""
+        raise NotImplementedError("Marching cubes not implemented yet.")
+
+
+Commands = 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")],
+]
+
+
+def entrypoint():
+    """Entrypoint for use with pyproject scripts."""
+    tyro.extras.set_accent_color("bright_yellow")
+    tyro.cli(Commands).main()
+
+
+if __name__ == "__main__":
+    entrypoint()