merging

nerfstudio/fields/dreamfusion_field.py

@@ -28,7 +28,6 @@ from nerfstudio.cameras.rays import RayBundle, RaySamples
 from nerfstudio.data.scene_box import SceneBox
 from nerfstudio.field_components.activations import trunc_exp
 from nerfstudio.field_components.field_heads import FieldHeadNames, PredNormalsFieldHead
-from nerfstudio.field_components.spatial_distortions import SpatialDistortion
 from nerfstudio.fields.base_field import Field
 from nerfstudio.utils import math
 
@@ -58,7 +57,6 @@ class DreamFusionField(Field):
         geo_feat_dim: output geo feat dimensions
         num_layers_color: number of hidden layers for color network
         hidden_dim_color: dimension of hidden layers for color network
-        spatial_distortion: spatial distortion to apply to the scene
     """
 
     def __init__(
@@ -69,16 +67,13 @@ class DreamFusionField(Field):
         geo_feat_dim: int = 15,
         num_layers_color: int = 3,
         hidden_dim_color: int = 64,
-        hidden_dim_transient: int = 64,
-        spatial_distortion: Optional[SpatialDistortion] = None,
+        hidden_dim_normal: int = 32,
     ) -> None:
         super().__init__()
 
         self.aabb = Parameter(aabb, requires_grad=False)
         self.geo_feat_dim = geo_feat_dim
 
-        self.spatial_distortion = spatial_distortion
-
         num_levels = 16
         max_res = 1024
         base_res = 16
@@ -122,7 +117,7 @@ class DreamFusionField(Field):
         # predicted normals
         self.mlp_pred_normals = tcnn.Network(
             n_input_dims=self.geo_feat_dim + self.position_encoding.n_output_dims,
-            n_output_dims=hidden_dim_transient,
+            n_output_dims=hidden_dim_normal,
             network_config={
                 "otype": "FullyFusedMLP",
                 "activation": "ReLU",
@@ -159,12 +154,7 @@ class DreamFusionField(Field):
 
     def get_density(self, ray_samples: RaySamples):
         """Computes and returns the densities."""
-        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 = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)
+        positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)
         self._sample_locations = positions
         if not self._sample_locations.requires_grad:
             self._sample_locations.requires_grad = True

nerfstudio/models/dreamfusion.py

@@ -344,8 +344,8 @@ class DreamFusionModel(Model):
             rgb=rgb, normals=pred_normals, light_direction=light_d, shading_weight=shading_weight, detach_normals=True
         )
 
-        outputs["normals"] = normals
-        outputs["pred_normals"] = pred_normals
+        outputs["normals"] = ((normals + 1.0) / 2.0) * accum_mask
+        outputs["pred_normals"] = ((pred_normals + 1.0) / 2.0) * accum_mask
         outputs["shaded"] = accum_mask * shaded + torch.ones_like(shaded) * accum_mask_inv
         outputs["shaded_albedo"] = shaded_albedo
         outputs["render"] = accum_mask * rgb + background

nerfstudio/pipelines/dreamfusion_pipeline.py

@@ -130,11 +130,13 @@ class DreamfusionPipeline(VanillaPipeline):
                 train_output.to(self.sd_device),
                 guidance_scale=int(self.config.guidance_scale),
             )
-            loss_dict["sds_loss"] = sds_loss.to(self.device)
+            sds_loss = sds_loss.to(self.device)
+
         # TODO: opacity penalty using transmittance, not accumultation
-        # if self.config.opacity_penalty:
-        #     accum_mean = np.mean(1.0 - accumulation)
-        #     sds_loss *= np.min((0.5, accum_mean))
+        if self.config.opacity_penalty:
+            accum_mean = torch.mean(1.0 - model_outputs["accumulation"])
+            sds_loss *= torch.clamp(accum_mean, min=0.5)
+        loss_dict["sds_loss"] = sds_loss.to(self.device)
 
         model_outputs["latents"] = latents
         model_outputs["grad"] = grad