Merge pull request #133 from Sarasra/main

update robustnerf into a trainable state with a robustnerf config.

Author: jonbarron

configs/360_robustnerf.gin

@@ -0,0 +1,24 @@
+Config.dataset_loader = 'llff'
+Config.near = 0.2
+Config.far = 1e6
+Config.factor = 4
+
+Config.patch_size = 16
+Config.data_loss_type = 'robustnerf'
+Config.robustnerf_inlier_quantile = 0.8
+Config.enable_robustnerf_loss = True
+
+Model.raydist_fn = @jnp.reciprocal
+Model.opaque_background = True
+
+PropMLP.warp_fn = @coord.contract
+PropMLP.net_depth = 4
+PropMLP.net_width = 256
+PropMLP.disable_density_normals = True
+PropMLP.disable_rgb = True
+
+NerfMLP.warp_fn = @coord.contract
+NerfMLP.net_depth = 8
+NerfMLP.net_width = 1024
+NerfMLP.disable_density_normals = True
+

internal/configs.py

@@ -92,12 +92,12 @@ class Config:
   orientation_loss_mult: float = 0.0  # Multiplier on the orientation loss.
   orientation_coarse_loss_mult: float = 0.0  # Coarser orientation loss weights.
   # RobustNerf loss hyperparameters
-  robustnerf_inlier_quantile = 0.5
-  enable_robutnerf_loss = False
-  robustnerf_inner_patch_size = 8
-  robustnerf_smoothed_filter_size = 3
-  robustnerf_smoothed_inlier_quantile = 0.5
-  robustnerf_inner_patch_inlier_quantile = 0.5
+  robustnerf_inlier_quantile: float = 0.5
+  enable_robustnerf_loss: bool = False
+  robustnerf_inner_patch_size: int = 8
+  robustnerf_smoothed_filter_size: int = 3
+  robustnerf_smoothed_inlier_quantile: float = 0.5
+  robustnerf_inner_patch_inlier_quantile: float = 0.5
   # What that loss is imposed on, options are 'normals' or 'normals_pred'.
   orientation_loss_target: str = 'normals_pred'
   predicted_normal_loss_mult: float = 0.0  # Mult. on the predicted normal loss.

internal/train_utils.py

@@ -104,7 +104,7 @@ def compute_data_loss(batch, renderings, rays, loss_threshold, config):
     elif config.data_loss_type == 'robustnerf':
         mask, robust_stats = robustnerf.robustnerf_mask(resid_sq, loss_threshold, 
                 config)
-        data_loss = data_loss * mask
+        data_loss = resid_sq * mask
         stats.update(robust_stats)
     else:
       assert False

train.py

@@ -125,7 +125,8 @@ def main(unused_argv):
         train_frac,
         loss_threshold,
     )
-    loss_threshold = jnp.mean(stats['loss_threshold'])
+    if config.enable_robustnerf_loss:
+        loss_threshold = jnp.mean(stats['loss_threshold'])
 
     if step % config.gc_every == 0:
       gc.collect()  # Disable automatic garbage collection for efficiency.