diff --git a/nerfactory/configs/base_configs.py b/nerfactory/configs/base_configs.py index 00b49eddc..6a661c91b 100644 --- a/nerfactory/configs/base_configs.py +++ b/nerfactory/configs/base_configs.py @@ -92,7 +92,7 @@ base_configs["proposal"] = Config( trainer=TrainerConfig(steps_per_eval_batch=500, steps_per_save=2000, mixed_precision=True), pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig( - dataparser=BlenderDataParserConfig(), train_num_rays_per_batch=8192, eval_num_rays_per_batch=8192 + dataparser=BlenderDataParserConfig(), train_num_rays_per_batch=4096, eval_num_rays_per_batch=8192 ), model=ProposalModelConfig(eval_num_rays_per_chunk=8192), ), @@ -102,7 +102,7 @@ base_configs["proposal"] = Config( "scheduler": None, } }, - viewer=ViewerConfig(enable=True), + viewer=ViewerConfig(enable=True, num_rays_per_chunk=2 << 15), logging=LoggingConfig(event_writer="none"), ) diff --git a/nerfactory/fields/compound_field.py b/nerfactory/fields/compound_field.py index 712eb15e8..41ae5adc9 100644 --- a/nerfactory/fields/compound_field.py +++ b/nerfactory/fields/compound_field.py @@ -139,9 +139,12 @@ class TCNNCompoundField(Field): def get_density(self, ray_samples: RaySamples): """Computes and returns the densities.""" - # TODO: add spatial distortion to ngp sampler - # positions = self.spatial_distortion(ray_samples.frustums.get_positions()) - positions = SceneBounds.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) + if self.spatial_distortion is not None: + positions = ray_samples.frustums.get_positions() + positions = self.spatial_distortion(positions) + positions = (positions + 2.0) / 4.0 + else: + positions = SceneBounds.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) positions_flat = positions.view(-1, 3) h = self.mlp_base(positions_flat).view(*ray_samples.frustums.shape, -1) density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1) @@ -213,7 +216,6 @@ class TorchCompoundField(Field): self.position_encoding = position_encoding self.direction_encoding = direction_encoding - self.spatial_distortion = spatial_distortion self.mlp_base = MLP( in_dim=self.position_encoding.get_out_dim(), @@ -234,9 +236,11 @@ class TorchCompoundField(Field): field_head.set_in_dim(self.mlp_head.get_out_dim()) # type: ignore def get_density(self, ray_samples: RaySamples): - # TODO: add spatial distortion to ngp sampler - # positions = self.spatial_distortion(ray_samples.frustums.get_positions()) - positions = ray_samples.frustums.get_positions() + if self.spatial_distortion is not None: + positions = ray_samples.frustums.get_positions() + positions = self.spatial_distortion(positions) + else: + positions = ray_samples.frustums.get_positions() encoded_xyz = self.position_encoding(positions) base_mlp_out = self.mlp_base(encoded_xyz) density = self.field_output_density(base_mlp_out) diff --git a/nerfactory/models/proposal.py b/nerfactory/models/proposal.py index 1aa632e1f..a99314ffa 100644 --- a/nerfactory/models/proposal.py +++ b/nerfactory/models/proposal.py @@ -28,6 +28,7 @@ from torchmetrics.functional import structural_similarity_index_measure from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity from nerfactory.cameras.rays import RayBundle +from nerfactory.fields.compound_field import TCNNCompoundField from nerfactory.fields.density_field import DensityField from nerfactory.fields.instant_ngp_field import TCNNInstantNGPField from nerfactory.fields.modules.field_heads import FieldHeadNames @@ -65,6 +66,8 @@ class ProposalModelConfig(ModelConfig): """Proposal loss multiplier.""" distortion_loss_mult: float = 0.01 """Distortion loss multiplier.""" + use_appearance_conditioning: bool = True + """Whether to use appearance conditioning.""" class ProposalModel(Model): @@ -81,7 +84,12 @@ class ProposalModel(Model): super().populate_modules() # Fields - self.field = TCNNInstantNGPField(self.scene_bounds.aabb, spatial_distortion=SceneContraction()) + if self.config.use_appearance_conditioning: + self.field = TCNNCompoundField( + self.scene_bounds.aabb, spatial_distortion=SceneContraction(), num_images=self.num_train_data + ) + else: + self.field = TCNNInstantNGPField(self.scene_bounds.aabb, spatial_distortion=SceneContraction()) self.proposal_network = DensityField(self.scene_bounds.aabb, spatial_distortion=SceneContraction()) # Collider