diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 6de1da68b..11a1e8b88 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -30,7 +30,7 @@ from nerfstudio.data.datamanagers.base_datamanager import ( VanillaDataManager, VanillaDataManagerConfig, ) -from nerfstudio.data.datamanagers.dreamfusion_datamanager import DreamFusionDataManagerConfig +from nerfstudio.data.datamanagers.random_cameras_datamanager import RandomCamerasDataManagerConfig from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig from nerfstudio.data.dataparsers.dycheck_dataparser import DycheckDataParserConfig @@ -458,7 +458,7 @@ method_configs["dreamfusion"] = TrainerConfig( max_num_iterations=30000, mixed_precision=True, pipeline=VanillaPipelineConfig( - datamanager=DreamFusionDataManagerConfig( + datamanager=RandomCamerasDataManagerConfig( horizontal_rotation_warmup=2000, ), model=DreamFusionModelConfig( diff --git a/nerfstudio/data/datamanagers/dreamfusion_datamanager.py b/nerfstudio/data/datamanagers/random_cameras_datamanager.py similarity index 96% rename from nerfstudio/data/datamanagers/dreamfusion_datamanager.py rename to nerfstudio/data/datamanagers/random_cameras_datamanager.py index 9a9091443..ba667b75f 100644 --- a/nerfstudio/data/datamanagers/dreamfusion_datamanager.py +++ b/nerfstudio/data/datamanagers/random_cameras_datamanager.py @@ -13,7 +13,7 @@ # limitations under the License. """ -Data manager for dreamfusion +Data manager without input images, only random camera poses """ from __future__ import annotations @@ -150,10 +150,10 @@ def random_train_pose( @dataclass -class DreamFusionDataManagerConfig(DataManagerConfig): +class RandomCamerasDataManagerConfig(DataManagerConfig): """Configuration for data manager that does not load from a dataset. Instead, it generates random poses.""" - _target: Type = field(default_factory=lambda: DreamFusionDataManager) + _target: Type = field(default_factory=lambda: RandomCamerasDataManager) train_resolution: int = 64 """Training resolution""" eval_resolution: int = 64 @@ -180,7 +180,7 @@ class DreamFusionDataManagerConfig(DataManagerConfig): """How many steps until the full horizontal rotation range is used""" -class DreamFusionDataManager(DataManager): # pylint: disable=abstract-method +class RandomCamerasDataManager(DataManager): # pylint: disable=abstract-method """Basic stored data manager implementation. This is pretty much a port over from our old dataloading utilities, and is a little jank @@ -193,12 +193,12 @@ class DreamFusionDataManager(DataManager): # pylint: disable=abstract-method config: the DataManagerConfig used to instantiate class """ - config: DreamFusionDataManagerConfig + config: RandomCamerasDataManagerConfig # pylint: disable=super-init-not-called def __init__( self, - config: DreamFusionDataManagerConfig, + config: RandomCamerasDataManagerConfig, device: Union[torch.device, str] = "cpu", test_mode: Literal["test", "val", "inference"] = "val", world_size: int = 1, @@ -215,7 +215,7 @@ class DreamFusionDataManager(DataManager): # pylint: disable=abstract-method self.input_data = False if self.config.data is not None: - CONSOLE.print("[red] --data should not be used with the DreamFusionDataManager[/red]") + CONSOLE.print("[red] --data should not be used with the RandomCamerasDataManager[/red]") sys.exit(1) print(type(config)) diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index 653929d94..b19cb6c02 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -38,6 +38,7 @@ from nerfstudio.engine.callbacks import ( TrainingCallbackAttributes, TrainingCallbackLocation, ) +from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager from nerfstudio.engine.optimizers import Optimizers from nerfstudio.pipelines.base_pipeline import VanillaPipeline from nerfstudio.utils import profiler, writer @@ -231,6 +232,7 @@ class Trainer: assert self.pipeline.datamanager.train_dataset is not None, "Missing DatsetInputs" if self.pipeline.datamanager.input_data: + assert isinstance(self.pipeline.datamanager, VanillaDataManager) self.pipeline.datamanager.train_dataparser_outputs.save_dataparser_transform( self.base_dir / "dataparser_transforms.json" ) diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 847e40841..e9d30a9ab 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -41,6 +41,7 @@ from nerfstudio.configs import base_config as cfg from nerfstudio.data.datamanagers.base_datamanager import ( DataManager, DataManagerConfig, + VanillaDataManager, ) from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes from nerfstudio.models.base_model import Model, ModelConfig @@ -346,6 +347,7 @@ class VanillaPipeline(Pipeline): """ self.eval() metrics_dict_list = [] + assert isinstance(self.datamanager, VanillaDataManager) num_images = len(self.datamanager.fixed_indices_eval_dataloader) with Progress( TextColumn("[progress.description]{task.description}"),