padding changes, some params

nerfstudio/configs/method_configs.py

@@ -83,7 +83,7 @@ method_configs["nerfacto"] = TrainerConfig(
     pipeline=VanillaPipelineConfig(
         datamanager=VanillaDataManagerConfig(
             dataparser=NerfstudioDataParserConfig(),
-            train_num_rays_per_batch=16000,
+            train_num_rays_per_batch=17000,
             eval_num_rays_per_batch=4096,
             camera_optimizer=CameraOptimizerConfig(
                 mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)

nerfstudio/data/datamanagers/base_datamanager.py

@@ -351,7 +351,7 @@ class VanillaDataManager(DataManager):  # pylint: disable=abstract-method
         # Otherwise, use the default pixel sampler
         if is_equirectangular.any():
             CONSOLE.print("[bold yellow]Warning: Some cameras are equirectangular, but using default pixel sampler.")
-        return PatchPixelSampler(*args, **kwargs,patch_size=8)
+        return PatchPixelSampler(*args, **kwargs,patch_size=16)
 
     def setup_train(self):
         """Sets up the data loaders for training"""

nerfstudio/model_components/losses.py

@@ -303,26 +303,28 @@ def depth_loss(
 
     raise NotImplementedError("Provided depth loss type not implemented.")
 
-def robust_rgb_loss(rgb_gt,rgb_outputs,patch_size=8):
+def robust_rgb_loss(rgb_gt,rgb_outputs,patch_size=16):
     out_patches = rgb_outputs.view(-1,1,patch_size,patch_size,3)
     gt_patches = rgb_gt.view(-1,1,patch_size,patch_size,3)
     device = rgb_outputs.device
     batch_size = out_patches.shape[0]
     residuals = torch.mean((out_patches - gt_patches)**2,dim=-1)
-    # with torch.no_grad():
-    med_residual = torch.median(residuals)
-    #equation 8
-    weight = (residuals<=med_residual).float()
-    #equation 9
-    blurred_w = (torch.nn.functional.conv2d(weight,(1/9.)*torch.ones((1,1,3,3),device=device),padding='same')>=0.5).float()
-    # expected_w = torch.mean(torch.mean(blurred_w,dim=-1),dim=-1)
-    #equation 10
-    # weight_r8 = (expected_w > 0.6).float()
-    # w_8x8 = weight_r8[...,None] * torch.ones((batch_size,8,8),device=device)
-    # final_w = torch.zeros((batch_size,patch_size,patch_size),device=device)
-    # final_w[:,
-    #         (patch_size//2-patch_size//4):(patch_size//2+patch_size//4),
-    #         (patch_size//2-patch_size//4):(patch_size//2+patch_size//4)] = w_8x8
-    # loss = torch.mean(residuals.squeeze()*blurred_w)
-    loss = torch.mean(residuals.squeeze()*blurred_w)
+    with torch.no_grad():
+        med_residual = torch.quantile(residuals,.9)
+        #equation 8
+        weight = (residuals<=med_residual).float()#B x patch x patch
+        #equation 9
+        weight = torch.nn.functional.pad(weight,(1,1,1,1),mode='replicate')
+        blurred_w = (torch.nn.functional.conv2d(weight,(1/9.)*torch.ones((1,1,3,3),device=device),
+                padding='valid')>=0.5).float()
+        #equation 10
+        expected_w = blurred_w.view(batch_size,-1).mean(1)
+        weight_r8 = (expected_w >= 0.6).float() #Bx1
+        #the paper uses the lines below and multiplies residuals by final_w, but I found that just using
+        #the value for all 16 pixels worked better.
+        # final_w = torch.zeros((batch_size,patch_size,patch_size),device=device)
+        # final_w[:,
+        #         (patch_size//2-patch_size//4):(patch_size//2+patch_size//4),
+        #         (patch_size//2-patch_size//4):(patch_size//2+patch_size//4)] = weight_r8[:,None,None]
+    # loss = torch.mean(residuals.squeeze()*weight_r8[:,None,None])
     return loss