Add appearance conditioning to proposal net (#544)

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"),
 )
 

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)

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