From 44d6431c433b204de455ef77021e4bde777535da Mon Sep 17 00:00:00 2001
From: Abhik Ahuja <ahujaabhik@berkeley.edu>
Date: Mon, 28 Nov 2022 14:49:46 -0800
Subject: [PATCH] Normals in dreamfusion

---
 nerfstudio/fields/dreamfusion_field.py | 39 +++++++++++---------------
 nerfstudio/models/dreamfusion.py       | 14 +++++++--
 2 files changed, 29 insertions(+), 24 deletions(-)

diff --git a/nerfstudio/fields/dreamfusion_field.py b/nerfstudio/fields/dreamfusion_field.py
index 35b6ac5cb..a6ecc44fb 100644
--- a/nerfstudio/fields/dreamfusion_field.py
+++ b/nerfstudio/fields/dreamfusion_field.py
@@ -85,7 +85,6 @@ class DreamFusionField(Field):
         num_layers_color: int = 3,
         hidden_dim_color: int = 64,
         hidden_dim_transient: int = 64,
-        use_pred_normals: bool = False,
         spatial_distortion: Optional[SpatialDistortion] = None,
     ) -> None:
         super().__init__()
@@ -94,7 +93,6 @@ class DreamFusionField(Field):
         self.geo_feat_dim = geo_feat_dim
 
         self.spatial_distortion = spatial_distortion
-        self.use_pred_normals = use_pred_normals
 
         num_levels = 16
         max_res = 1024
@@ -129,19 +127,18 @@ class DreamFusionField(Field):
         )
 
         # predicted normals
-        if self.use_pred_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,
-                network_config={
-                    "otype": "FullyFusedMLP",
-                    "activation": "ReLU",
-                    "output_activation": "None",
-                    "n_neurons": 64,
-                    "n_hidden_layers": 2,
-                },
-            )
-            self.field_head_pred_normals = PredNormalsFieldHead(in_dim=self.mlp_pred_normals.n_output_dims)
+        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,
+            network_config={
+                "otype": "FullyFusedMLP",
+                "activation": "ReLU",
+                "output_activation": "None",
+                "n_neurons": 64,
+                "n_hidden_layers": 2,
+            },
+        )
+        self.field_head_pred_normals = PredNormalsFieldHead(in_dim=self.mlp_pred_normals.n_output_dims)
 
         self.mlp_head = tcnn.Network(
             n_input_dims=self.geo_feat_dim,
@@ -182,19 +179,17 @@ class DreamFusionField(Field):
         outputs = {}
 
         directions = get_normalized_directions(ray_samples.frustums.directions)
-        directions_flat = directions.view(-1, 3)
 
         outputs_shape = ray_samples.frustums.directions.shape[:-1]
 
         # predicted normals
-        if self.use_pred_normals:
-            positions = ray_samples.frustums.get_positions()
+        positions = ray_samples.frustums.get_positions()
 
-            positions_flat = self.position_encoding(positions.view(-1, 3))
-            pred_normals_inp = torch.cat([positions_flat, density_embedding.view(-1, self.geo_feat_dim)], dim=-1)
+        positions_flat = self.position_encoding(positions.view(-1, 3))
+        pred_normals_inp = torch.cat([positions_flat, density_embedding.view(-1, self.geo_feat_dim)], dim=-1)
 
-            x = self.mlp_pred_normals(pred_normals_inp).view(*outputs_shape, -1).to(directions)
-            outputs[FieldHeadNames.PRED_NORMALS] = self.field_head_pred_normals(x)
+        x = self.mlp_pred_normals(pred_normals_inp).view(*outputs_shape, -1).to(directions)
+        outputs[FieldHeadNames.PRED_NORMALS] = self.field_head_pred_normals(x)
 
         h = density_embedding.view(-1, self.geo_feat_dim)
 
diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py
index e2939ed19..9867ac014 100644
--- a/nerfstudio/models/dreamfusion.py
+++ b/nerfstudio/models/dreamfusion.py
@@ -38,6 +38,7 @@ from nerfstudio.model_components.ray_samplers import PDFSampler, UniformSampler
 from nerfstudio.model_components.renderers import (
     AccumulationRenderer,
     DepthRenderer,
+    NormalsRenderer,
     RGBRenderer,
 )
 from nerfstudio.model_components.scene_colliders import AABBBoxCollider
@@ -87,6 +88,7 @@ class DreamFusionModel(Model):
         self.renderer_rgb = RGBRenderer(background_color=colors.WHITE)
         self.renderer_accumulation = AccumulationRenderer()
         self.renderer_depth = DepthRenderer()
+        self.renderer_normals = NormalsRenderer()
 
         # losses
         self.rgb_loss = MSELoss()
@@ -108,14 +110,22 @@ class DreamFusionModel(Model):
     def get_outputs(self, ray_bundle: RayBundle):
         # uniform sampling
         ray_samples_uniform = self.sampler_uniform(ray_bundle)
-        field_outputs = self.field(ray_samples_uniform)
+        field_outputs = self.field(ray_samples_uniform, compute_normals=True)
         weights = ray_samples_uniform.get_weights(field_outputs[FieldHeadNames.DENSITY])
 
         accumulation = self.renderer_accumulation(weights)
         depth = self.renderer_depth(weights, ray_samples_uniform)
         rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)
+        normals = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights)
+        pred_normals = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights)
 
-        outputs = {"rgb": rgb, "accumulation": accumulation, "depth": depth}
+        outputs = {
+            "rgb": rgb,
+            "accumulation": accumulation,
+            "depth": depth,
+            "normals": normals,
+            "pred_normals": pred_normals,
+        }
         return outputs
 
     def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]: