From 11975e9a434d5713f2bebcc1eaba9e872877df17 Mon Sep 17 00:00:00 2001
From: Abhik Ahuja <ahujaabhik@berkeley.edu>
Date: Wed, 30 Nov 2022 15:14:16 -0800
Subject: [PATCH] Configurable light direction sampling

---
 nerfstudio/configs/method_configs.py | 3 ---
 nerfstudio/models/dreamfusion.py     | 7 ++++++-
 nerfstudio/models/nerfacto.py        | 4 ++--
 3 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py
index 957650053..742739269 100644
--- a/nerfstudio/configs/method_configs.py
+++ b/nerfstudio/configs/method_configs.py
@@ -261,9 +261,6 @@ method_configs["dreamfusion"] = Config(
             dataparser=NerfstudioDataParserConfig(),
             train_num_rays_per_batch=4096,
             eval_num_rays_per_batch=4096,
-            camera_optimizer=CameraOptimizerConfig(
-                mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
-            ),
         ),
         model=DreamFusionModelConfig(),
     ),
diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py
index 55868cc04..f5869c5b9 100644
--- a/nerfstudio/models/dreamfusion.py
+++ b/nerfstudio/models/dreamfusion.py
@@ -63,6 +63,8 @@ class DreamFusionModelConfig(ModelConfig):
     """Orientation loss multipier on computed normals."""
     pred_normal_loss_mult: float = 0.001
     """Predicted normal loss multiplier."""
+    random_light_source: bool = True
+    """Randomizes light source per output."""
 
 
 class DreamFusionModel(Model):
@@ -130,7 +132,10 @@ class DreamFusionModel(Model):
         pred_normals = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights)
 
         # lambertian shading
-        light_d = ray_bundle.origins[0] + torch.randn(3, dtype=torch.float).to(normals)
+        if self.config.random_light_source:
+            light_d = ray_bundle.origins[0] + torch.randn(3, dtype=torch.float).to(normals)
+        else:
+            light_d = ray_bundle.origins[0]
         light_d = math.safe_normalize(light_d)
 
         ratio = 0.1
diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py
index 15d139054..2e2354458 100644
--- a/nerfstudio/models/nerfacto.py
+++ b/nerfstudio/models/nerfacto.py
@@ -98,9 +98,9 @@ class NerfactoModelConfig(ModelConfig):
     """Proposal loss multiplier."""
     distortion_loss_mult: float = 0.002
     """Distortion loss multiplier."""
-    orientation_loss_mult: float = 0.0001
+    orientation_loss_mult: float = 0.01
     """Orientation loss multipier on computed normals."""
-    pred_normal_loss_mult: float = 0.001
+    pred_normal_loss_mult: float = 0.0003
     """Predicted normal loss multiplier."""
     use_proposal_weight_anneal: bool = True
     """Whether to use proposal weight annealing."""