Shading hyperparameter adjustements

nerfstudio/models/dreamfusion.py

@@ -93,6 +93,7 @@ class DreamFusionModel(Model):
         self.num_samples = config.num_samples
         self.initialize_density = config.initialize_density
         self.train_normals = False
+        self.train_shaded = False
         super().__init__(config=config, **kwargs)
 
     def get_training_callbacks(
@@ -110,6 +111,11 @@ class DreamFusionModel(Model):
         ):
             self.train_normals = True
 
+        def start_shaded_training(
+            self, training_callback_attributes: TrainingCallbackAttributes, step: int  # pylint: disable=unused-argument
+        ):
+            self.train_shaded = True
+
         callbacks = [
             TrainingCallback(
                 where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],
@@ -123,6 +129,12 @@ class DreamFusionModel(Model):
                 func=start_training_normals,
                 args=[self, training_callback_attributes],
             ),
+            TrainingCallback(
+                where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],
+                iters=(1000,),
+                func=start_shaded_training,
+                args=[self, training_callback_attributes],
+            ),
         ]
         return callbacks
 
@@ -154,7 +166,7 @@ class DreamFusionModel(Model):
 
         # colliders
         if self.config.sphere_collider:
-            self.collider = SphereCollider(torch.Tensor([0, 0, 0]), 0.8)
+            self.collider = SphereCollider(torch.Tensor([0, 0, 0]), 1.0)
         else:
             self.collider = AABBBoxCollider(scene_box=self.scene_box)
 
@@ -206,7 +218,10 @@ class DreamFusionModel(Model):
             light_d = ray_bundle.origins[0]
         light_d = math.safe_normalize(light_d)
 
-        ratio = 0.3
+        if self.train_shaded:
+            ratio = 0.1
+        else:
+            ratio = 0.0
         # with torch.no_grad():
         lambertian = ratio + (1 - ratio) * (pred_normals.detach() @ light_d).clamp(min=0)
         shaded = lambertian.unsqueeze(-1).repeat(1, 3)
@@ -218,6 +233,14 @@ class DreamFusionModel(Model):
         outputs["shaded_albedo"] = shaded_albedo
         outputs["render"] = accum_mask * shaded_albedo + background
 
+        if ratio > 0:
+            if np.random.random_sample() > 0.5:
+                outputs["train_output"] = accum_mask * shaded + background
+            else:
+                outputs["train_output"] = accum_mask * shaded_albedo + background
+        else:
+            outputs["train_output"] = accum_mask * rgb + background
+
         # if self.train_normals:
 
         outputs["rendered_orientation_loss"] = orientation_loss(

nerfstudio/pipelines/dreamfusion_pipeline.py

@@ -56,7 +56,7 @@ class DreamfusionPipeline(VanillaPipeline):
         model_outputs = self.model(ray_bundle)
 
         # Just uses albedo for now
-        albedo_output = model_outputs["render"].view(1, 64, 64, 3).permute(0, 3, 1, 2)
+        albedo_output = model_outputs["train_output"].view(1, 64, 64, 3).permute(0, 3, 1, 2)
 
         accumulation = model_outputs["accumulation"].view(64, 64).detach().cpu().numpy()
         accumulation = np.clip(accumulation, 0.0, 1.0)
@@ -71,6 +71,9 @@ class DreamfusionPipeline(VanillaPipeline):
         plt.imsave("nerf_textureless.jpg", shaded)
 
         sds_loss, latents, grad = self.sd.sds_loss(self.text_embedding, albedo_output)
+        # TODO: opacity penalty using transmittance, not accumultation
+        accum_mean = np.mean(1.0 - accumulation)
+        sds_loss *= np.min((0.8, accum_mean))
 
         grad_scaler.scale(latents).backward(gradient=grad, retain_graph=True)
         # optimizers.scheduler_step_all(step)