From dcee1087ebd302b5a569147c773d101c3f577a6d Mon Sep 17 00:00:00 2001
From: terrance <terrance.3512@gmail.com>
Date: Fri, 16 Dec 2022 19:54:39 +0000
Subject: [PATCH] change position encoding to direction encoding for background
 color mlp

---
 nerfstudio/fields/dreamfusion_field.py | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/nerfstudio/fields/dreamfusion_field.py b/nerfstudio/fields/dreamfusion_field.py
index ca5afc803..a89d33a2c 100644
--- a/nerfstudio/fields/dreamfusion_field.py
+++ b/nerfstudio/fields/dreamfusion_field.py
@@ -43,7 +43,7 @@ from nerfstudio.field_components.spatial_distortions import (
     SpatialDistortion,
 )
 from nerfstudio.fields.base_field import Field
-from nerfstudio.utils import colormaps, colors, math
+from nerfstudio.utils import math
 
 try:
     import tinycudann as tcnn
@@ -104,6 +104,14 @@ class DreamFusionField(Field):
             encoding_config={"otype": "Frequency", "n_frequencies": 2},
         )
 
+        self.direction_encoding = tcnn.Encoding(
+            n_input_dims=3,
+            encoding_config={
+                "otype": "SphericalHarmonics",
+                "degree": 4,
+            },
+        )
+
         self.mlp_base = tcnn.NetworkWithInputEncoding(
             n_input_dims=3,
             n_output_dims=1 + self.geo_feat_dim,
@@ -139,7 +147,7 @@ class DreamFusionField(Field):
         self.field_head_pred_normals = PredNormalsFieldHead(in_dim=self.mlp_pred_normals.n_output_dims)
 
         self.mlp_background_color = tcnn.Network(
-            n_input_dims=self.position_encoding.n_output_dims,
+            n_input_dims=self.direction_encoding.n_output_dims,
             n_output_dims=3,
             network_config={
                 "otype": "FullyFusedMLP",
@@ -189,9 +197,9 @@ class DreamFusionField(Field):
         directions = get_normalized_directions(ray_bundle.directions)
 
         outputs_shape = ray_bundle.directions.shape[:-1]
+        directions_flat = self.direction_encoding(directions.view(-1, 3))
+        background_rgb = self.mlp_background_color(directions_flat).view(*outputs_shape, -1).to(directions)
 
-        positions_flat = self.position_encoding(directions.view(-1, 3))
-        background_rgb = self.mlp_background_color(positions_flat).view(*outputs_shape, -1).to(directions)
         return background_rgb
 
     def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None):