diff --git a/nerfstudio/models/base_model.py b/nerfstudio/models/base_model.py index ef36f53c8..e882603c3 100644 --- a/nerfstudio/models/base_model.py +++ b/nerfstudio/models/base_model.py @@ -169,6 +169,7 @@ class Model(nn.Module): num_rays = len(camera_ray_bundle) outputs_lists = defaultdict(list) for i in range(0, num_rays, num_rays_per_chunk): + print(self.config.eval_num_rays_per_chunk) start_idx = i end_idx = i + num_rays_per_chunk ray_bundle = camera_ray_bundle.get_row_major_sliced_ray_bundle(start_idx, end_idx) diff --git a/nerfstudio/utils/eval_utils.py b/nerfstudio/utils/eval_utils.py index d6aec65d8..5e3cf4c3d 100644 --- a/nerfstudio/utils/eval_utils.py +++ b/nerfstudio/utils/eval_utils.py @@ -20,7 +20,7 @@ from __future__ import annotations import os import sys from pathlib import Path -from typing import Tuple +from typing import Optional, Tuple import torch import yaml @@ -63,7 +63,9 @@ def eval_load_checkpoint(config: cfg.TrainerConfig, pipeline: Pipeline) -> Path: return load_path -def eval_setup(config_path: Path) -> Tuple[cfg.Config, Pipeline, Path]: +def eval_setup( + config_path: Path, double_nerfacto_nerf_samples: bool = False, eval_num_rays_per_chunk: Optional[int] = None +) -> Tuple[cfg.Config, Pipeline, Path]: """Shared setup for loading a saved pipeline for evaluation. Args: @@ -76,6 +78,17 @@ def eval_setup(config_path: Path) -> Tuple[cfg.Config, Pipeline, Path]: config = yaml.load(config_path.read_text(), Loader=yaml.Loader) assert isinstance(config, cfg.Config) + if eval_num_rays_per_chunk: + config.pipeline.model.eval_num_rays_per_chunk = eval_num_rays_per_chunk + + if config.method_name == "nerfacto" and double_nerfacto_nerf_samples: + config.pipeline.model.num_nerf_samples_per_ray *= 2 + console.print( + "[bold yellow]Warning: doubling the number of nerfacto nerf samples to " + f"{config.pipeline.model.num_nerf_samples_per_ray} for improved quality." + ) + console.print("[bold yellow]This logic is a hack and will be removed in the future.") + # load checkpoints from wherever they were saved # TODO: expose the ability to choose an arbitrary checkpoint config.trainer.load_dir = config.get_checkpoint_dir() diff --git a/scripts/render.py b/scripts/render.py index 5e6ea9d72..91dfb00eb 100644 --- a/scripts/render.py +++ b/scripts/render.py @@ -9,7 +9,7 @@ import json import logging import sys from pathlib import Path -from typing import Literal +from typing import Literal, Optional import dcargs import mediapy as media @@ -103,10 +103,18 @@ class RenderTrajectory: output_path: Path = Path("output.mp4") # How long the video should be. seconds: float = 5.0 + # A hack to double the number of samples for the nerfacto method. + double_nerfacto_nerf_samples: bool = True + # Specifies number of rays per chunk during eval. + eval_num_rays_per_chunk: Optional[int] = None def main(self) -> None: """Main function.""" - _, pipeline, _ = eval_setup(self.load_config) + _, pipeline, _ = eval_setup( + self.load_config, + double_nerfacto_nerf_samples=self.double_nerfacto_nerf_samples, + eval_num_rays_per_chunk=self.eval_num_rays_per_chunk, + ) seconds = self.seconds