Supports Neuralangelo, bakedangelo, neus-facto-angelo (autonomousvision#134)

* empty cache for eval

* load dtu highres images

* add adamW

* add multi-step with warpup lr scheduler

* add background model for bakedsdf

* make bakedsdf's beta anneal configurable

* add accumulate_grad_steps

* add numerical gradients

* make hash encoding configurable

* make frequency position encoding optional

* add mask for hash encoding

* add neuralangelo and config

* add bakedangelo

* fix typo

* update default config of neuralangelo & bakedangelo

* fix multi-gpu leak

* disable nuscenes data parser

* shared proposal network for background model

* use all images + return transform and scale

* zip-nerf interlevel loss

* change to inerlevel_loss_zip

* use proposal network for background model

* disable uniform_sampler

* add beta annealing to neus_factor

* add neus-facto-angelo

* update default config for bakedangelo

* fix typo

* update defulat config of neus-facto-angelo

* updates

Author: niujinshuchong

README.md

@@ -10,6 +10,13 @@
 SDFStudio is a unified and modular framework for neural implicit surface reconstruction, built on top of the awesome nerfstudio project. We provide a unified implementation of three major implicit surface reconstruction methods: UniSurf, VolSDF, and NeuS. SDFStudio also supports various scene representions, such as MLPs, Tri-plane, and Multi-res. feature grids, and multiple point sampling strategies such as surface-guided sampling as in UniSurf, and Voxel-surface guided sampling from NeuralReconW. It further integrates recent advances in the area such as the utillization of monocular cues (MonoSDF), geometry regularization (UniSurf) and multi-view consistency (Geo-NeuS). Thanks to the unified and modular implementation, SDFStudio makes it easy to transfer ideas from one method to another. For example, Mono-NeuS applies the idea from MonoSDF to NeuS, and Geo-VolSDF applies the idea from Geo-NeuS to VolSDF.
 
 # Updates
+
+**2023.06.16**: Add `bakedangelo` which combines `BakedSDF` with numerical gridents and progressive training of `Neuralangelo`.
+
+**2023.06.16**: Add `neus-facto-angelo` which combines `neus-facto` with numerical gridents and progressive training of `Neuralangelo`.
+
+**2023.06.16**: Support [Neuralangelo](https://research.nvidia.com/labs/dir/neuralangelo/).
+
 **2023.03.12**: Support [BakedSDF](https://bakedsdf.github.io/).
 
 **2022.12.28**: Support [Neural RGB-D Surface Reconstruction](https://dazinovic.github.io/neural-rgbd-surface-reconstruction/).

nerfstudio/cameras/rays.py

@@ -128,6 +128,21 @@ class RaySamples(TensorDataclass):
     times: Optional[TensorType[..., 1]] = None
     """Times at which rays are sampled"""
 
+    def get_alphas(self, densities: TensorType[..., "num_samples", 1]) -> TensorType[..., "num_samples", 1]:
+        """Return weights based on predicted densities
+
+        Args:
+            densities: Predicted densities for samples along ray
+
+        Returns:
+            Weights for each sample
+        """
+
+        delta_density = self.deltas * densities
+        alphas = 1 - torch.exp(-delta_density)
+
+        return alphas
+
     def get_weights(self, densities: TensorType[..., "num_samples", 1]) -> TensorType[..., "num_samples", 1]:
         """Return weights based on predicted densities
 

nerfstudio/configs/base_config.py

@@ -161,6 +161,8 @@ class TrainerConfig(PrintableConfig):
     """Optionally specify model config to load from; if none, will use the default config?"""
     load_scheduler: bool = True
     """Whether to load the lr scheduler state_dict if exists"""
+    accumulate_grad_steps: int = 1
+    """Number of gradient steps to accumulate before taking an optimizer step."""
 
 
 # Viewer related configs

nerfstudio/configs/method_configs.py

@@ -45,11 +45,12 @@
     PhototourismDataParserConfig,
 )
 from nerfstudio.data.dataparsers.sdfstudio_dataparser import SDFStudioDataParserConfig
-from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig
+from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig, AdamWOptimizerConfig
 from nerfstudio.engine.schedulers import (
     ExponentialSchedulerConfig,
     MultiStepSchedulerConfig,
     NeuSSchedulerConfig,
+    MultiStepWarmupSchedulerConfig,
 )
 from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind
 from nerfstudio.fields.sdf_field import SDFFieldConfig
@@ -58,6 +59,8 @@
 from nerfstudio.models.instant_ngp import InstantNGPModelConfig
 from nerfstudio.models.mipnerf import MipNerfModel
 from nerfstudio.models.nerfacto import NerfactoModelConfig
+from nerfstudio.models.neuralangelo import NeuralangeloModelConfig
+from nerfstudio.models.bakedangelo import BakedAngeloModelConfig
 from nerfstudio.models.neuralreconW import NeuralReconWModelConfig
 from nerfstudio.models.neus import NeuSModelConfig
 from nerfstudio.models.neus_acc import NeuSAccModelConfig
@@ -99,8 +102,147 @@
     "neus-facto-bigmlp": "NeuS-facto with big MLP, it is used in training heritage data with 8 gpus",
     "bakedsdf": "Implementation of BackedSDF with multi-res hash grids",
     "bakedsdf-mlp": "Implementation of BackedSDF with large MLPs",
+    "neuralangelo": "Implementation of Neuralangelo",
+    "bakedangelo": "Implementation of Neuralangelo with BakedSDF",
+    "neus-facto-angelo": "Implementation of Neuralangelo with neus-facto",
 }
 
+
+method_configs["bakedangelo"] = Config(
+    method_name="bakedangelo",
+    trainer=TrainerConfig(
+        steps_per_eval_image=5000,
+        steps_per_eval_batch=5000,
+        steps_per_save=20000,
+        steps_per_eval_all_images=1000000,  # set to a very large model so we don't eval with all images
+        max_num_iterations=1000_001,
+        mixed_precision=False,
+    ),
+    pipeline=VanillaPipelineConfig(
+        datamanager=VanillaDataManagerConfig(
+            dataparser=SDFStudioDataParserConfig(),
+            train_num_rays_per_batch=8192,
+            eval_num_rays_per_batch=1024,
+            camera_optimizer=CameraOptimizerConfig(
+                mode="off", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
+            ),
+        ),
+        model=BakedAngeloModelConfig(
+            near_plane=0.01,
+            far_plane=1000.0,
+            overwrite_near_far_plane=True,
+            sdf_field=SDFFieldConfig(
+                use_grid_feature=True,
+                num_layers=1,
+                num_layers_color=4,
+                hidden_dim=256,
+                hidden_dim_color=256,
+                geometric_init=True,
+                bias=1.5,
+                beta_init=0.1,
+                inside_outside=True,
+                use_appearance_embedding=True,
+                use_numerical_gradients=True,
+                base_res=64,
+                max_res=4096,
+                log2_hashmap_size=22,
+                hash_features_per_level=8,
+                hash_smoothstep=False,
+                use_position_encoding=False,
+            ),
+            eikonal_loss_mult=0.01,
+            background_model="grid",
+            proposal_weights_anneal_max_num_iters=10000,
+            use_anneal_beta=True,
+            eval_num_rays_per_chunk=1024,
+            use_spatial_varying_eikonal_loss=False,
+            steps_per_level=10_000,
+            curvature_loss_warmup_steps=20_000,
+            beta_anneal_end=0.0002,
+            beta_anneal_max_num_iters=1000_000,
+        ),
+    ),
+    optimizers={
+        "proposal_networks": {
+            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
+            "scheduler": MultiStepSchedulerConfig(max_steps=1000_000),
+        },
+        "fields": {
+            "optimizer": AdamWOptimizerConfig(lr=1e-3, eps=1e-15, weight_decay=1e-2),
+            "scheduler": MultiStepWarmupSchedulerConfig(warm_up_end=5000, milestones=[600_000, 800_000], gamma=0.1),
+        },
+        "field_background": {
+            "optimizer": AdamWOptimizerConfig(lr=1e-3, eps=1e-15),
+            "scheduler": MultiStepWarmupSchedulerConfig(warm_up_end=5000, milestones=[300_000, 400_000], gamma=0.1),
+        },
+    },
+    viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
+    vis="viewer",
+)
+
+
+method_configs["neuralangelo"] = Config(
+    method_name="neuralangelo",
+    trainer=TrainerConfig(
+        steps_per_eval_image=5000,
+        steps_per_eval_batch=5000,
+        steps_per_save=20000,
+        steps_per_eval_all_images=1000000,  # set to a very large model so we don't eval with all images
+        max_num_iterations=500_001,
+        mixed_precision=False,
+    ),
+    pipeline=VanillaPipelineConfig(
+        datamanager=VanillaDataManagerConfig(
+            dataparser=SDFStudioDataParserConfig(),
+            train_num_rays_per_batch=512,
+            eval_num_rays_per_batch=512,
+            camera_optimizer=CameraOptimizerConfig(
+                mode="off", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
+            ),
+        ),
+        model=NeuralangeloModelConfig(
+            sdf_field=SDFFieldConfig(
+                use_grid_feature=True,
+                num_layers=1,
+                num_layers_color=4,
+                hidden_dim=256,
+                hidden_dim_color=256,
+                geometric_init=True,
+                bias=0.5,
+                beta_init=0.3,
+                inside_outside=False,
+                use_appearance_embedding=False,
+                position_encoding_max_degree=6,
+                use_numerical_gradients=True,
+                base_res=64,
+                max_res=4096,
+                log2_hashmap_size=22,
+                hash_features_per_level=8,
+                hash_smoothstep=False,
+                use_position_encoding=False,
+            ),
+            background_model="mlp",
+            enable_progressive_hash_encoding=True,
+            enable_curvature_loss_schedule=True,
+            enable_numerical_gradients_schedule=True,
+        ),
+    ),
+    optimizers={
+        "fields": {
+            "optimizer": AdamWOptimizerConfig(lr=1e-3, weight_decay=0.01, eps=1e-15),
+            # "scheduler": NeuSSchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=500000),
+            "scheduler": MultiStepWarmupSchedulerConfig(warm_up_end=5000, milestones=[300_000, 400_000], gamma=0.1),
+        },
+        "field_background": {
+            "optimizer": AdamWOptimizerConfig(lr=1e-3, eps=1e-15),
+            "scheduler": MultiStepWarmupSchedulerConfig(warm_up_end=5000, milestones=[300_000, 400_000], gamma=0.1),
+        },
+    },
+    viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
+    vis="viewer",
+)
+
+
 method_configs["bakedsdf"] = Config(
     method_name="bakedsdf",
     trainer=TrainerConfig(
@@ -228,6 +370,77 @@
 )
 
 
+method_configs["neus-facto-angelo"] = Config(
+    method_name="neus-facto-angelo",
+    trainer=TrainerConfig(
+        steps_per_eval_image=5000,
+        steps_per_eval_batch=5000,
+        steps_per_save=20000,
+        steps_per_eval_all_images=1000000,  # set to a very large model so we don't eval with all images
+        max_num_iterations=1000_001,
+        mixed_precision=False,
+    ),
+    pipeline=VanillaPipelineConfig(
+        datamanager=VanillaDataManagerConfig(
+            dataparser=SDFStudioDataParserConfig(),
+            train_num_rays_per_batch=2048,
+            eval_num_rays_per_batch=1024,
+            camera_optimizer=CameraOptimizerConfig(
+                mode="off", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
+            ),
+        ),
+        model=NeuSFactoModelConfig(
+            near_plane=0.01,
+            far_plane=1000.0,
+            overwrite_near_far_plane=True,
+            sdf_field=SDFFieldConfig(
+                use_grid_feature=True,
+                num_layers=1,
+                num_layers_color=4,
+                hidden_dim=256,
+                hidden_dim_color=256,
+                geometric_init=True,
+                bias=0.5,
+                beta_init=0.3,
+                inside_outside=False,
+                use_appearance_embedding=True,
+                use_numerical_gradients=True,
+                base_res=64,
+                max_res=4096,
+                log2_hashmap_size=22,
+                hash_features_per_level=8,
+                hash_smoothstep=False,
+                use_position_encoding=False,
+            ),
+            background_model="grid",
+            eval_num_rays_per_chunk=1024,
+            level_init=8,
+            eikonal_loss_mult=0.01,
+            use_anneal_beta=True,
+            enable_progressive_hash_encoding=True,
+            enable_numerical_gradients_schedule=True,
+            enable_curvature_loss_schedule=True,
+            curvature_loss_multi=5e-4,
+        ),
+    ),
+    optimizers={
+        "proposal_networks": {
+            "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
+            "scheduler": MultiStepSchedulerConfig(max_steps=1000_000),
+        },
+        "fields": {
+            "optimizer": AdamOptimizerConfig(lr=1e-3, eps=1e-15),
+            "scheduler": MultiStepWarmupSchedulerConfig(warm_up_end=5000, milestones=[600_000, 800_000], gamma=0.1),
+        },
+        "field_background": {
+            "optimizer": AdamWOptimizerConfig(lr=1e-3, eps=1e-15),
+            "scheduler": MultiStepWarmupSchedulerConfig(warm_up_end=5000, milestones=[300_000, 400_000], gamma=0.1),
+        },
+    },
+    viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
+    vis="viewer",
+)
+
 method_configs["neus-facto"] = Config(
     method_name="neus-facto",
     trainer=TrainerConfig(

nerfstudio/data/datamanagers/base_datamanager.py

@@ -45,7 +45,7 @@
 from nerfstudio.data.dataparsers.mipnerf360_dataparser import Mipnerf360DataParserConfig
 from nerfstudio.data.dataparsers.monosdf_dataparser import MonoSDFDataParserConfig
 from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig
-from nerfstudio.data.dataparsers.nuscenes_dataparser import NuScenesDataParserConfig
+# from nerfstudio.data.dataparsers.nuscenes_dataparser import NuScenesDataParserConfig
 from nerfstudio.data.dataparsers.phototourism_dataparser import (
     PhototourismDataParserConfig,
 )
@@ -74,7 +74,7 @@
             "blender-data": BlenderDataParserConfig(),
             "friends-data": FriendsDataParserConfig(),
             "instant-ngp-data": InstantNGPDataParserConfig(),
-            "nuscenes-data": NuScenesDataParserConfig(),
+            # "nuscenes-data": NuScenesDataParserConfig(),
             "record3d-data": Record3DDataParserConfig(),
             "dnerf-data": DNeRFDataParserConfig(),
             "phototourism-data": PhototourismDataParserConfig(),

nerfstudio/data/dataparsers/nerfstudio_dataparser.py

@@ -62,6 +62,8 @@ class NerfstudioDataParserConfig(DataParserConfig):
     """Whether to automatically scale the poses to fit in +/- 1 bounding box."""
     train_split_percentage: float = 0.9
     """The percent of images to use for training. The remaining images are for eval."""
+    use_all_train_images: bool = False
+    """Whether to use all images for training. If True, all images are used for training."""
 
 
 @dataclass
@@ -169,6 +171,9 @@ def _generate_dataparser_outputs(self, split="train"):
         assert len(i_eval) == num_eval_images
         if split == "train":
             indices = i_train
+            if self.config.use_all_train_images:
+                indices = i_all
+                num_train_images = num_images
         elif split in ["val", "test"]:
             indices = i_eval
         else:
@@ -181,7 +186,7 @@ def _generate_dataparser_outputs(self, split="train"):
             orientation_method = self.config.orientation_method
 
         poses = torch.from_numpy(np.array(poses).astype(np.float32))
-        poses, _ = camera_utils.auto_orient_and_center_poses(
+        poses, transform_matrix = camera_utils.auto_orient_and_center_poses(
             poses,
             method=orientation_method,
             center_poses=self.config.center_poses,
@@ -190,9 +195,10 @@ def _generate_dataparser_outputs(self, split="train"):
         # Scale poses
         scale_factor = 1.0
         if self.config.auto_scale_poses:
-            scale_factor /= torch.max(torch.abs(poses[:, :3, 3]))
+            scale_factor /= float(torch.max(torch.abs(poses[:, :3, 3])))
+        scale_factor *= self.config.scale_factor
 
-        poses[:, :3, 3] *= scale_factor * self.config.scale_factor
+        poses[:, :3, 3] *= scale_factor
 
         # Choose image_filenames and poses based on split, but after auto orient and scaling the poses.
         image_filenames = [image_filenames[i] for i in indices]
@@ -247,11 +253,21 @@ def _generate_dataparser_outputs(self, split="train"):
         assert self.downscale_factor is not None
         cameras.rescale_output_resolution(scaling_factor=1.0 / self.downscale_factor)
 
+        if "applied_transform" in meta:
+            applied_transform = torch.tensor(meta["applied_transform"], dtype=transform_matrix.dtype)
+            transform_matrix = transform_matrix @ torch.cat(
+                [applied_transform, torch.tensor([[0, 0, 0, 1]], dtype=transform_matrix.dtype)], 0
+            )
+        if "applied_scale" in meta:
+            applied_scale = float(meta["applied_scale"])
+            scale_factor *= applied_scale
+        
         dataparser_outputs = DataparserOutputs(
             image_filenames=image_filenames,
             cameras=cameras,
             scene_box=scene_box,
             mask_filenames=mask_filenames if len(mask_filenames) > 0 else None,
+            metadata={"transform": transform_matrix, "scale_factor": scale_factor},
         )
         return dataparser_outputs
 

nerfstudio/data/dataparsers/sdfstudio_dataparser.py

@@ -167,6 +167,9 @@ class SDFStudioDataParserConfig(DataParserConfig):
     train_val_no_overlap: bool = False
     """remove selected / sampled validation images from training set"""
     auto_orient: bool = False
+    """automatically orient the scene such that the up direction is the same as the viewer's up direction"""
+    load_dtu_highres: bool = False
+    """load high resolution images from DTU dataset, should only be used for the preprocessed DTU dataset"""
 
 
 @dataclass
@@ -216,6 +219,14 @@ def _generate_dataparser_outputs(self, split="train"):  # pylint: disable=unused
             cy.append(intrinsics[1, 2])
             camera_to_worlds.append(camtoworld)
 
+            # here is hard coded for DTU high-res images
+            if self.config.load_dtu_highres:
+                image_filename = self.config.data / "image" / frame["rgb_path"].replace("_rgb", "")
+                intrinsics[:2, :] *= 1200 / 384.0
+                intrinsics[0, 2] += 200
+                height, width = 1200, 1600
+                meta["height"], meta["width"] = height, width
+            
             if self.config.include_mono_prior:
                 assert meta["has_mono_prior"]
                 # load mono depth