diff --git a/nerfactory/models/base.py b/nerfactory/models/base.py index e5207d88a..a95eeea3e 100644 --- a/nerfactory/models/base.py +++ b/nerfactory/models/base.py @@ -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]]: diff --git a/nerfactory/models/instant_ngp.py b/nerfactory/models/instant_ngp.py index 398a32229..98f42baf4 100644 --- a/nerfactory/models/instant_ngp.py +++ b/nerfactory/models/instant_ngp.py @@ -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) diff --git a/nerfactory/models/mipnerf.py b/nerfactory/models/mipnerf.py index dc4bba8d6..bd7e4b428 100644 --- a/nerfactory/models/mipnerf.py +++ b/nerfactory/models/mipnerf.py @@ -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: diff --git a/nerfactory/models/mipnerf_360.py b/nerfactory/models/mipnerf_360.py index 8ea2cd086..1a138a85a 100644 --- a/nerfactory/models/mipnerf_360.py +++ b/nerfactory/models/mipnerf_360.py @@ -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: diff --git a/nerfactory/models/nerfw.py b/nerfactory/models/nerfw.py index afe0705cf..93206ef3d 100644 --- a/nerfactory/models/nerfw.py +++ b/nerfactory/models/nerfw.py @@ -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) diff --git a/nerfactory/models/semantic_nerf.py b/nerfactory/models/semantic_nerf.py index 09a1d0eed..b175d8711 100644 --- a/nerfactory/models/semantic_nerf.py +++ b/nerfactory/models/semantic_nerf.py @@ -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() diff --git a/nerfactory/models/tensorf.py b/nerfactory/models/tensorf.py index 98a36b1a1..b054c16c2 100644 --- a/nerfactory/models/tensorf.py +++ b/nerfactory/models/tensorf.py @@ -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()) diff --git a/nerfactory/models/vanilla_nerf.py b/nerfactory/models/vanilla_nerf.py index 0f5dd1548..a0abb5b2f 100644 --- a/nerfactory/models/vanilla_nerf.py +++ b/nerfactory/models/vanilla_nerf.py @@ -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: