Add test mode to eval CLI

nerfstudio/data/datamanagers/full_images_datamanager.py

@@ -84,7 +84,7 @@ def __init__(
  self,
  config: FullImageDatamanagerConfig,
  device: Union[torch.device, str] = "cpu",
- test_mode: Literal["test", "val", "inference"] = "val",
+ test_mode: Literal["test", "val", "inference", "train"] = "val",
  world_size: int = 1,
  local_rank: int = 0,
  **kwargs,
@@ -95,7 +95,12 @@ def __init__(
  self.local_rank = local_rank
  self.sampler = None
  self.test_mode = test_mode
- self.test_split = "test" if test_mode in ["test", "inference"] else "val"
+ if test_mode in ["test", "inference"]:
+ self.test_split = "test"
+ elif test_mode == "train":
+ self.test_split = "train"
+ else:
+ self.test_split = "val"
  self.dataparser_config = self.config.dataparser
  if self.config.data is not None:
  self.config.dataparser.data = Path(self.config.data)

nerfstudio/pipelines/base_pipeline.py

@@ -243,7 +243,7 @@ def __init__(
  self,
  config: VanillaPipelineConfig,
  device: str,
- test_mode: Literal["test", "val", "inference"] = "val",
+ test_mode: Literal["test", "val", "inference", "train"] = "val",
  world_size: int = 1,
  local_rank: int = 0,
  grad_scaler: Optional[GradScaler] = None,

nerfstudio/scripts/eval.py

@@ -21,7 +21,7 @@
 import json
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Optional
+from typing import Literal, Optional
 
 import tyro
 
@@ -39,10 +39,12 @@ class ComputePSNR:
  output_path: Path = Path("output.json")
  # Optional path to save rendered outputs to.
  render_output_path: Optional[Path] = None
+ # Split to test the data on
+ test_mode: Optional[Literal["test", "val", "inference", "train"]] = "test"
 
  def main(self) -> None:
  """Main function."""
- config, pipeline, checkpoint_path, _ = eval_setup(self.load_config)
+ config, pipeline, checkpoint_path, _ = eval_setup(self.load_config, test_mode=self.test_mode)
  assert self.output_path.suffix == ".json"
  if self.render_output_path is not None:
  self.render_output_path.mkdir(parents=True, exist_ok=True)

nerfstudio/utils/eval_utils.py

@@ -67,7 +67,7 @@ def eval_load_checkpoint(config: TrainerConfig, pipeline: Pipeline) -> Tuple[Pat
 def eval_setup(
  config_path: Path,
  eval_num_rays_per_chunk: Optional[int] = None,
- test_mode: Literal["test", "val", "inference"] = "test",
+ test_mode: Optional[Literal["test", "val", "inference", "train"]] = "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 +79,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
+ 'train': loads train/train dataset into memory - useful for evaluation on `train` set.
  update_config_callback: Callback to update the config before loading the pipeline