From 77b7bfb1b1670daf885d6c06ae66a841949b6c88 Mon Sep 17 00:00:00 2001 From: Jonas Kulhanek Date: Wed, 1 Mar 2023 11:15:01 +0100 Subject: [PATCH] Implement ForwardRef generic interface --- .../debugging_tools/local_logger.md | 3 +- .../pipelines/datamanagers.md | 4 +- nerfstudio/cameras/camera_optimizers.py | 4 +- nerfstudio/configs/base_config.py | 39 ++++++++++++++----- nerfstudio/configs/experiment_config.py | 5 ++- .../data/datamanagers/base_datamanager.py | 4 +- .../dataparsers/arkitscenes_dataparser.py | 5 +-- .../data/dataparsers/base_dataparser.py | 4 +- .../data/dataparsers/blender_dataparser.py | 3 +- .../data/dataparsers/dnerf_dataparser.py | 5 +-- .../data/dataparsers/dycheck_dataparser.py | 5 +-- .../dataparsers/instant_ngp_dataparser.py | 5 +-- .../data/dataparsers/minimal_dataparser.py | 5 +-- .../data/dataparsers/nerfstudio_dataparser.py | 5 +-- .../data/dataparsers/nuscenes_dataparser.py | 5 +-- .../dataparsers/phototourism_dataparser.py | 5 +-- .../data/dataparsers/scannet_dataparser.py | 5 +-- .../data/dataparsers/sdfstudio_dataparser.py | 5 +-- .../data/dataparsers/sitcoms3d_dataparser.py | 5 +-- nerfstudio/engine/schedulers.py | 2 +- nerfstudio/engine/trainer.py | 4 +- nerfstudio/models/base_model.py | 5 +-- nerfstudio/models/depth_nerfacto.py | 7 ++++ nerfstudio/models/instant_ngp.py | 2 +- nerfstudio/models/mipnerf.py | 14 +++++-- nerfstudio/models/nerfacto.py | 3 +- nerfstudio/models/nerfplayer_nerfacto.py | 5 +++ nerfstudio/models/nerfplayer_ngp.py | 4 ++ nerfstudio/models/semantic_nerfw.py | 4 ++ nerfstudio/models/tensorf.py | 4 +- nerfstudio/models/vanilla_nerf.py | 3 +- nerfstudio/pipelines/base_pipeline.py | 18 ++++----- nerfstudio/pipelines/dynamic_batch.py | 13 ++++++- 33 files changed, 120 insertions(+), 89 deletions(-) diff --git a/docs/developer_guides/debugging_tools/local_logger.md b/docs/developer_guides/debugging_tools/local_logger.md index 92461154d..544b0084e 100644 --- a/docs/developer_guides/debugging_tools/local_logger.md +++ b/docs/developer_guides/debugging_tools/local_logger.md @@ -8,10 +8,9 @@ A skeleton of the local writer config is defined below. """nerfstudio/configs/base_config.py"""" @dataclass -class LocalWriterConfig(InstantiateConfig): +class LocalWriterConfig(InstantiateConfig[writer.LocalWriter]): """Local Writer config""" - _target: Type = writer.LocalWriter enable: bool = False stats_to_track: Tuple[writer.EventName, ...] = ( writer.EventName.ITER_TRAIN_TIME, diff --git a/docs/developer_guides/pipelines/datamanagers.md b/docs/developer_guides/pipelines/datamanagers.md index 32c55b63e..5d4d8c3eb 100644 --- a/docs/developer_guides/pipelines/datamanagers.md +++ b/docs/developer_guides/pipelines/datamanagers.md @@ -44,14 +44,12 @@ We've implemented a VanillaDataManager that implements the standard logic of mos ```python @dataclass -class VanillaDataManagerConfig(InstantiateConfig): +class VanillaDataManagerConfig(InstantiateConfig["VanillaDataManager"]): """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: VanillaDataManager) - """target class to instantiate""" dataparser: AnnotatedDataParserUnion = BlenderDataParserConfig() """specifies the dataparser used to unpack the data""" train_num_rays_per_batch: int = 1024 diff --git a/nerfstudio/cameras/camera_optimizers.py b/nerfstudio/cameras/camera_optimizers.py index afd002a71..67799b6af 100644 --- a/nerfstudio/cameras/camera_optimizers.py +++ b/nerfstudio/cameras/camera_optimizers.py @@ -39,11 +39,9 @@ from nerfstudio.utils import poses as pose_utils @dataclass -class CameraOptimizerConfig(InstantiateConfig): +class CameraOptimizerConfig(InstantiateConfig["CameraOptimizer"]): """Configuration of optimization for camera poses.""" - _target: Type = field(default_factory=lambda: CameraOptimizer) - mode: Literal["off", "SO3xR3", "SE3"] = "off" """Pose optimization strategy to use. If enabled, we recommend SO3xR3.""" diff --git a/nerfstudio/configs/base_config.py b/nerfstudio/configs/base_config.py index a2e695c14..1df2cdd6d 100644 --- a/nerfstudio/configs/base_config.py +++ b/nerfstudio/configs/base_config.py @@ -19,11 +19,11 @@ from __future__ import annotations import warnings -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path -from typing import Any, List, Optional, Tuple, Type +from typing import Any, ForwardRef, Generic, List, Optional, Tuple, Type -from typing_extensions import Literal +from typing_extensions import Literal, TypeVar # model instances from nerfstudio.utils import writer @@ -47,15 +47,37 @@ class PrintableConfig: # pylint: disable=too-few-public-methods return "\n ".join(lines) +T = TypeVar("T") + # Base instantiate configs @dataclass -class InstantiateConfig(PrintableConfig): # pylint: disable=too-few-public-methods +class InstantiateConfig(PrintableConfig, Generic[T]): # pylint: disable=too-few-public-methods """Config class for instantiating an the class specified in the _target attribute.""" - _target: Type + _target: Optional[Type] = field(repr=False, default=None) - def setup(self, **kwargs) -> Any: + def __post_init__(self): + if self._target is None: + for base in getattr(self, "__orig_bases__", []): + if hasattr(base, "__origin__") and issubclass(base.__origin__, InstantiateConfig): + value = base.__args__[0] + if isinstance(value, ForwardRef): + if value.__forward_evaluated__: + value = value.__forward_value__ + elif value.__forward_module__ is None: + value.__forward_module__ = type(self).__module__ + value = getattr(value, "_evaluate")(None, None, set()) + self._target = value + break + else: + raise RuntimeError( + f"Could not find correct type to instantiate for {type(self)}. " + "Make sure to inherit InstantiateConfig[T], where T is the type to instantiate." + ) + + def setup(self, **kwargs) -> T: """Returns the instantiated object using the config.""" + assert self._target is not None return self._target(self, **kwargs) @@ -77,11 +99,9 @@ class MachineConfig(PrintableConfig): @dataclass -class LocalWriterConfig(InstantiateConfig): +class LocalWriterConfig(InstantiateConfig[writer.LocalWriter]): """Local Writer config""" - _target: Type = writer.LocalWriter - """target class to instantiate""" enable: bool = False """if True enables local logging, else disables""" stats_to_track: Tuple[writer.EventName, ...] = ( @@ -102,6 +122,7 @@ class LocalWriterConfig(InstantiateConfig): Args: banner_messages: List of strings that always print at the bottom of screen. """ + assert self._target is not None return self._target(self, banner_messages=banner_messages, **kwargs) diff --git a/nerfstudio/configs/experiment_config.py b/nerfstudio/configs/experiment_config.py index 98703ee40..ad93a511b 100644 --- a/nerfstudio/configs/experiment_config.py +++ b/nerfstudio/configs/experiment_config.py @@ -23,7 +23,7 @@ from typing import Any, Dict, Optional import yaml from rich.console import Console -from typing_extensions import Literal +from typing_extensions import Literal, TypeVar from nerfstudio.configs.base_config import ( InstantiateConfig, @@ -37,10 +37,11 @@ from nerfstudio.engine.schedulers import SchedulerConfig from nerfstudio.pipelines.base_pipeline import VanillaPipelineConfig CONSOLE = Console(width=120) +T = TypeVar("T") @dataclass -class ExperimentConfig(InstantiateConfig): +class ExperimentConfig(InstantiateConfig[T]): """Full config contents for running an experiment. Any experiment types (like training) will be subclassed from this, and must have their _target field defined accordingly.""" diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 9f7e6272a..887d13919 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -258,14 +258,12 @@ class DataManager(nn.Module): @dataclass -class VanillaDataManagerConfig(InstantiateConfig): +class VanillaDataManagerConfig(InstantiateConfig["VanillaDataManager"]): """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: VanillaDataManager) - """Target class to instantiate.""" dataparser: AnnotatedDataParserUnion = BlenderDataParserConfig() """Specifies the dataparser used to unpack the data.""" train_num_rays_per_batch: int = 1024 diff --git a/nerfstudio/data/dataparsers/arkitscenes_dataparser.py b/nerfstudio/data/dataparsers/arkitscenes_dataparser.py index 1143fe003..9e1af356e 100644 --- a/nerfstudio/data/dataparsers/arkitscenes_dataparser.py +++ b/nerfstudio/data/dataparsers/arkitscenes_dataparser.py @@ -23,6 +23,7 @@ import numpy as np import torch from nerfstudio.cameras.cameras import Cameras, CameraType +from nerfstudio.configs.base_config import InstantiateConfig from nerfstudio.data.dataparsers.base_dataparser import ( DataParser, DataParserConfig, @@ -60,14 +61,12 @@ def traj_string_to_matrix(traj_string: str): @dataclass -class ARKitScenesDataParserConfig(DataParserConfig): +class ARKitScenesDataParserConfig(DataParserConfig, InstantiateConfig["ARKitScenes"]): """ARKitScenes dataset config. ARKitScenes dataset (http://github.com/apple/ARKitScenes) is a large-scale 3D dataset of indoor scenes. This dataparser uses 3D deteciton subset of the ARKitScenes dataset. """ - _target: Type = field(default_factory=lambda: ARKitScenes) - """target class to instantiate""" data: Path = Path("data/ARKitScenes/3dod/Validation/41069021") """Path to ARKitScenes folder with densely extracted scenes.""" scene_scale: float = 1.0 diff --git a/nerfstudio/data/dataparsers/base_dataparser.py b/nerfstudio/data/dataparsers/base_dataparser.py index 66df055b4..c42c0a258 100644 --- a/nerfstudio/data/dataparsers/base_dataparser.py +++ b/nerfstudio/data/dataparsers/base_dataparser.py @@ -91,11 +91,9 @@ class DataparserOutputs: @dataclass -class DataParserConfig(cfg.InstantiateConfig): +class DataParserConfig: """Basic dataset config""" - _target: Type = field(default_factory=lambda: DataParser[DataParserConfig]) - """_target: target class to instantiate""" data: Path = Path() """Directory specifying location of data.""" diff --git a/nerfstudio/data/dataparsers/blender_dataparser.py b/nerfstudio/data/dataparsers/blender_dataparser.py index 37ccbe07f..ddd813fae 100644 --- a/nerfstudio/data/dataparsers/blender_dataparser.py +++ b/nerfstudio/data/dataparsers/blender_dataparser.py @@ -24,6 +24,7 @@ import numpy as np import torch from nerfstudio.cameras.cameras import Cameras, CameraType +from nerfstudio.configs.base_config import InstantiateConfig from nerfstudio.data.dataparsers.base_dataparser import ( DataParser, DataParserConfig, @@ -35,7 +36,7 @@ from nerfstudio.utils.io import load_from_json @dataclass -class BlenderDataParserConfig(DataParserConfig): +class BlenderDataParserConfig(DataParserConfig, InstantiateConfig["Blender"]): """Blender dataset parser config""" _target: Type = field(default_factory=lambda: Blender) diff --git a/nerfstudio/data/dataparsers/dnerf_dataparser.py b/nerfstudio/data/dataparsers/dnerf_dataparser.py index 5d8d4db0e..5c8cb41a1 100644 --- a/nerfstudio/data/dataparsers/dnerf_dataparser.py +++ b/nerfstudio/data/dataparsers/dnerf_dataparser.py @@ -24,6 +24,7 @@ import numpy as np import torch from nerfstudio.cameras.cameras import Cameras, CameraType +from nerfstudio.configs.base_config import InstantiateConfig from nerfstudio.data.dataparsers.base_dataparser import ( DataParser, DataParserConfig, @@ -35,11 +36,9 @@ from nerfstudio.utils.io import load_from_json @dataclass -class DNeRFDataParserConfig(DataParserConfig): +class DNeRFDataParserConfig(DataParserConfig, InstantiateConfig["DNeRF"]): """D-NeRF dataset parser config""" - _target: Type = field(default_factory=lambda: DNeRF) - """target class to instantiate""" data: Path = Path("data/dnerf/lego") """Directory specifying location of data.""" scale_factor: float = 1.0 diff --git a/nerfstudio/data/dataparsers/dycheck_dataparser.py b/nerfstudio/data/dataparsers/dycheck_dataparser.py index 6239fc780..98dcfb62e 100644 --- a/nerfstudio/data/dataparsers/dycheck_dataparser.py +++ b/nerfstudio/data/dataparsers/dycheck_dataparser.py @@ -26,6 +26,7 @@ import torch from rich.progress import Console from nerfstudio.cameras.cameras import Cameras, CameraType +from nerfstudio.configs.base_config import InstantiateConfig from nerfstudio.data.dataparsers.base_dataparser import ( DataParser, DataParserConfig, @@ -180,11 +181,9 @@ def _rescale_depth(depth_raw: np.ndarray, cam: Dict) -> np.ndarray: @dataclass -class DycheckDataParserConfig(DataParserConfig): +class DycheckDataParserConfig(DataParserConfig, InstantiateConfig["Dycheck"]): """Dycheck (https://arxiv.org/abs/2210.13445) dataset parser config""" - _target: Type = field(default_factory=lambda: Dycheck) - """target class to instantiate""" data: Path = Path("data/iphone/mochi-high-five") """Directory specifying location of data.""" scale_factor: float = 5.0 diff --git a/nerfstudio/data/dataparsers/instant_ngp_dataparser.py b/nerfstudio/data/dataparsers/instant_ngp_dataparser.py index a8281c536..1a4dfec1e 100644 --- a/nerfstudio/data/dataparsers/instant_ngp_dataparser.py +++ b/nerfstudio/data/dataparsers/instant_ngp_dataparser.py @@ -27,6 +27,7 @@ from rich.console import Console from nerfstudio.cameras import camera_utils from nerfstudio.cameras.cameras import Cameras, CameraType +from nerfstudio.configs.base_config import InstantiateConfig from nerfstudio.data.dataparsers.base_dataparser import ( DataParser, DataParserConfig, @@ -39,11 +40,9 @@ CONSOLE = Console(width=120) @dataclass -class InstantNGPDataParserConfig(DataParserConfig): +class InstantNGPDataParserConfig(DataParserConfig, InstantiateConfig["InstantNGP"]): """Instant-NGP dataset parser config""" - _target: Type = field(default_factory=lambda: InstantNGP) - """target class to instantiate""" data: Path = Path("data/ours/posterv2") """Directory or explicit json file path specifying location of data.""" scene_scale: float = 0.3333 diff --git a/nerfstudio/data/dataparsers/minimal_dataparser.py b/nerfstudio/data/dataparsers/minimal_dataparser.py index 506ccfc43..4fec65b84 100644 --- a/nerfstudio/data/dataparsers/minimal_dataparser.py +++ b/nerfstudio/data/dataparsers/minimal_dataparser.py @@ -24,6 +24,7 @@ import numpy as np import torch from nerfstudio.cameras.cameras import Cameras +from nerfstudio.configs.base_config import InstantiateConfig from nerfstudio.data.dataparsers.base_dataparser import ( DataParser, DataParserConfig, @@ -34,11 +35,9 @@ from nerfstudio.data.scene_box import SceneBox @dataclass -class MinimalDataParserConfig(DataParserConfig): +class MinimalDataParserConfig(DataParserConfig, InstantiateConfig["MinimalDataParser"]): """Minimal dataset config""" - _target: Type = field(default_factory=lambda: MinimalDataParser) - """target class to instantiate""" data: Path = Path("/home/nikhil/nerfstudio-main/tests/data/lego_test/minimal_parser") diff --git a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py index b6e78a0d8..c3f1957bd 100644 --- a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py +++ b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py @@ -28,6 +28,7 @@ from typing_extensions import Literal from nerfstudio.cameras import camera_utils from nerfstudio.cameras.cameras import CAMERA_MODEL_TO_TYPE, Cameras, CameraType +from nerfstudio.configs.base_config import InstantiateConfig from nerfstudio.data.dataparsers.base_dataparser import ( DataParser, DataParserConfig, @@ -41,11 +42,9 @@ MAX_AUTO_RESOLUTION = 1600 @dataclass -class NerfstudioDataParserConfig(DataParserConfig): +class NerfstudioDataParserConfig(DataParserConfig, InstantiateConfig["Nerfstudio"]): """Nerfstudio dataset config""" - _target: Type = field(default_factory=lambda: Nerfstudio) - """target class to instantiate""" data: Path = Path("data/nerfstudio/poster") """Directory or explicit json file path specifying location of data.""" scale_factor: float = 1.0 diff --git a/nerfstudio/data/dataparsers/nuscenes_dataparser.py b/nerfstudio/data/dataparsers/nuscenes_dataparser.py index da15208d2..281f6758e 100644 --- a/nerfstudio/data/dataparsers/nuscenes_dataparser.py +++ b/nerfstudio/data/dataparsers/nuscenes_dataparser.py @@ -26,6 +26,7 @@ from nuscenes.nuscenes import NuScenes as NuScenesDatabase from typing_extensions import Literal from nerfstudio.cameras.cameras import Cameras, CameraType +from nerfstudio.configs.base_config import InstantiateConfig from nerfstudio.data.dataparsers.base_dataparser import ( DataParser, DataParserConfig, @@ -50,7 +51,7 @@ def rotation_translation_to_pose(r_quat, t_vec): @dataclass -class NuScenesDataParserConfig(DataParserConfig): +class NuScenesDataParserConfig(DataParserConfig, InstantiateConfig["NuScenes"]): """NuScenes dataset config. NuScenes (https://www.nuscenes.org/nuscenes) is an autonomous driving dataset containing 1000 20s clips. Each clip was recorded with a suite of sensors including 6 surround cameras. @@ -59,8 +60,6 @@ class NuScenesDataParserConfig(DataParserConfig): To create these masks use scripts/datasets/process_nuscenes_masks.py. """ - _target: Type = field(default_factory=lambda: NuScenes) - """target class to instantiate""" data: Path = Path("scene-0103") # TODO: rename to scene but keep checkpoint saving name? """Name of the scene.""" data_dir: Path = Path("/mnt/local/NuScenes") diff --git a/nerfstudio/data/dataparsers/phototourism_dataparser.py b/nerfstudio/data/dataparsers/phototourism_dataparser.py index 68ea89326..a06123115 100644 --- a/nerfstudio/data/dataparsers/phototourism_dataparser.py +++ b/nerfstudio/data/dataparsers/phototourism_dataparser.py @@ -27,6 +27,7 @@ from typing_extensions import Literal from nerfstudio.cameras import camera_utils from nerfstudio.cameras.cameras import Cameras, CameraType +from nerfstudio.configs.base_config import InstantiateConfig from nerfstudio.data.dataparsers.base_dataparser import ( DataParser, DataParserConfig, @@ -45,11 +46,9 @@ CONSOLE = Console(width=120) @dataclass -class PhototourismDataParserConfig(DataParserConfig): +class PhototourismDataParserConfig(DataParserConfig, InstantiateConfig["Phototourism"]): """Phototourism dataset parser config""" - _target: Type = field(default_factory=lambda: Phototourism) - """target class to instantiate""" data: Path = Path("data/phototourism/brandenburg-gate") """Directory specifying location of data.""" scale_factor: float = 3.0 diff --git a/nerfstudio/data/dataparsers/scannet_dataparser.py b/nerfstudio/data/dataparsers/scannet_dataparser.py index 80d0dd688..7453a3bb8 100644 --- a/nerfstudio/data/dataparsers/scannet_dataparser.py +++ b/nerfstudio/data/dataparsers/scannet_dataparser.py @@ -23,6 +23,7 @@ import numpy as np import torch from nerfstudio.cameras.cameras import Cameras, CameraType +from nerfstudio.configs.base_config import InstantiateConfig from nerfstudio.data.dataparsers.base_dataparser import ( DataParser, DataParserConfig, @@ -32,7 +33,7 @@ from nerfstudio.data.scene_box import SceneBox @dataclass -class ScanNetDataParserConfig(DataParserConfig): +class ScanNetDataParserConfig(DataParserConfig, InstantiateConfig["ScanNet"]): """ScanNet dataset config. ScanNet dataset (https://www.scan-net.org/) is a large-scale 3D dataset of indoor scenes. This dataparser assumes that the dense stream was extracted from .sens files. @@ -43,8 +44,6 @@ class ScanNetDataParserConfig(DataParserConfig): - pose/ """ - _target: Type = field(default_factory=lambda: ScanNet) - """target class to instantiate""" data: Path = Path("data/scannet/scene0423_02") """Path to ScanNet folder with densely extracted scenes.""" scene_scale: float = 1.0 diff --git a/nerfstudio/data/dataparsers/sdfstudio_dataparser.py b/nerfstudio/data/dataparsers/sdfstudio_dataparser.py index 460a4a709..f203194ca 100644 --- a/nerfstudio/data/dataparsers/sdfstudio_dataparser.py +++ b/nerfstudio/data/dataparsers/sdfstudio_dataparser.py @@ -24,6 +24,7 @@ from rich.console import Console from nerfstudio.cameras import camera_utils from nerfstudio.cameras.cameras import Cameras, CameraType +from nerfstudio.configs.base_config import InstantiateConfig from nerfstudio.data.dataparsers.base_dataparser import ( DataParser, DataParserConfig, @@ -36,11 +37,9 @@ CONSOLE = Console() @dataclass -class SDFStudioDataParserConfig(DataParserConfig): +class SDFStudioDataParserConfig(DataParserConfig, InstantiateConfig["SDFStudio"]): """Scene dataset parser config""" - _target: Type = field(default_factory=lambda: SDFStudio) - """target class to instantiate""" data: Path = Path("data/DTU/scan65") """Directory specifying location of data.""" include_mono_prior: bool = False diff --git a/nerfstudio/data/dataparsers/sitcoms3d_dataparser.py b/nerfstudio/data/dataparsers/sitcoms3d_dataparser.py index 426fe6d21..82dd05d59 100644 --- a/nerfstudio/data/dataparsers/sitcoms3d_dataparser.py +++ b/nerfstudio/data/dataparsers/sitcoms3d_dataparser.py @@ -28,6 +28,7 @@ import torch from rich.console import Console from nerfstudio.cameras.cameras import Cameras, CameraType +from nerfstudio.configs.base_config import InstantiateConfig from nerfstudio.data.dataparsers.base_dataparser import ( DataParser, DataParserConfig, @@ -41,11 +42,9 @@ CONSOLE = Console() @dataclass -class Sitcoms3DDataParserConfig(DataParserConfig): +class Sitcoms3DDataParserConfig(DataParserConfig, InstantiateConfig["Sitcoms3D"]): """sitcoms3D dataset parser config""" - _target: Type = field(default_factory=lambda: Sitcoms3D) - """target class to instantiate""" data: Path = Path("data/sitcoms3d/TBBT-big_living_room") """Directory specifying location of data.""" include_semantics: bool = True diff --git a/nerfstudio/engine/schedulers.py b/nerfstudio/engine/schedulers.py index 0b6586ffc..e436682a5 100644 --- a/nerfstudio/engine/schedulers.py +++ b/nerfstudio/engine/schedulers.py @@ -28,7 +28,7 @@ from nerfstudio.configs.base_config import InstantiateConfig @dataclass -class SchedulerConfig(InstantiateConfig): +class SchedulerConfig(InstantiateConfig["Scheduler"]): """Basic scheduler config""" _target: Type = field(default_factory=lambda: Scheduler) diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index f900dd863..f427b405b 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -57,11 +57,9 @@ TORCH_DEVICE = Union[torch.device, str] # pylint: disable=invalid-name @dataclass -class TrainerConfig(ExperimentConfig): +class TrainerConfig(ExperimentConfig["Trainer"]): """Configuration for training regimen""" - _target: Type = field(default_factory=lambda: Trainer) - """target class to instantiate""" steps_per_save: int = 1000 """Number of steps between saves.""" steps_per_eval_batch: int = 500 diff --git a/nerfstudio/models/base_model.py b/nerfstudio/models/base_model.py index aa2181b38..c4069790b 100644 --- a/nerfstudio/models/base_model.py +++ b/nerfstudio/models/base_model.py @@ -34,14 +34,13 @@ from nerfstudio.data.scene_box import SceneBox from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes from nerfstudio.model_components.scene_colliders import NearFarCollider +TModel = TypeVar("TModel", bound="Model") # Model related configs @dataclass -class ModelConfig(InstantiateConfig): +class ModelConfig(InstantiateConfig[TModel]): """Configuration for model instantiation""" - _target: Type = field(default_factory=lambda: Model) - """target class to instantiate""" enable_collider: bool = True """Whether to create a scene collider to filter rays.""" collider_params: Optional[Dict[str, float]] = to_immutable_dict({"near_plane": 2.0, "far_plane": 6.0}) diff --git a/nerfstudio/models/depth_nerfacto.py b/nerfstudio/models/depth_nerfacto.py index 683c89e1f..056580a38 100644 --- a/nerfstudio/models/depth_nerfacto.py +++ b/nerfstudio/models/depth_nerfacto.py @@ -18,6 +18,7 @@ Nerfacto augmented with depth supervision. from __future__ import annotations +import typing from dataclasses import dataclass, field from typing import Dict, Tuple, Type @@ -49,6 +50,9 @@ class DepthNerfactoModelConfig(NerfactoModelConfig): depth_loss_type: DepthLossType = DepthLossType.DS_NERF """Depth loss type.""" + def setup(self, **kwargs) -> DepthNerfactoModel: + return typing.cast(DepthNerfactoModel, super().setup(**kwargs)) + class DepthNerfactoModel(NerfactoModel): """Depth loss augmented nerfacto model. @@ -59,6 +63,9 @@ class DepthNerfactoModel(NerfactoModel): config: DepthNerfactoModelConfig + def __init__(self, config: DepthNerfactoModelConfig, scene_box: SceneBox, num_train_data: int, **kwargs) -> None: + super().__init__(config, scene_box, num_train_data, **kwargs) + def populate_modules(self): """Set the fields and modules.""" super().populate_modules() diff --git a/nerfstudio/models/instant_ngp.py b/nerfstudio/models/instant_ngp.py index 21746a683..920d88930 100644 --- a/nerfstudio/models/instant_ngp.py +++ b/nerfstudio/models/instant_ngp.py @@ -50,7 +50,7 @@ from nerfstudio.utils import colormaps @dataclass -class InstantNGPModelConfig(ModelConfig): +class InstantNGPModelConfig(ModelConfig["NGPModel"]): """Instant NGP Model Config""" _target: Type = field( diff --git a/nerfstudio/models/mipnerf.py b/nerfstudio/models/mipnerf.py index 992390351..b467d15df 100644 --- a/nerfstudio/models/mipnerf.py +++ b/nerfstudio/models/mipnerf.py @@ -17,7 +17,8 @@ Implementation of mip-NeRF. """ from __future__ import annotations -from typing import Dict, List, Tuple +import dataclasses +from typing import Dict, List, Tuple, Type import torch from torch.nn import Parameter @@ -36,12 +37,19 @@ from nerfstudio.model_components.renderers import ( DepthRenderer, RGBRenderer, ) -from nerfstudio.models.base_model import Model +from nerfstudio.models.base_model import Model, ModelConfig from nerfstudio.models.vanilla_nerf import VanillaModelConfig from nerfstudio.utils import colormaps, colors, misc -class MipNerfModel(Model[VanillaModelConfig]): +class MipNerfModelConfig(VanillaModelConfig): + _target: Type = dataclasses.field(default_factory=lambda: MipNerfModel) + + def setup(self, **kwargs) -> MipNerfModel: + return typing.cast(MipNerfModel, super().setup(**kwargs)) + + +class MipNerfModel(Model[MipNerfModelConfig]): """mip-NeRF model Args: diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 543b3ad60..e6ab09f7a 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -63,10 +63,9 @@ from nerfstudio.utils import colormaps @dataclass -class NerfactoModelConfig(ModelConfig): +class NerfactoModelConfig(ModelConfig["NerfactoModel"]): """Nerfacto Model Config""" - _target: Type = field(default_factory=lambda: NerfactoModel) near_plane: float = 0.05 """How far along the ray to start sampling.""" far_plane: float = 1000.0 diff --git a/nerfstudio/models/nerfplayer_nerfacto.py b/nerfstudio/models/nerfplayer_nerfacto.py index e10b3cad2..805f7e7df 100644 --- a/nerfstudio/models/nerfplayer_nerfacto.py +++ b/nerfstudio/models/nerfplayer_nerfacto.py @@ -19,6 +19,7 @@ NeRFPlayer (https://arxiv.org/abs/2210.15947) implementation with nerfacto backb from __future__ import annotations import functools +import typing from dataclasses import dataclass, field from typing import Dict, List, Type @@ -60,6 +61,7 @@ class NerfplayerNerfactoModelConfig(NerfactoModelConfig): """Nerfplayer Model Config with Nerfacto backbone""" _target: Type = field(default_factory=lambda: NerfplayerNerfactoModel) + near_plane: float = 0.05 """How far along the ray to start sampling.""" far_plane: float = 1000.0 @@ -88,6 +90,9 @@ class NerfplayerNerfactoModelConfig(NerfactoModelConfig): depth_weight: float = 1e-1 """depth loss balancing weight for feature channels.""" + def setup(self, **kwargs) -> NerfplayerNerfactoModel: + return typing.cast(NerfplayerNerfactoModel, super().setup(**kwargs)) + class NerfplayerNerfactoModel(NerfactoModel): """Nerfplayer model with Nerfacto backbone. diff --git a/nerfstudio/models/nerfplayer_ngp.py b/nerfstudio/models/nerfplayer_ngp.py index 09f68d67f..ec778a21c 100644 --- a/nerfstudio/models/nerfplayer_ngp.py +++ b/nerfstudio/models/nerfplayer_ngp.py @@ -18,6 +18,7 @@ Implementation of NeRFPlayer (https://arxiv.org/abs/2210.15947) with InstantNGP from __future__ import annotations +import typing from dataclasses import dataclass, field from typing import Type @@ -81,6 +82,9 @@ class NerfplayerNGPModelConfig(InstantNGPModelConfig): disable_viewing_dependent: bool = True """Disable viewing dependent effects.""" + def setup(self, **kwargs) -> NerfplayerNGPModel: + return typing.cast(NerfplayerNGPModel, super().setup(**kwargs)) + class NerfplayerNGPModel(NGPModel): """NeRFPlayer Model with InstantNGP backbone. diff --git a/nerfstudio/models/semantic_nerfw.py b/nerfstudio/models/semantic_nerfw.py index a51bbe7d1..86890dbda 100644 --- a/nerfstudio/models/semantic_nerfw.py +++ b/nerfstudio/models/semantic_nerfw.py @@ -18,6 +18,7 @@ Semantic NeRF-W implementation which should be fast enough to view in the viewer from __future__ import annotations +import typing from dataclasses import dataclass, field from typing import Dict, List, Tuple, Type @@ -64,6 +65,9 @@ class SemanticNerfWModelConfig(NerfactoModelConfig): semantic_loss_weight: float = 1.0 pass_semantic_gradients: bool = False + def setup(self, **kwargs) -> SemanticNerfWModel: + return typing.cast(SemanticNerfWModel, super().setup(**kwargs)) + class SemanticNerfWModel(Model[SemanticNerfWModelConfig]): """Nerfacto model diff --git a/nerfstudio/models/tensorf.py b/nerfstudio/models/tensorf.py index 5784a763f..cd7a9ee10 100644 --- a/nerfstudio/models/tensorf.py +++ b/nerfstudio/models/tensorf.py @@ -57,11 +57,9 @@ from nerfstudio.utils import colormaps, colors, misc @dataclass -class TensoRFModelConfig(ModelConfig): +class TensoRFModelConfig(ModelConfig["TensoRFModel"]): """TensoRF model config""" - _target: Type = field(default_factory=lambda: TensoRFModel) - """target class to instantiate""" init_resolution: int = 128 """initial render resolution""" final_resolution: int = 300 diff --git a/nerfstudio/models/vanilla_nerf.py b/nerfstudio/models/vanilla_nerf.py index 034d465d5..de3ab6769 100644 --- a/nerfstudio/models/vanilla_nerf.py +++ b/nerfstudio/models/vanilla_nerf.py @@ -45,10 +45,9 @@ from nerfstudio.utils import colormaps, colors, misc @dataclass -class VanillaModelConfig(ModelConfig): +class VanillaModelConfig(ModelConfig["NeRFModel"]): """Vanilla Model Config""" - _target: Type = field(default_factory=lambda: NeRFModel) num_coarse_samples: int = 64 """Number of samples in coarse field evaluation""" num_importance_samples: int = 128 diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 063237abd..b9a702f2c 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -189,22 +189,20 @@ class Pipeline(nn.Module): """ +TModelConfig = TypeVar("TModelConfig", bound=ModelConfig) + + @dataclass -class VanillaPipelineConfig(cfg.InstantiateConfig): +class VanillaPipelineConfig(cfg.InstantiateConfig["VanillaPipeline"]): """Configuration for pipeline instantiation""" - _target: Type = field(default_factory=lambda: VanillaPipeline) - """target class to instantiate""" datamanager: VanillaDataManagerConfig = VanillaDataManagerConfig() """specifies the datamanager config""" - model: ModelConfig = ModelConfig() + model: ModelConfig = field(default_factory=ModelConfig[Model]) """specifies the model config""" -TPipelineConfig = TypeVar("TPipelineConfig", bound=VanillaPipelineConfig, default=VanillaPipelineConfig) - - -class VanillaPipeline(Pipeline, Generic[TPipelineConfig]): +class VanillaPipeline(Pipeline): """The pipeline class for the vanilla nerf setup of multiple cameras for one or a few scenes. config: configuration to instantiate pipeline @@ -221,11 +219,11 @@ class VanillaPipeline(Pipeline, Generic[TPipelineConfig]): model: The model that will be used """ - config: TPipelineConfig + config: VanillaPipelineConfig def __init__( self, - config: TPipelineConfig, + config: VanillaPipelineConfig, device: str, test_mode: Literal["test", "val", "inference"] = "val", world_size: int = 1, diff --git a/nerfstudio/pipelines/dynamic_batch.py b/nerfstudio/pipelines/dynamic_batch.py index 0bec74f81..3fb94ead9 100644 --- a/nerfstudio/pipelines/dynamic_batch.py +++ b/nerfstudio/pipelines/dynamic_batch.py @@ -16,12 +16,16 @@ A pipeline that dynamically chooses the number of rays to sample. """ +from __future__ import annotations + +import typing from dataclasses import dataclass, field from typing import Type import torch -from typing_extensions import Literal +from typing_extensions import Literal, TypeVar +from nerfstudio.configs.base_config import InstantiateConfig from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager from nerfstudio.pipelines.base_pipeline import VanillaPipeline, VanillaPipelineConfig @@ -31,17 +35,22 @@ class DynamicBatchPipelineConfig(VanillaPipelineConfig): """Dynamic Batch Pipeline Config""" _target: Type = field(default_factory=lambda: DynamicBatchPipeline) + target_num_samples: int = 262144 # 1 << 18 """The target number of samples to use for an entire batch of rays.""" max_num_samples_per_ray: int = 1024 # 1 << 10 """The maximum number of samples to be placed along a ray.""" + def setup(self, **kwargs) -> DynamicBatchPipeline: + return typing.cast(DynamicBatchPipeline, super().setup(**kwargs)) -class DynamicBatchPipeline(VanillaPipeline[DynamicBatchPipelineConfig]): + +class DynamicBatchPipeline(VanillaPipeline): """Pipeline with logic for changing the number of rays per batch.""" # pylint: disable=abstract-method + config: DynamicBatchPipelineConfig datamanager: VanillaDataManager dynamic_num_rays_per_batch: int