fix: background net should condition on rays_d

nerf/network.py

@@ -52,7 +52,7 @@ class NeRFNetwork(NeRFRenderer):
         if self.bg_radius > 0:
             self.num_layers_bg = num_layers_bg   
             self.hidden_dim_bg = hidden_dim_bg
-            self.encoder_bg, self.in_dim_bg = get_encoder('frequency', input_dim=2)
+            self.encoder_bg, self.in_dim_bg = get_encoder('frequency', input_dim=3)
             self.bg_net = MLP(self.in_dim_bg, 3, hidden_dim_bg, num_layers_bg, bias=True)
             
         else:
@@ -80,7 +80,7 @@ class NeRFNetwork(NeRFRenderer):
         return sigma, albedo
     
     # ref: https://github.com/zhaofuq/Instant-NSR/blob/main/nerf/network_sdf.py#L192
-    def finite_difference_normal(self, x, epsilon=5e-4):
+    def finite_difference_normal(self, x, epsilon=1e-2):
         # x: [N, 3]
         dx_pos, _ = self.common_forward((x + torch.tensor([[epsilon, 0.00, 0.00]], device=x.device)).clamp(-self.bound, self.bound))
         dx_neg, _ = self.common_forward((x + torch.tensor([[-epsilon, 0.00, 0.00]], device=x.device)).clamp(-self.bound, self.bound))
@@ -148,10 +148,9 @@ class NeRFNetwork(NeRFRenderer):
         }
 
 
-    def background(self, x, d):
-        # x: [N, 2], in [-1, 1]
+    def background(self, d):
 
-        h = self.encoder_bg(x) # [N, C]
+        h = self.encoder_bg(d) # [N, C]
         
         h = self.bg_net(h)
 

nerf/network_grid.py

@@ -57,7 +57,7 @@ class NeRFNetwork(NeRFRenderer):
             
             # use a very simple network to avoid it learning the prompt...
             # self.encoder_bg, self.in_dim_bg = get_encoder('tiledgrid', input_dim=2, num_levels=4, desired_resolution=2048)
-            self.encoder_bg, self.in_dim_bg = get_encoder('frequency', input_dim=2)
+            self.encoder_bg, self.in_dim_bg = get_encoder('frequency', input_dim=3)
 
             self.bg_net = MLP(self.in_dim_bg, 3, hidden_dim_bg, num_layers_bg, bias=True)
             
@@ -87,7 +87,7 @@ class NeRFNetwork(NeRFRenderer):
         return sigma, albedo
     
     # ref: https://github.com/zhaofuq/Instant-NSR/blob/main/nerf/network_sdf.py#L192
-    def finite_difference_normal(self, x, epsilon=5e-4):
+    def finite_difference_normal(self, x, epsilon=1e-2):
         # x: [N, 3]
         dx_pos, _ = self.common_forward((x + torch.tensor([[epsilon, 0.00, 0.00]], device=x.device)).clamp(-self.bound, self.bound))
         dx_neg, _ = self.common_forward((x + torch.tensor([[-epsilon, 0.00, 0.00]], device=x.device)).clamp(-self.bound, self.bound))
@@ -155,10 +155,9 @@ class NeRFNetwork(NeRFRenderer):
         }
 
 
-    def background(self, x, d):
-        # x: [N, 2], in [-1, 1]
+    def background(self, d):
 
-        h = self.encoder_bg(x) # [N, C]
+        h = self.encoder_bg(d) # [N, C]
         
         h = self.bg_net(h)
 

nerf/network_tcnn.py

@@ -4,6 +4,7 @@ import torch.nn.functional as F
 
 from activation import trunc_exp
 from .renderer import NeRFRenderer
+from encoding import get_encoder
 
 import numpy as np
 import tinycudann as tcnn
@@ -65,19 +66,9 @@ class NeRFNetwork(NeRFRenderer):
             self.num_layers_bg = num_layers_bg   
             self.hidden_dim_bg = hidden_dim_bg
 
-            self.encoder_bg = tcnn.Encoding(
-                n_input_dims=2,
-                encoding_config={
-                    "otype": "HashGrid",
-                    "n_levels": 4,
-                    "n_features_per_level": 2,
-                    "log2_hashmap_size": 16,
-                    "base_resolution": 16,
-                    "per_level_scale": 1.5,
-                },
-            )
-
-            self.bg_net = MLP(8, 3, hidden_dim_bg, num_layers_bg, bias=True)
+            self.encoder_bg, self.in_dim_bg = get_encoder('frequency', input_dim=3)
+
+            self.bg_net = MLP(self.in_dim_bg, 3, hidden_dim_bg, num_layers_bg, bias=True)
             
         else:
             self.bg_net = None
@@ -156,11 +147,10 @@ class NeRFNetwork(NeRFRenderer):
         }
 
 
-    def background(self, x, d):
+    def background(self, d):
         # x: [N, 2], in [-1, 1]
 
-        h = (x + 1) / (2 * 1) # to [0, 1]
-        h = self.encoder_bg(h) # [N, C]
+        h = self.encoder_bg(d) # [N, C]
         
         h = self.bg_net(h)
 

nerf/renderer.py

@@ -420,8 +420,8 @@ class NeRFRenderer(nn.Module):
         # mix background color
         if self.bg_radius > 0:
             # use the bg model to calculate bg_color
-            sph = raymarching.sph_from_ray(rays_o, rays_d, self.bg_radius) # [N, 2] in [-1, 1]
-            bg_color = self.background(sph, rays_d.reshape(-1, 3)) # [N, 3]
+            # sph = raymarching.sph_from_ray(rays_o, rays_d, self.bg_radius) # [N, 2] in [-1, 1]
+            bg_color = self.background(rays_d.reshape(-1, 3)) # [N, 3]
         elif bg_color is None:
             bg_color = 1
             
@@ -526,8 +526,8 @@ class NeRFRenderer(nn.Module):
         if self.bg_radius > 0:
             
             # use the bg model to calculate bg_color
-            sph = raymarching.sph_from_ray(rays_o, rays_d, self.bg_radius) # [N, 2] in [-1, 1]
-            bg_color = self.background(sph, rays_d) # [N, 3]
+            # sph = raymarching.sph_from_ray(rays_o, rays_d, self.bg_radius) # [N, 2] in [-1, 1]
+            bg_color = self.background(rays_d) # [N, 3]
 
         elif bg_color is None:
             bg_color = 1

nerf/utils.py

@@ -343,6 +343,9 @@ class Trainer(object):
         pred_rgb = outputs['image'].reshape(B, H, W, 3).permute(0, 3, 1, 2).contiguous() # [1, 3, H, W]
         # torch.cuda.synchronize(); print(f'[TIME] nerf render {time.time() - _t:.4f}s')
         
+        # print(shading)
+        # torch_vis_2d(pred_rgb[0])
+        
         # text embeddings
         if self.opt.dir_text:
             dirs = data['dir'] # [B,]