From 246cb618ee31283a12dff2f9a858dfc74083ea5c Mon Sep 17 00:00:00 2001 From: Jonas Kulhanek Date: Fri, 22 Sep 2023 21:27:36 +0200 Subject: [PATCH 1/3] ns-render - render all images --- nerfstudio/scripts/render.py | 208 ++++++++++++++++++++++++++++++++- nerfstudio/utils/eval_utils.py | 7 +- 2 files changed, 211 insertions(+), 4 deletions(-) diff --git a/nerfstudio/scripts/render.py b/nerfstudio/scripts/render.py index 99a7cae01d..c09ffcf17b 100644 --- a/nerfstudio/scripts/render.py +++ b/nerfstudio/scripts/render.py @@ -23,7 +23,8 @@ import struct import shutil import sys -from contextlib import ExitStack +import gzip +from contextlib import ExitStack, contextmanager from dataclasses import dataclass, field from pathlib import Path from typing import Any, Dict, List, Literal, Optional, Union @@ -52,8 +53,11 @@ get_path_from_json, get_spiral_path, ) -from nerfstudio.cameras.cameras import Cameras, CameraType -from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager +from nerfstudio.cameras.cameras import Cameras, CameraType, RayBundle +from nerfstudio.data.datasets.base_dataset import Dataset +from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager, VanillaDataManagerConfig +from nerfstudio.data.utils.dataloaders import FixedIndicesEvalDataloader +from nerfstudio.engine.trainer import TrainerConfig from nerfstudio.data.scene_box import OrientedBox from nerfstudio.model_components import renderers from nerfstudio.pipelines.base_pipeline import Pipeline @@ -585,11 +589,209 @@ def main(self) -> None: ) +@contextmanager +def _disable_datamanager_setup(cls): + """ + Disables setup_train or setup_eval for faster initialization. + """ + old_setup_train = getattr(cls, "setup_train") + old_setup_eval = getattr(cls, "setup_eval") + setattr(cls, "setup_train", lambda *args, **kwargs: None) + setattr(cls, "setup_eval", lambda *args, **kwargs: None) + yield cls + setattr(cls, "setup_train", old_setup_train) + setattr(cls, "setup_eval", old_setup_eval) + + +@dataclass +class DatasetRender(BaseRender): + """Render all images in the dataset.""" + + output_path: Path = Path("renders") + """Path to output video file.""" + data: Optional[Path] = None + """Override path to the dataset.""" + downscale_factor: Optional[float] = None + """Scaling factor to apply to the camera image resolution.""" + split: Literal["train", "val", "test", "train+test"] = "test" + """Split to render.""" + rendered_output_names: Optional[List[str]] = field(default_factory=lambda: None) + """Name of the renderer outputs to use. rgb, depth, raw-depth, gt-rgb etc. By default all outputs are rendered.""" + + def main(self): + config: TrainerConfig + + def update_config(config: TrainerConfig) -> TrainerConfig: + data_manager_config = config.pipeline.datamanager + assert isinstance(data_manager_config, VanillaDataManagerConfig) + data_manager_config.eval_image_indices = None + data_manager_config.eval_num_images_to_sample_from = -1 + data_manager_config.eval_num_times_to_repeat_images = -1 + data_manager_config.train_num_images_to_sample_from = -1 + data_manager_config.train_num_times_to_repeat_images = -1 + data_manager_config.data = self.data + if self.downscale_factor is not None: + assert hasattr(data_manager_config.dataparser, "downscale_factor") + setattr(data_manager_config.dataparser, "downscale_factor", self.downscale_factor) + return config + + config, pipeline, _, _ = eval_setup( + self.load_config, + eval_num_rays_per_chunk=self.eval_num_rays_per_chunk, + test_mode="inference", + update_config_callback=update_config, + ) + data_manager_config = config.pipeline.datamanager + assert isinstance(data_manager_config, VanillaDataManagerConfig) + + for split in self.split.split("+"): + datamanager: VanillaDataManager + dataset: Dataset + if split == "train": + with _disable_datamanager_setup(data_manager_config._target): # pylint: disable=protected-access + datamanager = data_manager_config.setup(test_mode="test", device=pipeline.device) + + dataset = datamanager.train_dataset + dataparser_outputs = getattr(dataset, "_dataparser_outputs", datamanager.train_dataparser_outputs) + else: + with _disable_datamanager_setup(data_manager_config._target): # pylint: disable=protected-access + datamanager = data_manager_config.setup(test_mode=split, device=pipeline.device) + + dataset = datamanager.eval_dataset + dataparser_outputs = getattr(dataset, "_dataparser_outputs", None) + if dataparser_outputs is None: + dataparser_outputs = datamanager.dataparser.get_dataparser_outputs(split=datamanager.test_split) + dataloader = FixedIndicesEvalDataloader( + input_dataset=dataset, + device=datamanager.device, + num_workers=datamanager.world_size * 4, + ) + images_root = Path(os.path.commonprefix(dataparser_outputs.image_filenames)) + with Progress( + TextColumn(f":movie_camera: Rendering split {split} :movie_camera:"), + BarColumn(), + TaskProgressColumn( + text_format="[progress.percentage]{task.completed}/{task.total:>.0f}({task.percentage:>3.1f}%)", + show_speed=True, + ), + ItersPerSecColumn(suffix="fps"), + TimeRemainingColumn(elapsed_when_finished=False, compact=False), + TimeElapsedColumn(), + ) as progress: + for camera_idx, (ray_bundle, batch) in enumerate(progress.track(dataloader, total=len(dataset))): + ray_bundle: RayBundle + with torch.no_grad(): + outputs = pipeline.model.get_outputs_for_camera_ray_bundle(ray_bundle) + + gt_batch = batch.copy() + gt_batch["rgb"] = gt_batch.pop("image") + all_outputs = ( + list(outputs.keys()) + + [f"raw-{x}" for x in outputs.keys()] + + [f"gt-{x}" for x in gt_batch.keys()] + + [f"raw-gt-{x}" for x in gt_batch.keys()] + ) + rendered_output_names = self.rendered_output_names + if rendered_output_names is None: + rendered_output_names = ["gt-rgb"] + list(outputs.keys()) + for rendered_output_name in rendered_output_names: + if rendered_output_name not in all_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: {all_outputs}", justify="center" + ) + sys.exit(1) + + is_raw = False + is_depth = rendered_output_name.find("depth") != -1 + image_name = f"{camera_idx:05d}" + + # Try to get the original filename + image_name = ( + dataparser_outputs.image_filenames[camera_idx].with_suffix("").relative_to(images_root) + ) + + output_path = self.output_path / split / rendered_output_name / image_name + output_path.parent.mkdir(exist_ok=True, parents=True) + + output_name = rendered_output_name + if output_name.startswith("raw-"): + output_name = output_name[4:] + is_raw = True + if output_name.startswith("gt-"): + output_name = output_name[3:] + output_image = gt_batch[output_name] + else: + output_image = outputs[output_name] + if is_depth: + # Divide by the dataparser scale factor + output_image.div_(dataparser_outputs.dataparser_scale) + else: + if output_name.startswith("gt-"): + output_name = output_name[3:] + output_image = gt_batch[output_name] + else: + output_image = outputs[output_name] + del output_name + + # Map to color spaces / numpy + if is_raw: + output_image = output_image.cpu().numpy() + elif is_depth: + output_image = ( + colormaps.apply_depth_colormap( + output_image, + accumulation=outputs["accumulation"], + near_plane=self.depth_near_plane, + far_plane=self.depth_far_plane, + colormap_options=self.colormap_options, + ) + .cpu() + .numpy() + ) + else: + output_image = ( + colormaps.apply_colormap( + image=output_image, + colormap_options=self.colormap_options, + ) + .cpu() + .numpy() + ) + + # Save to file + if is_raw: + with gzip.open(output_path.with_suffix(".npy.gz"), "wb") as f: + np.save(f, output_image) + if self.image_format == "png": + media.write_image(output_path.with_suffix(".png"), output_image, fmt="png") + elif self.image_format == "jpeg": + media.write_image( + output_path.with_suffix(".jpg"), output_image, fmt="jpeg", quality=self.jpeg_quality + ) + else: + raise ValueError(f"Unknown image format {self.image_format}") + + table = Table( + title=None, + show_header=False, + box=box.MINIMAL, + title_style=style.Style(bold=True), + ) + for split in self.split.split("+"): + table.add_row(f"Outputs {split}", str(self.output_path / split)) + CONSOLE.print(Panel(table, title="[bold][green]:tada: Render on split {} Complete :tada:[/bold]", expand=False)) + + Commands = tyro.conf.FlagConversionOff[ Union[ Annotated[RenderCameraPath, tyro.conf.subcommand(name="camera-path")], Annotated[RenderInterpolated, tyro.conf.subcommand(name="interpolate")], Annotated[SpiralRender, tyro.conf.subcommand(name="spiral")], + Annotated[DatasetRender, tyro.conf.subcommand(name="dataset")], ] ] diff --git a/nerfstudio/utils/eval_utils.py b/nerfstudio/utils/eval_utils.py index 13948678bb..7e04005368 100644 --- a/nerfstudio/utils/eval_utils.py +++ b/nerfstudio/utils/eval_utils.py @@ -20,7 +20,7 @@ import os import sys from pathlib import Path -from typing import Literal, Optional, Tuple +from typing import Literal, Optional, Tuple, Callable import torch import yaml @@ -69,6 +69,7 @@ def eval_setup( config_path: Path, eval_num_rays_per_chunk: Optional[int] = None, test_mode: Literal["test", "val", "inference"] = "test", + update_config_callback: Optional[Callable[[TrainerConfig], TrainerConfig]] = None, ) -> Tuple[TrainerConfig, Pipeline, Path, int]: """Shared setup for loading a saved pipeline for evaluation. @@ -79,6 +80,7 @@ def eval_setup( 'val': loads train/val datasets into memory 'test': loads train/test dataset into memory 'inference': does not load any dataset into memory + update_config_callback: Callback to update the config before loading the pipeline Returns: @@ -92,6 +94,9 @@ def eval_setup( if eval_num_rays_per_chunk: config.pipeline.model.eval_num_rays_per_chunk = eval_num_rays_per_chunk + if update_config_callback is not None: + config = update_config_callback(config) + # load checkpoints from wherever they were saved # TODO: expose the ability to choose an arbitrary checkpoint config.load_dir = config.get_checkpoint_dir() From c7c81de40b7004aeeafc12029b3dd273d715318d Mon Sep 17 00:00:00 2001 From: Jonas Kulhanek Date: Tue, 26 Sep 2023 11:33:31 +0200 Subject: [PATCH 2/3] Fix commonprefix->commonpath --- nerfstudio/scripts/render.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nerfstudio/scripts/render.py b/nerfstudio/scripts/render.py index c09ffcf17b..6c9f06a44b 100644 --- a/nerfstudio/scripts/render.py +++ b/nerfstudio/scripts/render.py @@ -666,7 +666,7 @@ def update_config(config: TrainerConfig) -> TrainerConfig: device=datamanager.device, num_workers=datamanager.world_size * 4, ) - images_root = Path(os.path.commonprefix(dataparser_outputs.image_filenames)) + images_root = Path(os.path.commonpath(dataparser_outputs.image_filenames)) with Progress( TextColumn(f":movie_camera: Rendering split {split} :movie_camera:"), BarColumn(), From 5a9ad703253abaa7ae21f3a4c9ea5a6fa8801e30 Mon Sep 17 00:00:00 2001 From: Jonas Kulhanek Date: Tue, 26 Sep 2023 13:03:02 +0200 Subject: [PATCH 3/3] Fix raw rendering --- nerfstudio/scripts/render.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nerfstudio/scripts/render.py b/nerfstudio/scripts/render.py index 6c9f06a44b..ecb7961cc4 100644 --- a/nerfstudio/scripts/render.py +++ b/nerfstudio/scripts/render.py @@ -766,7 +766,7 @@ def update_config(config: TrainerConfig) -> TrainerConfig: if is_raw: with gzip.open(output_path.with_suffix(".npy.gz"), "wb") as f: np.save(f, output_image) - if self.image_format == "png": + elif self.image_format == "png": media.write_image(output_path.with_suffix(".png"), output_image, fmt="png") elif self.image_format == "jpeg": media.write_image(