option for using black bg in test dataset

nerfstudio-project · Apr 9, 2023 · 658df40 · 658df40
1 parent 0426623
commit 658df40
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/examples/datasets/nerf_360_v2.py b/examples/datasets/nerf_360_v2.py
@@ -281,7 +281,10 @@ def preprocess(self, data):
                 color_bkgd = torch.zeros(3, device=self.images.device)
         else:
             # just use white during inference
-            color_bkgd = torch.ones(3, device=self.images.device)
+            if self.color_bkgd_aug == "black":
+                color_bkgd = torch.zeros(3, device=self.images.device)
+            else:
+                color_bkgd = torch.ones(3, device=self.images.device)
 
         return {
             "pixels": pixels,  # [n_rays, 3] or [h, w, 3]