Remove dataparser dependency (#1529)

nerfstudio/configs/experiment_config.py

@@ -72,7 +72,7 @@ class ExperimentConfig(InstantiateConfig):
     vis: Literal["viewer", "wandb", "tensorboard", "viewer+wandb", "viewer+tensorboard"] = "wandb"
     """Which visualizer to use."""
     data: Optional[Path] = None
-    """Alias for --pipeline.datamanager.dataparser.data"""
+    """Alias for --pipeline.datamanager.data"""
     relative_model_dir: Path = Path("nerfstudio_models/")
     """Relative path to save all checkpoints."""
 
@@ -96,9 +96,12 @@ class ExperimentConfig(InstantiateConfig):
     def set_experiment_name(self) -> None:
         """Dynamically set the experiment name"""
         if self.experiment_name is None:
-            datapath = self.pipeline.datamanager.dataparser.data
-            datapath = datapath.parent if datapath.is_file() else datapath
-            self.experiment_name = str(datapath).replace("../", "").replace("/", "-")
+            datapath = self.pipeline.datamanager.data
+            if datapath is not None:
+                datapath = datapath.parent if datapath.is_file() else datapath
+                self.experiment_name = str(datapath.stem)
+            else:
+                self.experiment_name = "unnamed"
 
     def get_base_dir(self) -> Path:
         """Retrieve the base directory to set relative paths"""

nerfstudio/data/datamanagers/base_datamanager.py

@@ -20,6 +20,7 @@ from __future__ import annotations
 
 from abc import abstractmethod
 from dataclasses import dataclass, field
+from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
 import torch
@@ -95,6 +96,21 @@ AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[  # Omit prefixes of
 same as the vanilla union, but results in shorter subcommand names."""
 
 
+@dataclass
+class DataManagerConfig(InstantiateConfig):
+    """Configuration for data manager instantiation; DataManager is in charge of keeping the train/eval dataparsers;
+    After instantiation, data manager holds both train/eval datasets and is in charge of returning unpacked
+    train/eval data at each iteration
+    """
+
+    _target: Type = field(default_factory=lambda: DataManager)
+    """Target class to instantiate."""
+    data: Optional[Path] = None
+    """Source of data, may not be used by all models."""
+    camera_optimizer: Optional[CameraOptimizerConfig] = None
+    """Specifies the camera pose optimizer used during training. Helpful if poses are noisy."""
+
+
 class DataManager(nn.Module):
     """Generic data manager's abstract class
 
@@ -213,12 +229,10 @@ class DataManager(nn.Module):
         """Sets up the data manager for training.
 
         Here you will define any subclass specific object attributes from the attribute"""
-        raise NotImplementedError
 
     @abstractmethod
     def setup_eval(self):
         """Sets up the data manager for evaluation"""
-        raise NotImplementedError
 
     @abstractmethod
     def next_train(self, step: int) -> Tuple:
@@ -241,6 +255,20 @@ class DataManager(nn.Module):
         """Returns the next eval image."""
         raise NotImplementedError
 
+    @abstractmethod
+    def get_train_rays_per_batch(self) -> int:
+        """Returns the number of rays per batch for training."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_eval_rays_per_batch(self) -> int:
+        """Returns the number of rays per batch for evaluation."""
+        raise NotImplementedError
+
+    def get_datapath(self) -> Optional[Path]:  # pylint:disable=no-self-use
+        """Returns the path to the data. This is used to determine where to save camera paths."""
+        return None
+
     def get_training_callbacks(  # pylint:disable=no-self-use
         self, training_callback_attributes: TrainingCallbackAttributes  # pylint: disable=unused-argument
     ) -> List[TrainingCallback]:
@@ -258,11 +286,8 @@ class DataManager(nn.Module):
 
 
 @dataclass
-class VanillaDataManagerConfig(InstantiateConfig):
-    """Configuration for data manager instantiation; DataManager is in charge of keeping the train/eval dataparsers;
-    After instantiation, data manager holds both train/eval datasets and is in charge of returning unpacked
-    train/eval data at each iteration
-    """
+class VanillaDataManagerConfig(DataManagerConfig):
+    """A basic data manager"""
 
     _target: Type = field(default_factory=lambda: VanillaDataManager)
     """Target class to instantiate."""
@@ -333,7 +358,12 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
         self.sampler = None
         self.test_mode = test_mode
         self.test_split = "test" if test_mode in ["test", "inference"] else "val"
-        self.dataparser = self.config.dataparser.setup()
+        self.dataparser = self.config.dataparser
+        if self.config.data is not None:
+            self.config.dataparser.data = Path(self.config.data)
+        else:
+            self.config.data = self.config.dataparser.data
+        self.dataparser = self.dataparser.setup()
         self.train_dataparser_outputs = self.dataparser.get_dataparser_outputs(split="train")
 
         self.train_dataset = self.create_train_dataset()
@@ -455,6 +485,15 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
             return image_idx, camera_ray_bundle, batch
         raise ValueError("No more eval images")
 
+    def get_train_rays_per_batch(self) -> int:
+        return self.config.train_num_rays_per_batch
+
+    def get_eval_rays_per_batch(self) -> int:
+        return self.config.eval_num_rays_per_batch
+
+    def get_datapath(self) -> Path:
+        return self.config.dataparser.data
+
     def get_param_groups(self) -> Dict[str, List[Parameter]]:  # pylint: disable=no-self-use
         """Get the param groups for the data manager.
         Returns:

nerfstudio/engine/trainer.py

@@ -126,20 +126,8 @@ class Trainer:
         # directory to save checkpoints
         self.checkpoint_dir: Path = config.get_checkpoint_dir()
         CONSOLE.log(f"Saving checkpoints to: {self.checkpoint_dir}")
-        # set up viewer if enabled
-        viewer_log_path: Path = self.base_dir / config.viewer.relative_log_filename
-        self.viewer_state, banner_messages = None, None
-        if self.config.is_viewer_enabled() and local_rank == 0:
-            self.viewer_state, banner_messages = viewer_utils.setup_viewer(
-                config.viewer, log_filename=viewer_log_path, datapath=config.pipeline.datamanager.dataparser.data
-            )
-        self._check_viewer_warnings()
-        # set up writers/profilers if enabled
-        writer_log_path: Path = self.base_dir / config.logging.relative_log_dir
-        writer.setup_event_writer(config.is_wandb_enabled(), config.is_tensorboard_enabled(), log_dir=writer_log_path)
-        writer.setup_local_writer(config.logging, max_iter=config.max_num_iterations, banner_messages=banner_messages)
-        writer.put_config(name="config", config_dict=dataclasses.asdict(config), step=0)
-        profiler.setup_profiler(config.logging)
+
+        self.viewer_state = None
 
     def setup(self, test_mode: Literal["test", "val", "inference"] = "val") -> None:
         """Setup the Trainer by calling other setup functions.
@@ -165,6 +153,28 @@ class Trainer:
             )
         )
 
+        # set up viewer if enabled
+        viewer_log_path = self.base_dir / self.config.viewer.relative_log_filename
+        self.viewer_state, banner_messages = None, None
+        if self.config.is_viewer_enabled() and self.local_rank == 0:
+            datapath = self.pipeline.datamanager.get_datapath()
+            if datapath is None:
+                datapath = self.base_dir
+            self.viewer_state, banner_messages = viewer_utils.setup_viewer(
+                self.config.viewer, log_filename=viewer_log_path, datapath=datapath
+            )
+        self._check_viewer_warnings()
+        # set up writers/profilers if enabled
+        writer_log_path = self.base_dir / self.config.logging.relative_log_dir
+        writer.setup_event_writer(
+            self.config.is_wandb_enabled(), self.config.is_tensorboard_enabled(), log_dir=writer_log_path
+        )
+        writer.setup_local_writer(
+            self.config.logging, max_iter=self.config.max_num_iterations, banner_messages=banner_messages
+        )
+        writer.put_config(name="config", config_dict=dataclasses.asdict(self.config), step=0)
+        profiler.setup_profiler(self.config.logging)
+
     def setup_optimizers(self) -> Optimizers:
         """Helper to set up the optimizers
 
@@ -172,9 +182,9 @@ class Trainer:
             The optimizers object given the trainer config.
         """
         optimizer_config = self.config.optimizers.copy()
-        camera_optimizer_config = self.config.pipeline.datamanager.camera_optimizer
         param_groups = self.pipeline.get_param_groups()
-        if camera_optimizer_config.mode != "off":
+        camera_optimizer_config = self.config.pipeline.datamanager.camera_optimizer
+        if camera_optimizer_config is not None and camera_optimizer_config.mode != "off":
             assert camera_optimizer_config.param_group not in optimizer_config
             optimizer_config[camera_optimizer_config.param_group] = {
                 "optimizer": camera_optimizer_config.optimizer,
@@ -215,7 +225,7 @@ class Trainer:
                 if step > 1:
                     writer.put_time(
                         name=EventName.TRAIN_RAYS_PER_SEC,
-                        duration=self.config.pipeline.datamanager.train_num_rays_per_batch / train_t.duration,
+                        duration=self.pipeline.datamanager.get_train_rays_per_batch() / train_t.duration,
                         step=step,
                         avg_over_steps=True,
                     )
@@ -249,7 +259,7 @@ class Trainer:
 
     @check_main_thread
     def _always_render(self, step: int) -> None:
-        if self.config.is_viewer_enabled():
+        if self.viewer_state is not None:
             while True:
                 self.viewer_state.vis["renderingState/isTraining"].write(False)
                 self._update_viewer_state(step)
@@ -289,7 +299,7 @@ class Trainer:
         """
         assert self.viewer_state is not None
         with TimeWriter(writer, EventName.ITER_VIS_TIME, step=step) as _:
-            num_rays_per_batch: int = self.config.pipeline.datamanager.train_num_rays_per_batch
+            num_rays_per_batch: int = self.pipeline.datamanager.get_train_rays_per_batch()
             try:
                 self.viewer_state.update_scene(self, step, self.pipeline.model, num_rays_per_batch)
             except RuntimeError:
@@ -308,7 +318,7 @@ class Trainer:
             vis_t: timer object carrying time to execute visualization step
             step: current step
         """
-        train_num_rays_per_batch: int = self.config.pipeline.datamanager.train_num_rays_per_batch
+        train_num_rays_per_batch: int = self.pipeline.datamanager.get_train_rays_per_batch()
         writer.put_time(
             name=EventName.TRAIN_RAYS_PER_SEC,
             duration=train_num_rays_per_batch / (train_t.duration - vis_t.duration),

nerfstudio/pipelines/base_pipeline.py

@@ -40,6 +40,7 @@ from typing_extensions import Literal
 from nerfstudio.configs import base_config as cfg
 from nerfstudio.data.datamanagers.base_datamanager import (
     DataManager,
+    DataManagerConfig,
     VanillaDataManager,
     VanillaDataManagerConfig,
 )
@@ -193,7 +194,7 @@ class VanillaPipelineConfig(cfg.InstantiateConfig):
 
     _target: Type = field(default_factory=lambda: VanillaPipeline)
     """target class to instantiate"""
-    datamanager: VanillaDataManagerConfig = VanillaDataManagerConfig()
+    datamanager: DataManagerConfig = VanillaDataManagerConfig()
     """specifies the datamanager config"""
     model: ModelConfig = ModelConfig()
     """specifies the model config"""
@@ -264,15 +265,16 @@ class VanillaPipeline(Pipeline):
         model_outputs = self.model(ray_bundle)
         metrics_dict = self.model.get_metrics_dict(model_outputs, batch)
 
-        camera_opt_param_group = self.config.datamanager.camera_optimizer.param_group
-        if camera_opt_param_group in self.datamanager.get_param_groups():
-            # Report the camera optimization metrics
-            metrics_dict["camera_opt_translation"] = (
-                self.datamanager.get_param_groups()[camera_opt_param_group][0].data[:, :3].norm()
-            )
-            metrics_dict["camera_opt_rotation"] = (
-                self.datamanager.get_param_groups()[camera_opt_param_group][0].data[:, 3:].norm()
-            )
+        if self.config.datamanager.camera_optimizer is not None:
+            camera_opt_param_group = self.config.datamanager.camera_optimizer.param_group
+            if camera_opt_param_group in self.datamanager.get_param_groups():
+                # Report the camera optimization metrics
+                metrics_dict["camera_opt_translation"] = (
+                    self.datamanager.get_param_groups()[camera_opt_param_group][0].data[:, :3].norm()
+                )
+                metrics_dict["camera_opt_rotation"] = (
+                    self.datamanager.get_param_groups()[camera_opt_param_group][0].data[:, 3:].norm()
+                )
 
         loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict)
 

scripts/train.py

@@ -219,8 +219,8 @@ def main(config: TrainerConfig) -> None:
 
     config.set_timestamp()
     if config.data:
-        CONSOLE.log("Using --data alias for --data.pipeline.datamanager.dataparser.data")
-        config.pipeline.datamanager.dataparser.data = config.data
+        CONSOLE.log("Using --data alias for --data.pipeline.datamanager.data")
+        config.pipeline.datamanager.data = config.data
 
     if config.load_config:
         CONSOLE.log(f"Loading pre-set config from: {config.load_config}")

scripts/viewer/run_viewer.py

@@ -58,11 +58,11 @@ class RunViewer:
 
         self._start_viewer(config, pipeline)
 
-    def _start_viewer(self, config: TrainerConfig, pipeline):
+    def _start_viewer(self, config: TrainerConfig, pipeline: Pipeline):
         base_dir = config.get_base_dir()
         viewer_log_path = base_dir / config.viewer.relative_log_filename
         viewer_state, banner_messages = viewer_utils.setup_viewer(
-            config.viewer, log_filename=viewer_log_path, datapath=config.pipeline.datamanager.dataparser.data
+            config.viewer, log_filename=viewer_log_path, datapath=pipeline.datamanager.get_datapath()
         )
 
         # We don't need logging, but writer.GLOBAL_BUFFER needs to be populated
@@ -76,16 +76,16 @@ class RunViewer:
         )
         while True:
             viewer_state.vis["renderingState/isTraining"].write(False)
-            self._update_viewer_state(viewer_state, config, pipeline)
+            self._update_viewer_state(viewer_state, pipeline)
 
-    def _update_viewer_state(self, viewer_state: viewer_utils.ViewerState, config: TrainerConfig, pipeline: Pipeline):
+    def _update_viewer_state(self, viewer_state: viewer_utils.ViewerState, pipeline: Pipeline):
         """Updates the viewer state by rendering out scene with current pipeline
         Returns the time taken to render scene.
 
         """
         # NOTE: step must be > 0 otherwise the rendering would not happen
         step = 1
-        num_rays_per_batch = config.pipeline.datamanager.train_num_rays_per_batch
+        num_rays_per_batch = pipeline.datamanager.get_train_rays_per_batch()
         with TimeWriter(writer, EventName.ITER_VIS_TIME) as _:
             try:
                 viewer_state.update_scene(self, step, pipeline.model, num_rays_per_batch)