mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-15 12:23:46 +08:00
326 lines
12 KiB
Python
326 lines
12 KiB
Python
# Copyright 2022 The Nerfstudio Team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""
|
|
Put all the method implementations in one location.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Dict
|
|
|
|
import tyro
|
|
from nerfacc import ContractionType
|
|
|
|
from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig
|
|
from nerfstudio.configs.base_config import TrainerConfig, ViewerConfig
|
|
from nerfstudio.configs.experiment_config import ExperimentConfig
|
|
from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManagerConfig
|
|
from nerfstudio.data.datamanagers.semantic_datamanager import SemanticDataManagerConfig
|
|
from nerfstudio.data.datamanagers.variable_res_datamanager import (
|
|
VariableResDataManagerConfig,
|
|
)
|
|
from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig
|
|
from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig
|
|
from nerfstudio.data.dataparsers.friends_dataparser import FriendsDataParserConfig
|
|
from nerfstudio.data.dataparsers.instant_ngp_dataparser import (
|
|
InstantNGPDataParserConfig,
|
|
)
|
|
from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig
|
|
from nerfstudio.data.dataparsers.phototourism_dataparser import (
|
|
PhototourismDataParserConfig,
|
|
)
|
|
from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig
|
|
from nerfstudio.engine.schedulers import SchedulerConfig
|
|
from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind
|
|
from nerfstudio.models.instant_ngp import InstantNGPModelConfig
|
|
from nerfstudio.models.mipnerf import MipNerfModel
|
|
from nerfstudio.models.nerfacto import NerfactoModelConfig
|
|
from nerfstudio.models.nerfacto_4k import Nerfacto4KModelConfig
|
|
from nerfstudio.models.semantic_nerfw import SemanticNerfWModelConfig
|
|
from nerfstudio.models.tensorf import TensoRFModelConfig
|
|
from nerfstudio.models.vanilla_nerf import NeRFModel, VanillaModelConfig
|
|
from nerfstudio.pipelines.base_pipeline import VanillaPipelineConfig
|
|
from nerfstudio.pipelines.dynamic_batch import DynamicBatchPipelineConfig
|
|
|
|
method_configs: Dict[str, ExperimentConfig] = {}
|
|
descriptions = {
|
|
"nerfacto": "Recommended real-time model tuned for real captures. This model will be continually updated.",
|
|
"instant-ngp": "Implementation of Instant-NGP. Recommended real-time model for unbounded scenes.",
|
|
"instant-ngp-bounded": "Implementation of Instant-NGP. Recommended for bounded real and synthetic scenes",
|
|
"mipnerf": "High quality model for bounded scenes. (slow)",
|
|
"semantic-nerfw": "Predicts semantic segmentations and filters out transient objects.",
|
|
"vanilla-nerf": "Original NeRF model. (slow)",
|
|
"tensorf": "tensorf",
|
|
"dnerf": "Dynamic-NeRF model. (slow)",
|
|
"phototourism": "Uses the Phototourism data.",
|
|
"nerfacto-4k": "4k nerf with nerfacto",
|
|
}
|
|
|
|
method_configs["nerfacto-4k"] = ExperimentConfig(
|
|
method_name="nerfacto-4k",
|
|
trainer=TrainerConfig(
|
|
steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True
|
|
),
|
|
pipeline=VanillaPipelineConfig(
|
|
datamanager=VanillaDataManagerConfig(
|
|
dataparser=NerfstudioDataParserConfig(),
|
|
train_num_rays_per_batch=4096,
|
|
eval_num_rays_per_batch=4096,
|
|
camera_optimizer=CameraOptimizerConfig(
|
|
mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
|
|
),
|
|
pixel_sampling="patch",
|
|
pixel_patch_size=8,
|
|
),
|
|
model=Nerfacto4KModelConfig(eval_num_rays_per_chunk=1 << 15),
|
|
),
|
|
optimizers={
|
|
"proposal_networks": {
|
|
"optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
|
"scheduler": None,
|
|
},
|
|
"fields": {
|
|
"optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
|
"scheduler": None,
|
|
},
|
|
},
|
|
viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
|
|
vis="viewer",
|
|
)
|
|
|
|
method_configs["nerfacto"] = ExperimentConfig(
|
|
method_name="nerfacto",
|
|
trainer=TrainerConfig(
|
|
steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True
|
|
),
|
|
pipeline=VanillaPipelineConfig(
|
|
datamanager=VanillaDataManagerConfig(
|
|
dataparser=NerfstudioDataParserConfig(),
|
|
train_num_rays_per_batch=4096,
|
|
eval_num_rays_per_batch=4096,
|
|
camera_optimizer=CameraOptimizerConfig(
|
|
mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
|
|
),
|
|
),
|
|
model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 15),
|
|
),
|
|
optimizers={
|
|
"proposal_networks": {
|
|
"optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
|
"scheduler": None,
|
|
},
|
|
"fields": {
|
|
"optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
|
"scheduler": None,
|
|
},
|
|
},
|
|
viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
|
|
vis="viewer",
|
|
)
|
|
|
|
method_configs["instant-ngp"] = ExperimentConfig(
|
|
method_name="instant-ngp",
|
|
trainer=TrainerConfig(
|
|
steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True
|
|
),
|
|
pipeline=DynamicBatchPipelineConfig(
|
|
datamanager=VanillaDataManagerConfig(dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=8192),
|
|
model=InstantNGPModelConfig(eval_num_rays_per_chunk=8192),
|
|
),
|
|
optimizers={
|
|
"fields": {
|
|
"optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
|
"scheduler": None,
|
|
}
|
|
},
|
|
viewer=ViewerConfig(num_rays_per_chunk=64000),
|
|
vis="viewer",
|
|
)
|
|
|
|
|
|
method_configs["instant-ngp-bounded"] = ExperimentConfig(
|
|
method_name="instant-ngp-bounded",
|
|
trainer=TrainerConfig(
|
|
steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True
|
|
),
|
|
pipeline=DynamicBatchPipelineConfig(
|
|
datamanager=VanillaDataManagerConfig(dataparser=InstantNGPDataParserConfig(), train_num_rays_per_batch=8192),
|
|
model=InstantNGPModelConfig(
|
|
eval_num_rays_per_chunk=8192,
|
|
contraction_type=ContractionType.AABB,
|
|
render_step_size=0.001,
|
|
max_num_samples_per_ray=48,
|
|
near_plane=0.01,
|
|
background_color="black",
|
|
),
|
|
),
|
|
optimizers={
|
|
"fields": {
|
|
"optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
|
"scheduler": None,
|
|
}
|
|
},
|
|
viewer=ViewerConfig(num_rays_per_chunk=64000),
|
|
vis="viewer",
|
|
)
|
|
|
|
|
|
method_configs["mipnerf"] = ExperimentConfig(
|
|
method_name="mipnerf",
|
|
pipeline=VanillaPipelineConfig(
|
|
datamanager=VanillaDataManagerConfig(dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=1024),
|
|
model=VanillaModelConfig(
|
|
_target=MipNerfModel,
|
|
loss_coefficients={"rgb_loss_coarse": 0.1, "rgb_loss_fine": 1.0},
|
|
num_coarse_samples=128,
|
|
num_importance_samples=128,
|
|
eval_num_rays_per_chunk=1024,
|
|
),
|
|
),
|
|
optimizers={
|
|
"fields": {
|
|
"optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08),
|
|
"scheduler": None,
|
|
}
|
|
},
|
|
)
|
|
|
|
method_configs["semantic-nerfw"] = ExperimentConfig(
|
|
method_name="semantic-nerfw",
|
|
trainer=TrainerConfig(
|
|
steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True
|
|
),
|
|
pipeline=VanillaPipelineConfig(
|
|
datamanager=SemanticDataManagerConfig(
|
|
dataparser=FriendsDataParserConfig(), train_num_rays_per_batch=4096, eval_num_rays_per_batch=8192
|
|
),
|
|
model=SemanticNerfWModelConfig(eval_num_rays_per_chunk=1 << 16),
|
|
),
|
|
optimizers={
|
|
"proposal_networks": {
|
|
"optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
|
"scheduler": None,
|
|
},
|
|
"fields": {
|
|
"optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
|
"scheduler": None,
|
|
},
|
|
},
|
|
viewer=ViewerConfig(num_rays_per_chunk=1 << 16),
|
|
vis="viewer",
|
|
)
|
|
|
|
method_configs["vanilla-nerf"] = ExperimentConfig(
|
|
method_name="vanilla-nerf",
|
|
pipeline=VanillaPipelineConfig(
|
|
datamanager=VanillaDataManagerConfig(
|
|
dataparser=BlenderDataParserConfig(),
|
|
),
|
|
model=VanillaModelConfig(_target=NeRFModel),
|
|
),
|
|
optimizers={
|
|
"fields": {
|
|
"optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08),
|
|
"scheduler": None,
|
|
},
|
|
"temporal_distortion": {
|
|
"optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08),
|
|
"scheduler": None,
|
|
},
|
|
},
|
|
)
|
|
|
|
method_configs["tensorf"] = ExperimentConfig(
|
|
method_name="tensorf",
|
|
trainer=TrainerConfig(mixed_precision=False),
|
|
pipeline=VanillaPipelineConfig(
|
|
datamanager=VanillaDataManagerConfig(
|
|
dataparser=BlenderDataParserConfig(),
|
|
),
|
|
model=TensoRFModelConfig(),
|
|
),
|
|
optimizers={
|
|
"fields": {
|
|
"optimizer": AdamOptimizerConfig(lr=0.001),
|
|
"scheduler": SchedulerConfig(lr_final=0.0001, max_steps=30000),
|
|
},
|
|
"encodings": {
|
|
"optimizer": AdamOptimizerConfig(lr=0.02),
|
|
"scheduler": SchedulerConfig(lr_final=0.002, max_steps=30000),
|
|
},
|
|
},
|
|
)
|
|
|
|
method_configs["dnerf"] = ExperimentConfig(
|
|
method_name="dnerf",
|
|
pipeline=VanillaPipelineConfig(
|
|
datamanager=VanillaDataManagerConfig(dataparser=DNeRFDataParserConfig()),
|
|
model=VanillaModelConfig(
|
|
_target=NeRFModel,
|
|
enable_temporal_distortion=True,
|
|
temporal_distortion_params={"kind": TemporalDistortionKind.DNERF},
|
|
),
|
|
),
|
|
optimizers={
|
|
"fields": {
|
|
"optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08),
|
|
"scheduler": None,
|
|
},
|
|
"temporal_distortion": {
|
|
"optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08),
|
|
"scheduler": None,
|
|
},
|
|
},
|
|
)
|
|
|
|
method_configs["phototourism"] = ExperimentConfig(
|
|
method_name="phototourism",
|
|
trainer=TrainerConfig(
|
|
steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True
|
|
),
|
|
pipeline=VanillaPipelineConfig(
|
|
datamanager=VariableResDataManagerConfig( # NOTE: one of the only differences with nerfacto
|
|
dataparser=PhototourismDataParserConfig(), # NOTE: one of the only differences with nerfacto
|
|
train_num_rays_per_batch=4096,
|
|
eval_num_rays_per_batch=4096,
|
|
camera_optimizer=CameraOptimizerConfig(
|
|
mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
|
|
),
|
|
),
|
|
model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 15),
|
|
),
|
|
optimizers={
|
|
"proposal_networks": {
|
|
"optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
|
"scheduler": None,
|
|
},
|
|
"fields": {
|
|
"optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
|
"scheduler": None,
|
|
},
|
|
},
|
|
viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
|
|
vis="viewer",
|
|
)
|
|
|
|
AnnotatedBaseConfigUnion = tyro.conf.SuppressFixed[ # Don't show unparseable (fixed) arguments in helptext.
|
|
tyro.conf.FlagConversionOff[
|
|
tyro.extras.subcommand_type_from_defaults(defaults=method_configs, descriptions=descriptions)
|
|
]
|
|
]
|
|
"""Union[] type over config types, annotated with default instances for use with
|
|
tyro.cli(). Allows the user to pick between one of several base configurations, and
|
|
then override values in it."""
|