diff --git a/README.md b/README.md index 871fd6acd..b68024abe 100644 --- a/README.md +++ b/README.md @@ -208,7 +208,7 @@ We support logging to weights and biases. To enable wandb logging, add the flag ## 5. Rendering a trajectory during inference ``` -ns-eval render-trajectory --load-config=outputs/blender_lego/instant_ngp/{MOST_RECENT_RUN}/config.yml--traj=spiral --output-path=output.mp4 +ns-render --load-config=outputs/blender_lego/instant_ngp/{MOST_RECENT_RUN}/config.yml--traj=spiral --output-path=output.mp4 ``` ## 6. In-depth guide diff --git a/docs/quickstart/first_nerf.md b/docs/quickstart/first_nerf.md index 41f6abe6f..4e9f645a7 100644 --- a/docs/quickstart/first_nerf.md +++ b/docs/quickstart/first_nerf.md @@ -47,7 +47,7 @@ ns-train nerfacto --vis viewer --viewer.zmq-port 8001 --viewer.websocket-port 80 ## Visualizing training runs -If you using a fast NeRF variant (ie. Instant-NGP), we reccomend using our viewer. See our [viewer docs](viewer_quickstart.md) for more details. The viewer will allow interactive visualization of training in realtime. +If you are using a fast NeRF variant (ie. Instant-NGP), we recommend using our viewer. See our [viewer docs](viewer_quickstart.md) for more details. The viewer will allow interactive visualization of training in realtime. Additionally, if you run everything with the default configuration, by default, we use [TensorBoard](https://www.tensorflow.org/tensorboard) to log all training curves, test images, and other stats. Once the job is launched, you will be able to track training by launching the tensorboard in `outputs/blender_lego/vanilla_nerf//`. @@ -60,7 +60,7 @@ tensorboard --logdir outputs To evaluate the trained NeRF, we provide an evaluation script that allows you to do benchmarking (see our [benchmarking workflow](../developer_guides/benchmarking.md)) or to render out the scene with a custom trajectory and save the output to a video. ```bash -ns-eval render-trajectory --load-config=outputs/blender_lego/instant_ngp/2022-07-07_230905/config.yml --traj=spiral --output-path=output.mp4 +ns-render --load-config=outputs/blender_lego/instant_ngp/2022-07-07_230905/config.yml --traj=spiral --output-path=output.mp4 ``` Please note, this quickstart allows you to preform everything in a headless manner. We also provide a web-based viewer that allows you to easily monitor training or render out trajectories. See our [viewer docs](viewer_quickstart.md) for more. diff --git a/nerfstudio/utils/eval_utils.py b/nerfstudio/utils/eval_utils.py new file mode 100644 index 000000000..d6aec65d8 --- /dev/null +++ b/nerfstudio/utils/eval_utils.py @@ -0,0 +1,93 @@ +# 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. + +""" +Evaluation utils +""" +from __future__ import annotations + +import os +import sys +from pathlib import Path +from typing import Tuple + +import torch +import yaml +from rich.console import Console + +from nerfstudio.configs import base_config as cfg +from nerfstudio.pipelines.base_pipeline import Pipeline + +console = Console(width=120) + + +def eval_load_checkpoint(config: cfg.TrainerConfig, pipeline: Pipeline) -> Path: + ## TODO: ideally eventually want to get this to be the same as whatever is used to load train checkpoint too + """Helper function to load checkpointed pipeline + + Args: + config (DictConfig): Configuration of pipeline to load + pipeline (Pipeline): Pipeline instance of which to load weights + """ + assert config.load_dir is not None + if config.load_step is None: + console.print("Loading latest checkpoint from load_dir") + # NOTE: this is specific to the checkpoint name format + if not os.path.exists(config.load_dir): + console.rule("Error", style="red") + console.print(f"No checkpoint directory found at {config.load_dir}, ", justify="center") + console.print( + "Please make sure the checkpoint exists, they should be generated periodically during training", + justify="center", + ) + sys.exit(1) + load_step = sorted(int(x[x.find("-") + 1 : x.find(".")]) for x in os.listdir(config.load_dir))[-1] + else: + load_step = config.load_step + load_path = config.load_dir / f"step-{load_step:09d}.ckpt" + assert load_path.exists(), f"Checkpoint {load_path} does not exist" + loaded_state = torch.load(load_path, map_location="cpu") + pipeline.load_pipeline(loaded_state["pipeline"]) + console.print(f":white_check_mark: Done loading checkpoint from {load_path}") + return load_path + + +def eval_setup(config_path: Path) -> Tuple[cfg.Config, Pipeline, Path]: + """Shared setup for loading a saved pipeline for evaluation. + + Args: + config_path: Path to config YAML file. + + Returns: + Loaded config, pipeline module, and corresponding checkpoint. + """ + # load save config + config = yaml.load(config_path.read_text(), Loader=yaml.Loader) + assert isinstance(config, cfg.Config) + + # load checkpoints from wherever they were saved + # TODO: expose the ability to choose an arbitrary checkpoint + config.trainer.load_dir = config.get_checkpoint_dir() + config.pipeline.datamanager.eval_image_indices = None + + # setup pipeline (which includes the DataManager) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + pipeline = config.pipeline.setup(device=device, test_mode=True) + assert isinstance(pipeline, Pipeline) + pipeline.eval() + + # load checkpointed information + checkpoint_path = eval_load_checkpoint(config.trainer, pipeline) + + return config, pipeline, checkpoint_path diff --git a/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPanel.jsx b/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPanel.jsx index 4f50d3cf0..adb3ee15e 100644 --- a/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPanel.jsx +++ b/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPanel.jsx @@ -468,7 +468,7 @@ export default function CameraPanel(props) { // Copy the text inside the text field const config_filename = `${config_base_dir}/config.yml`; const camera_path_filename = `${config_base_dir}/camera_path.json`; - const cmd = `python scripts/eval.py render-trajectory --load-config ${config_filename} --traj filename --camera-path-filename ${camera_path_filename} --output-path output.mp4`; + const cmd = `ns-render --load-config ${config_filename} --traj filename --camera-path-filename ${camera_path_filename} --output-path output.mp4`; navigator.clipboard.writeText(cmd); const camera_path_payload = { diff --git a/pyproject.toml b/pyproject.toml index 6ef497c47..cc138dcc9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -80,6 +80,7 @@ ns-process-data = "scripts.process_data:entrypoint" ns-download-data = "scripts.downloads.download_data:entrypoint" ns-train = "scripts.train:entrypoint" ns-eval = "scripts.eval:entrypoint" +ns-render = "scripts.render:entrypoint" ns-dev-test = "scripts.debugging.run_actions:entrypoint" [options] diff --git a/scripts/completions/install.py b/scripts/completions/install.py index 91a595603..2caf88ef1 100755 --- a/scripts/completions/install.py +++ b/scripts/completions/install.py @@ -28,6 +28,7 @@ ENTRYPOINTS = [ "ns-download-data", "ns-train", "ns-eval", + "ns-render", "ns-dev-test", ] diff --git a/scripts/eval.py b/scripts/eval.py index 148fa408b..33b0e7540 100755 --- a/scripts/eval.py +++ b/scripts/eval.py @@ -4,149 +4,21 @@ eval.py """ from __future__ import annotations -import dataclasses import json import logging -import os -import sys from dataclasses import dataclass from pathlib import Path -from typing import Literal, Tuple, Union import dcargs -import mediapy as media -import torch -import yaml from rich.console import Console -from rich.progress import ( - BarColumn, - Progress, - TaskProgressColumn, - TextColumn, - TimeRemainingColumn, -) -from typing_extensions import assert_never -from nerfstudio.cameras.camera_paths import get_path_from_json, get_spiral_path -from nerfstudio.cameras.cameras import Cameras -from nerfstudio.configs import base_config as cfg -from nerfstudio.pipelines.base_pipeline import Pipeline -from nerfstudio.utils.rich_utils import ItersPerSecColumn +from nerfstudio.utils.eval_utils import eval_setup console = Console(width=120) logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=logging.INFO) -def _load_checkpoint(config: cfg.TrainerConfig, pipeline: Pipeline) -> Path: - """Helper function to load checkpointed pipeline - - Args: - config (DictConfig): Configuration of pipeline to load - pipeline (Pipeline): Pipeline instance of which to load weights - """ - assert config.load_dir is not None - if config.load_step is None: - console.print("Loading latest checkpoint from load_dir") - # NOTE: this is specific to the checkpoint name format - if not os.path.exists(config.load_dir): - console.rule("Error", style="red") - console.print(f"No checkpoint directory found at {config.load_dir}, ", justify="center") - console.print( - "Please make sure the checkpoint exists, they should be generated periodically during training", - justify="center", - ) - sys.exit(1) - load_step = sorted(int(x[x.find("-") + 1 : x.find(".")]) for x in os.listdir(config.load_dir))[-1] - else: - load_step = config.load_step - load_path = config.load_dir / f"step-{load_step:09d}.ckpt" - assert load_path.exists(), f"Checkpoint {load_path} does not exist" - loaded_state = torch.load(load_path, map_location="cpu") - pipeline.load_pipeline(loaded_state["pipeline"]) - console.print(f":white_check_mark: Done loading checkpoint from {load_path}") - return load_path - - -def _render_trajectory_video( - pipeline: Pipeline, - cameras: Cameras, - output_filename: Path, - rendered_output_name: str, - rendered_resolution_scaling_factor: float = 1.0, - seconds: float = 5.0, -) -> None: - """Helper function to create a video of the spiral trajectory. - - Args: - pipeline: Pipeline to evaluate with. - cameras: Cameras to render. - output_filename: Name of the output file. - rendered_output_name: Name of the renderer output to use. - rendered_resolution_scaling_factor: Scaling factor to apply to the camera image resolution. - seconds: Number for the output video. - """ - console.print("[bold green]Creating trajectory video") - images = [] - cameras.rescale_output_resolution(rendered_resolution_scaling_factor) - - progress = Progress( - TextColumn(":movie_camera: Rendering :movie_camera:"), - BarColumn(), - TaskProgressColumn(show_speed=True), - ItersPerSecColumn(suffix="fps"), - TimeRemainingColumn(elapsed_when_finished=True, compact=True), - ) - with progress: - for camera_idx in progress.track(range(cameras.size), description=""): - camera_ray_bundle = cameras.generate_rays(camera_indices=camera_idx).to(pipeline.device) - with torch.no_grad(): - outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) - if rendered_output_name not in outputs: - console.rule("Error", style="red") - console.print(f"Could not find {rendered_output_name} in the model outputs", justify="center") - console.print(f"Please set --rendered_output_name to one of: {outputs.keys()}", justify="center") - sys.exit(1) - image = outputs[rendered_output_name].cpu().numpy() - images.append(image) - - fps = len(images) / seconds - with console.status("[yellow]Saving video", spinner="bouncingBall"): - media.write_video(output_filename, images, fps=fps) - console.rule("[green] :tada: :tada: :tada: Success :tada: :tada: :tada:") - console.print(f"[green]Saved video to {output_filename}", justify="center") - - -def _eval_setup(config_path: Path) -> Tuple[cfg.Config, Pipeline, Path]: - """Shared setup for loading a saved pipeline for evaluation. - - Args: - config_path: Path to config YAML file. - - Returns: - Loaded config, pipeline module, and corresponding checkpoint. - """ - # load save config - config = yaml.load(config_path.read_text(), Loader=yaml.Loader) - assert isinstance(config, cfg.Config) - - # load checkpoints from wherever they were saved - # TODO: expose the ability to choose an arbitrary checkpoint - config.trainer.load_dir = config.get_checkpoint_dir() - config.pipeline.datamanager.eval_image_indices = None - - # setup pipeline (which includes the DataManager) - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - pipeline = config.pipeline.setup(device=device, test_mode=True) - assert isinstance(pipeline, Pipeline) - pipeline.eval() - - # load checkpointed information - checkpoint_path = _load_checkpoint(config.trainer, pipeline) - - return config, pipeline, checkpoint_path - - @dataclass class ComputePSNR: """Load a checkpoint, compute some PSNR metrics, and save to a JSON.""" @@ -158,7 +30,7 @@ class ComputePSNR: def main(self) -> None: """Main function.""" - config, pipeline, checkpoint_path = _eval_setup(self.load_config) + config, pipeline, checkpoint_path = eval_setup(self.load_config) assert self.output_path.suffix == ".json" metrics_dict = pipeline.get_average_eval_image_metrics() @@ -174,71 +46,10 @@ class ComputePSNR: console.print(f"Saved results to: {self.output_path}") -@dataclasses.dataclass -class RenderTrajectory: - """Load a checkpoint, render a trajectory, and save to a video file.""" - - # Path to config YAML file. - load_config: Path - # Name of the renderer output to use. rgb, depth, etc. - rendered_output_name: str = "rgb" - # Trajectory to render. - traj: Literal["spiral", "interp", "filename"] = "spiral" - # Scaling factor to apply to the camera image resolution. - downscale_factor: int = 1 - # Filename of the camera path to render. - camera_path_filename: Path = Path("camera_path.json") - # Name of the output file. - output_path: Path = Path("output.mp4") - # How long the video should be. - seconds: float = 5.0 - - def main(self) -> None: - """Main function.""" - _, pipeline, _ = _eval_setup(self.load_config) - - seconds = self.seconds - - # TODO(ethan): use camera information from parsing args - if self.traj == "spiral": - camera_start = pipeline.datamanager.eval_dataloader.get_camera(image_idx=0) - # TODO(ethan): pass in the up direction of the camera - camera_path = get_spiral_path(camera_start, steps=30, radius=0.1) - elif self.traj == "interp": - # cameras_a = pipeline.datamanager.eval_dataloader.get_camera(image_idx=0) - # cameras_b = pipeline.datamanager.eval_dataloader.get_camera(image_idx=10) - # camera_path = get_interpolated_camera_path(cameras, steps=30) - raise NotImplementedError("Interpolated camera path not implemented.") - elif self.traj == "filename": - with open(self.camera_path_filename, "r", encoding="utf-8") as f: - camera_path = json.load(f) - seconds = camera_path["seconds"] - camera_path = get_path_from_json(camera_path) - else: - assert_never(self.traj) - - _render_trajectory_video( - pipeline, - camera_path, - output_filename=self.output_path, - rendered_output_name=self.rendered_output_name, - rendered_resolution_scaling_factor=1.0 / self.downscale_factor, - seconds=seconds, - ) - - def entrypoint(): """Entrypoint for use with pyproject scripts.""" - # A Union over dataclass types will create a subcommand for each type. - # - # TODO: it would make sense to split this script up into separate scripts. - # - To reduce duplicate code, some "upstream" refactor could also simplify the - # shared `_load_checkpoint()` and `_eval_setup()` helpers. The high-level - # operations implemented by each seem fairly universal; ideally the checkpoint - # loading logic, for example, would be the same as what's used for loading a - # checkpoint when resuming a training run. dcargs.extras.set_accent_color("bright_yellow") - dcargs.cli(Union[ComputePSNR, RenderTrajectory]).main() + dcargs.cli(ComputePSNR).main() if __name__ == "__main__": diff --git a/scripts/render.py b/scripts/render.py new file mode 100644 index 000000000..8da733855 --- /dev/null +++ b/scripts/render.py @@ -0,0 +1,148 @@ +#!/usr/bin/env python +""" +render.py +""" +from __future__ import annotations + +import dataclasses +import json +import logging +import sys +from pathlib import Path +from typing import Literal + +import dcargs +import mediapy as media +import torch +from rich.console import Console +from rich.progress import ( + BarColumn, + Progress, + TaskProgressColumn, + TextColumn, + TimeRemainingColumn, +) +from typing_extensions import assert_never + +from nerfstudio.cameras.camera_paths import get_path_from_json, get_spiral_path +from nerfstudio.cameras.cameras import Cameras +from nerfstudio.configs.base_config import Config # pylint: disable=unused-import +from nerfstudio.pipelines.base_pipeline import Pipeline +from nerfstudio.utils.eval_utils import eval_setup +from nerfstudio.utils.rich_utils import ItersPerSecColumn + +console = Console(width=120) + +logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=logging.INFO) + + +def _render_trajectory_video( + pipeline: Pipeline, + cameras: Cameras, + output_filename: Path, + rendered_output_name: str, + rendered_resolution_scaling_factor: float = 1.0, + seconds: float = 5.0, +) -> None: + """Helper function to create a video of the spiral trajectory. + + Args: + pipeline: Pipeline to evaluate with. + cameras: Cameras to render. + output_filename: Name of the output file. + rendered_output_name: Name of the renderer output to use. + rendered_resolution_scaling_factor: Scaling factor to apply to the camera image resolution. + seconds: Number for the output video. + """ + console.print("[bold green]Creating trajectory video") + images = [] + cameras.rescale_output_resolution(rendered_resolution_scaling_factor) + + progress = Progress( + TextColumn(":movie_camera: Rendering :movie_camera:"), + BarColumn(), + TaskProgressColumn(show_speed=True), + ItersPerSecColumn(suffix="fps"), + TimeRemainingColumn(elapsed_when_finished=True, compact=True), + ) + with progress: + for camera_idx in progress.track(range(cameras.size), description=""): + camera_ray_bundle = cameras.generate_rays(camera_indices=camera_idx).to(pipeline.device) + with torch.no_grad(): + outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) + if rendered_output_name not in outputs: + console.rule("Error", style="red") + console.print(f"Could not find {rendered_output_name} in the model outputs", justify="center") + console.print(f"Please set --rendered_output_name to one of: {outputs.keys()}", justify="center") + sys.exit(1) + image = outputs[rendered_output_name].cpu().numpy() + images.append(image) + + fps = len(images) / seconds + with console.status("[yellow]Saving video", spinner="bouncingBall"): + media.write_video(output_filename, images, fps=fps) + console.rule("[green] :tada: :tada: :tada: Success :tada: :tada: :tada:") + console.print(f"[green]Saved video to {output_filename}", justify="center") + + +@dataclasses.dataclass +class RenderTrajectory: + """Load a checkpoint, render a trajectory, and save to a video file.""" + + # Path to config YAML file. + load_config: Path + # Name of the renderer output to use. rgb, depth, etc. + rendered_output_name: str = "rgb" + # Trajectory to render. + traj: Literal["spiral", "interp", "filename"] = "spiral" + # Scaling factor to apply to the camera image resolution. + downscale_factor: int = 1 + # Filename of the camera path to render. + camera_path_filename: Path = Path("camera_path.json") + # Name of the output file. + output_path: Path = Path("output.mp4") + # How long the video should be. + seconds: float = 5.0 + + def main(self) -> None: + """Main function.""" + _, pipeline, _ = eval_setup(self.load_config) + + seconds = self.seconds + + # TODO(ethan): use camera information from parsing args + if self.traj == "spiral": + camera_start = pipeline.datamanager.eval_dataloader.get_camera(image_idx=0) + # TODO(ethan): pass in the up direction of the camera + camera_path = get_spiral_path(camera_start, steps=30, radius=0.1) + elif self.traj == "interp": + # cameras_a = pipeline.datamanager.eval_dataloader.get_camera(image_idx=0) + # cameras_b = pipeline.datamanager.eval_dataloader.get_camera(image_idx=10) + # camera_path = get_interpolated_camera_path(cameras, steps=30) + raise NotImplementedError("Interpolated camera path not implemented.") + elif self.traj == "filename": + with open(self.camera_path_filename, "r", encoding="utf-8") as f: + camera_path = json.load(f) + seconds = camera_path["seconds"] + camera_path = get_path_from_json(camera_path) + else: + assert_never(self.traj) + + _render_trajectory_video( + pipeline, + camera_path, + output_filename=self.output_path, + rendered_output_name=self.rendered_output_name, + rendered_resolution_scaling_factor=1.0 / self.downscale_factor, + seconds=seconds, + ) + + +def entrypoint(): + """Entrypoint for use with pyproject scripts.""" + dcargs.extras.set_accent_color("bright_yellow") + dcargs.cli(RenderTrajectory).main() + + +if __name__ == "__main__": + entrypoint()