diff --git a/nerfstudio/model_components/exporters.py b/nerfstudio/model_components/exporters.py new file mode 100644 index 000000000..947254500 --- /dev/null +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index c08d82dbe..58955d64e 100644 --- a/pyproject.toml +++ b/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" diff --git a/scripts/exporter.py b/scripts/exporter.py new file mode 100644 index 000000000..55a8a1b28 --- /dev/null +++ b/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()