diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index a3649458a..2c829e092 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -37,13 +37,12 @@ from nerfstudio.data.dataparsers.friends_dataparser import FriendsDataParserConf from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind -from nerfstudio.models.base_model import VanillaModelConfig from nerfstudio.models.instant_ngp import InstantNGPModelConfig from nerfstudio.models.mipnerf import MipNerfModel from nerfstudio.models.nerfacto import NerfactoModelConfig from nerfstudio.models.semantic_nerfw import SemanticNerfWModelConfig from nerfstudio.models.tensorf import TensoRFModelConfig -from nerfstudio.models.vanilla_nerf import NeRFModel +from nerfstudio.models.vanilla_nerf import NeRFModel, VanillaModelConfig from nerfstudio.pipelines.base_pipeline import VanillaPipelineConfig from nerfstudio.pipelines.dynamic_batch import DynamicBatchPipelineConfig diff --git a/nerfstudio/models/base_model.py b/nerfstudio/models/base_model.py index 57d0d3bea..6be459343 100644 --- a/nerfstudio/models/base_model.py +++ b/nerfstudio/models/base_model.py @@ -32,7 +32,6 @@ from nerfstudio.configs.base_config import InstantiateConfig from nerfstudio.configs.config_utils import to_immutable_dict from nerfstudio.data.scene_box import SceneBox from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes -from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind from nerfstudio.model_components.scene_colliders import NearFarCollider @@ -210,18 +209,3 @@ class Model(nn.Module): """ state = {key.replace("module.", ""): value for key, value in loaded_state["model"].items()} self.load_state_dict(state) # type: ignore - - -@dataclass -class VanillaModelConfig(ModelConfig): - """Vanilla Model Config""" - - num_coarse_samples: int = 64 - """Number of samples in coarse field evaluation""" - num_importance_samples: int = 128 - """Number of samples in fine field evaluation""" - - enable_temporal_distortion: bool = False - """Specifies whether or not to include ray warping based on time.""" - temporal_distortion_params: Dict[str, Any] = to_immutable_dict({"kind": TemporalDistortionKind.DNERF}) - """Parameters to instantiate temporal distortion with""" diff --git a/nerfstudio/models/tensorf.py b/nerfstudio/models/tensorf.py index a0c771d2c..e3fa1ee65 100644 --- a/nerfstudio/models/tensorf.py +++ b/nerfstudio/models/tensorf.py @@ -46,12 +46,12 @@ from nerfstudio.model_components.renderers import ( RGBRenderer, ) from nerfstudio.model_components.scene_colliders import AABBBoxCollider -from nerfstudio.models.base_model import Model, VanillaModelConfig +from nerfstudio.models.base_model import Model, ModelConfig from nerfstudio.utils import colormaps, colors, misc @dataclass -class TensoRFModelConfig(VanillaModelConfig): +class TensoRFModelConfig(ModelConfig): """TensoRF model config""" _target: Type = field(default_factory=lambda: TensoRFModel) diff --git a/nerfstudio/models/vanilla_nerf.py b/nerfstudio/models/vanilla_nerf.py index 2c5227ee4..3e731524f 100644 --- a/nerfstudio/models/vanilla_nerf.py +++ b/nerfstudio/models/vanilla_nerf.py @@ -18,7 +18,8 @@ Implementation of vanilla nerf. from __future__ import annotations -from typing import Dict, List, Tuple +from dataclasses import dataclass, field +from typing import Any, Dict, List, Tuple, Type import torch from torch.nn import Parameter @@ -27,8 +28,10 @@ from torchmetrics.functional import structural_similarity_index_measure from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity from nerfstudio.cameras.rays import RayBundle +from nerfstudio.configs.config_utils import to_immutable_dict from nerfstudio.field_components.encodings import NeRFEncoding from nerfstudio.field_components.field_heads import FieldHeadNames +from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind from nerfstudio.fields.vanilla_nerf_field import NeRFField from nerfstudio.model_components.losses import MSELoss from nerfstudio.model_components.ray_samplers import PDFSampler, UniformSampler @@ -41,6 +44,22 @@ from nerfstudio.models.base_model import Model, ModelConfig from nerfstudio.utils import colormaps, colors, misc +@dataclass +class VanillaModelConfig(ModelConfig): + """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 + """Number of samples in fine field evaluation""" + + enable_temporal_distortion: bool = False + """Specifies whether or not to include ray warping based on time.""" + temporal_distortion_params: Dict[str, Any] = to_immutable_dict({"kind": TemporalDistortionKind.DNERF}) + """Parameters to instantiate temporal distortion with""" + + class NeRFModel(Model): """Vanilla NeRF model @@ -50,7 +69,7 @@ class NeRFModel(Model): def __init__( self, - config: ModelConfig, + config: VanillaModelConfig, **kwargs, ) -> None: self.field_coarse = None