tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 10:34:29 +08:00
Code

merge populate functions (#353)

Browse Source
This commit is contained in:
evonneng 2022-09-06 16:44:34 -07:00 committed by GitHub
parent 28d5b3bda1
commit fcba56b9b6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
8 changed files with 40 additions and 63 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
nerfactory/models/base.py
View File

@ -30,6 +30,7 @@ from nerfactory.cameras.rays import RayBundle
from nerfactory.configs import base as cfg
from nerfactory.datamanagers.structs import SceneBounds
from nerfactory.fields.density_fields.density_grid import DensityGrid
from nerfactory.models.modules.scene_colliders import NearFarCollider
from nerfactory.utils.callbacks import TrainingCallback, TrainingCallbackAttributes
from nerfactory.utils.misc import get_masked_dict, is_not_none
@ -64,9 +65,7 @@ class Model(nn.Module):
self.density_field = None
self.kwargs = kwargs
self.collider = None
self.populate_density_field()
self.populate_fields()
self.populate_misc_modules() # populate the modules
self.populate_modules() # populate the modules
self.callbacks = None
# to keep track of which device the nn.Module is on
self.device_indicator_param = nn.Parameter(torch.empty(0))
@ -82,10 +81,12 @@ class Model(nn.Module):
"""Returns a list of callbacks that run functions at the specified training iterations."""
return []
def populate_density_field(self):
"""Set the scene density field to use."""
if self.config.enable_density_field:
def populate_modules(self):
"""Set the necessary modules to get the network working."""
# default instantiates optional modules that are common among many networks
# NOTE: call `super().populate_modules()` in subclasses
if self.config.enable_density_field:
self.density_field = DensityGrid(
center=self.config.density_field_params["center"],
base_scale=self.config.density_field_params["base_scale"],
@ -93,13 +94,10 @@ class Model(nn.Module):
resolution=self.config.density_field_params["resolution"],
update_every_num_iters=self.config.density_field_params["update_every_num_iters"],
)
def populate_fields(self):
"""Set the fields."""
@abstractmethod
def populate_misc_modules(self):
"""Initializes any additional modules that are part of the network."""
if self.config.enable_collider:
self.collider = NearFarCollider(
near_plane=self.config.collider_params["near_plane"], far_plane=self.config.collider_params["far_plane"]
)
@abstractmethod
def get_param_groups(self) -> Dict[str, List[Parameter]]:

6
nerfactory/models/instant_ngp.py
View File

@ -70,12 +70,12 @@ class NGPModel(Model):
)
]
def populate_fields(self):
"""Set the fields."""
def populate_modules(self):
"""Set the fields and modules."""
super().populate_modules()
# torch or tiny-cuda-nn version
self.field = field_implementation_to_class[self.config.field_implementation](self.scene_bounds.aabb)
def populate_misc_modules(self):
# samplers
self.sampler = NGPSpacedSampler(num_samples=self.config.num_samples, density_field=self.density_field)

14
nerfactory/models/mipnerf.py
View File

@ -32,7 +32,6 @@ from nerfactory.fields.modules.field_heads import FieldHeadNames
from nerfactory.fields.nerf_field import NeRFField
from nerfactory.models.base import Model
from nerfactory.models.modules.ray_sampler import PDFSampler, UniformSampler
from nerfactory.models.modules.scene_colliders import NearFarCollider
from nerfactory.optimizers.loss import MSELoss
from nerfactory.renderers.renderers import (
AccumulationRenderer,
@ -53,9 +52,11 @@ class MipNerfModel(Model):
self.field = None
super().__init__(config=config, **kwargs)
def populate_fields(self):
"""Set the fields."""
def populate_modules(self):
"""Set the fields and modules"""
super().populate_modules()
# setting up fields
position_encoding = NeRFEncoding(
in_dim=3, num_frequencies=16, min_freq_exp=0.0, max_freq_exp=16.0, include_input=True
)
@ -67,7 +68,6 @@ class MipNerfModel(Model):
position_encoding=position_encoding, direction_encoding=direction_encoding, use_integrated_encoding=True
)
def populate_misc_modules(self):
# samplers
self.sampler_uniform = UniformSampler(num_samples=self.config.num_coarse_samples)
self.sampler_pdf = PDFSampler(num_samples=self.config.num_importance_samples, include_original=False)
@ -85,12 +85,6 @@ class MipNerfModel(Model):
self.ssim = structural_similarity_index_measure
self.lpips = LearnedPerceptualImagePatchSimilarity()
# colliders
if self.config.enable_collider:
self.collider = NearFarCollider(
near_plane=self.config.collider_params["near_plane"], far_plane=self.config.collider_params["far_plane"]
)
def get_param_groups(self) -> Dict[str, List[Parameter]]:
param_groups = {}
if self.field is None:

14
nerfactory/models/mipnerf_360.py
View File

@ -35,7 +35,6 @@ from nerfactory.fields.nerf_field import NeRFField
from nerfactory.models.base import Model
from nerfactory.models.modules.ray_losses import distortion_loss
from nerfactory.models.modules.ray_sampler import PDFSampler, UniformSampler
from nerfactory.models.modules.scene_colliders import NearFarCollider
from nerfactory.optimizers.loss import MSELoss
from nerfactory.renderers.renderers import (
AccumulationRenderer,
@ -56,9 +55,11 @@ class MipNerf360Model(Model):
self.field = None
super().__init__(config=config, **kwargs)
def populate_fields(self):
"""Set the fields."""
def populate_modules(self):
"""Set the fields and modules"""
super().populate_modules()
# setting up fields
position_encoding = NeRFEncoding(
in_dim=3, num_frequencies=10, min_freq_exp=0.0, max_freq_exp=8.0, include_input=True
)
@ -73,7 +74,6 @@ class MipNerf360Model(Model):
spatial_distortion=SceneContraction(),
)
def populate_misc_modules(self):
# samplers
self.sampler_uniform = UniformSampler(num_samples=self.config.num_coarse_samples)
self.sampler_pdf = PDFSampler(num_samples=self.config.num_importance_samples, include_original=False)
@ -91,12 +91,6 @@ class MipNerf360Model(Model):
self.ssim = structural_similarity_index_measure
self.lpips = LearnedPerceptualImagePatchSimilarity()
# colliders
if self.config.enable_collider:
self.collider = NearFarCollider(
near_plane=self.config.collider_params["near_plane"], far_plane=self.config.collider_params["far_plane"]
)
def get_param_groups(self) -> Dict[str, List[Parameter]]:
param_groups = {}
if self.field is None:

7
nerfactory/models/nerfw.py
View File

@ -68,9 +68,11 @@ class NerfWModel(Model):
self.transient_embedding_dim = 16
super().__init__(config=config, **kwargs)
def populate_fields(self):
"""Set the fields."""
def populate_modules(self):
"""Set the fields and modules"""
super().populate_modules()
# setting up fields
position_encoding = NeRFEncoding(
in_dim=3, num_frequencies=10, min_freq_exp=0.0, max_freq_exp=8.0, include_input=True
)
@ -87,7 +89,6 @@ class NerfWModel(Model):
transient_embedding_dim=self.transient_embedding_dim,
)
def populate_misc_modules(self):
# samplers
self.sampler_uniform = UniformSampler(num_samples=self.config.num_coarse_samples)
self.sampler_pdf = PDFSampler(num_samples=self.config.num_importance_samples)

10
nerfactory/models/semantic_nerf.py
View File

@ -41,8 +41,11 @@ class SemanticNerfModel(NeRFModel):
self.stuff_colors = semantics.stuff_colors
super().__init__(config=config, **kwargs)
def populate_fields(self):
"""Set the fields."""
def populate_modules(self):
"""Set the fields and modules"""
super().populate_modules()
# setting up fields
position_encoding = NeRFEncoding(
in_dim=3, num_frequencies=10, min_freq_exp=0.0, max_freq_exp=8.0, include_input=True
)
@ -56,9 +59,6 @@ class SemanticNerfModel(NeRFModel):
num_semantic_classes, position_encoding=position_encoding, direction_encoding=direction_encoding
)
def populate_misc_modules(self):
super().populate_misc_modules()
# renderers
self.renderer_semantic = SemanticRenderer()

14
nerfactory/models/tensorf.py
View File

@ -34,7 +34,6 @@ from nerfactory.fields.modules.field_heads import FieldHeadNames
from nerfactory.fields.nerf_field import NeRFField
from nerfactory.models.base import Model
from nerfactory.models.modules.ray_sampler import PDFSampler, UniformSampler
from nerfactory.models.modules.scene_colliders import NearFarCollider
from nerfactory.optimizers.loss import L1Loss, MSELoss
from nerfactory.optimizers.optimizers import Optimizers
from nerfactory.renderers.renderers import (
@ -111,8 +110,11 @@ class TensoRFModel(Model):
]
return callbacks
def populate_misc_modules(self):
# fields
def populate_modules(self):
"""Set the fields and modules"""
super().populate_modules()
# setting up fields
position_encoding = TensorVMEncoding(
resolution=self.init_resolution,
num_components=96,
@ -149,12 +151,6 @@ class TensoRFModel(Model):
self.ssim = structural_similarity_index_measure
self.lpips = LearnedPerceptualImagePatchSimilarity()
# colliders
if self.config.enable_collider:
self.collider = NearFarCollider(
near_plane=self.config.collider_params["near_plane"], far_plane=self.config.collider_params["far_plane"]
)
def get_param_groups(self) -> Dict[str, List[Parameter]]:
param_groups = {}
param_groups["fields"] = list(self.field.parameters())

14
nerfactory/models/vanilla_nerf.py
View File

@ -33,7 +33,6 @@ from nerfactory.fields.modules.field_heads import FieldHeadNames
from nerfactory.fields.nerf_field import NeRFField
from nerfactory.models.base import Model
from nerfactory.models.modules.ray_sampler import PDFSampler, UniformSampler
from nerfactory.models.modules.scene_colliders import NearFarCollider
from nerfactory.optimizers.loss import MSELoss
from nerfactory.renderers.renderers import (
AccumulationRenderer,
@ -90,9 +89,11 @@ class NeRFModel(Model):
]
return callbacks
def populate_fields(self):
"""Set the fields."""
def populate_modules(self):
"""Set the fields and modules"""
super().populate_modules()
# fields
position_encoding = NeRFEncoding(
in_dim=3, num_frequencies=10, min_freq_exp=0.0, max_freq_exp=8.0, include_input=True
)
@ -103,7 +104,6 @@ class NeRFModel(Model):
self.field_coarse = NeRFField(position_encoding=position_encoding, direction_encoding=direction_encoding)
self.field_fine = NeRFField(position_encoding=position_encoding, direction_encoding=direction_encoding)
def populate_misc_modules(self):
# samplers
self.sampler_uniform = UniformSampler(
num_samples=self.config.num_coarse_samples, density_field=self.density_field
@ -123,12 +123,6 @@ class NeRFModel(Model):
self.ssim = structural_similarity_index_measure
self.lpips = LearnedPerceptualImagePatchSimilarity()
# colliders
if self.config.enable_collider:
self.collider = NearFarCollider(
near_plane=self.config.collider_params["near_plane"], far_plane=self.config.collider_params["far_plane"]
)
def get_param_groups(self) -> Dict[str, List[Parameter]]:
param_groups = {}
if self.field_coarse is None or self.field_fine is None: