make robust and patch_size model params (#1386)

* make robust and patch_size model params

* change defaults back, turns off robust by default

* Make nerfacto-robust a separate config using same model

nerfstudio/configs/method_configs.py

@@ -80,6 +80,40 @@ method_configs["nerfacto"] = TrainerConfig(
     steps_per_save=2000,
     max_num_iterations=30000,
     mixed_precision=True,
+    pipeline=VanillaPipelineConfig(
+        datamanager=VanillaDataManagerConfig(
+            dataparser=NerfstudioDataParserConfig(),
+            train_num_rays_per_batch=4096,
+            eval_num_rays_per_batch=4096,
+            camera_optimizer=CameraOptimizerConfig(
+                mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
+            ),
+        ),
+        model=NerfactoModelConfig(
+            eval_num_rays_per_chunk=1 << 15,
+            robust=False,
+        ),
+    ),
+    optimizers={
+        "proposal_networks": {
+            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
+            "scheduler": None,
+        },
+        "fields": {
+            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
+            "scheduler": None,
+        },
+    },
+    viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
+    vis="viewer",
+)
+
+method_configs["nerfacto-robust"] = TrainerConfig(
+    method_name="nerfacto-robust",
+    steps_per_eval_batch=500,
+    steps_per_save=2000,
+    max_num_iterations=30000,
+    mixed_precision=True,
     pipeline=VanillaPipelineConfig(
         datamanager=VanillaDataManagerConfig(
             dataparser=NerfstudioDataParserConfig(),
@@ -89,7 +123,11 @@ method_configs["nerfacto"] = TrainerConfig(
                 mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
             ),
         ),
-        model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 15),
+        model=NerfactoModelConfig(
+            eval_num_rays_per_chunk=1 << 15,
+            robust=True,
+            patch_size=16
+        ),
     ),
     optimizers={
         "proposal_networks": {

nerfstudio/model_components/losses.py

@@ -303,7 +303,7 @@ def depth_loss(
 
     raise NotImplementedError("Provided depth loss type not implemented.")
 
-def robust_rgb_loss(rgb_gt,rgb_outputs,patch_size=16):
+def robust_rgb_loss(rgb_gt,rgb_outputs,patch_size):
     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

nerfstudio/models/nerfacto.py

@@ -116,6 +116,10 @@ class NerfactoModelConfig(ModelConfig):
     """Whether use single jitter or not for the proposal networks."""
     predict_normals: bool = False
     """Whether to predict normals or not."""
+    robust: bool = False
+    """Whether to use robust RGB loss or not."""
+    patch_size: int = 16
+    """Patch size to use for robust RGB loss."""
 
 
 class NerfactoModel(Model):
@@ -301,7 +305,10 @@ class NerfactoModel(Model):
     def get_loss_dict(self, outputs, batch, metrics_dict=None):
         loss_dict = {}
         image = batch["image"].to(self.device)
-        loss_dict["rgb_loss"] = robust_rgb_loss(image, outputs["rgb"])
+        if self.config.robust:
+            loss_dict["rgb_loss"] = robust_rgb_loss(image, outputs["rgb"], self.config.patch_size)
+        else:
+            loss_dict["rgb_loss"] = self.rgb_loss(image, outputs["rgb"])
         if self.training:
             loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss(
                 outputs["weights_list"], outputs["ray_samples_list"]