Update NGP related models with nerfacc 0.5.2

nerfstudio/cameras/cameras.py

@@ -633,7 +633,7 @@ def _generate_rays_from_coords(
             if distortion_params is not None:
                 mask = (self.camera_type[true_indices] != CameraType.EQUIRECTANGULAR.value).squeeze(-1)  # (num_rays)
                 coord_mask = torch.stack([mask, mask, mask], dim=0)
-                if mask.any():
+                if mask.any() and (distortion_params != 0).any():
                     coord_stack[coord_mask, :] = camera_utils.radial_and_tangential_undistort(
                         coord_stack[coord_mask, :].reshape(3, -1, 2),
                         distortion_params[mask, :],

nerfstudio/configs/method_configs.py

@@ -21,7 +21,6 @@
 from typing import Dict
 
 import tyro
-from nerfacc import ContractionType
 
 from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig
 from nerfstudio.configs.base_config import ViewerConfig
@@ -235,16 +234,20 @@
     max_num_iterations=30000,
     mixed_precision=True,
     pipeline=DynamicBatchPipelineConfig(
-        datamanager=VanillaDataManagerConfig(dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=8192),
+        datamanager=VanillaDataManagerConfig(
+            dataparser=NerfstudioDataParserConfig(),
+            train_num_rays_per_batch=4096,
+            eval_num_rays_per_batch=4096,
+        ),
         model=InstantNGPModelConfig(eval_num_rays_per_chunk=8192),
     ),
     optimizers={
         "fields": {
             "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
-            "scheduler": None,
+            "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000),
         }
     },
-    viewer=ViewerConfig(num_rays_per_chunk=64000),
+    viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
     vis="viewer",
 )
 
@@ -259,20 +262,22 @@
         datamanager=VanillaDataManagerConfig(dataparser=InstantNGPDataParserConfig(), train_num_rays_per_batch=8192),
         model=InstantNGPModelConfig(
             eval_num_rays_per_chunk=8192,
-            contraction_type=ContractionType.AABB,
+            grid_levels=1,
+            alpha_thre=0.0,
+            cone_angle=0.0,
             render_step_size=0.001,
-            max_num_samples_per_ray=48,
+            disable_scene_contraction=True,
             near_plane=0.01,
             background_color="black",
         ),
     ),
     optimizers={
         "fields": {
             "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
-            "scheduler": None,
+            "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000),
         }
     },
-    viewer=ViewerConfig(num_rays_per_chunk=64000),
+    viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
     vis="viewer",
 )
 
@@ -470,9 +475,10 @@
         datamanager=DepthDataManagerConfig(dataparser=DycheckDataParserConfig(), train_num_rays_per_batch=8192),
         model=NerfplayerNGPModelConfig(
             eval_num_rays_per_chunk=8192,
-            contraction_type=ContractionType.AABB,
+            grid_levels=1,
+            alpha_thre=0.0,
             render_step_size=0.001,
-            max_num_samples_per_ray=48,
+            disable_scene_contraction=True,
             near_plane=0.01,
         ),
     ),

nerfstudio/fields/instant_ngp_field.py

@@ -21,7 +21,6 @@
 
 import numpy as np
 import torch
-from nerfacc import ContractionType, contract
 from torch.nn.parameter import Parameter
 from torchtyping import TensorType
 
@@ -30,6 +29,10 @@
 from nerfstudio.field_components.activations import trunc_exp
 from nerfstudio.field_components.embedding import Embedding
 from nerfstudio.field_components.field_heads import FieldHeadNames
+from nerfstudio.field_components.spatial_distortions import (
+    SceneContraction,
+    SpatialDistortion,
+)
 from nerfstudio.fields.base_field import Field, shift_directions_for_tcnn
 
 try:
@@ -69,16 +72,16 @@ def __init__(
         use_appearance_embedding: Optional[bool] = False,
         num_images: Optional[int] = None,
         appearance_embedding_dim: int = 32,
-        contraction_type: ContractionType = ContractionType.UN_BOUNDED_SPHERE,
         num_levels: int = 16,
         log2_hashmap_size: int = 19,
         max_res: int = 2048,
+        spatial_distortion: Optional[SpatialDistortion] = SceneContraction(),
     ) -> None:
         super().__init__()
 
         self.aabb = Parameter(aabb, requires_grad=False)
         self.geo_feat_dim = geo_feat_dim
-        self.contraction_type = contraction_type
+        self.spatial_distortion = spatial_distortion
 
         self.use_appearance_embedding = use_appearance_embedding
         if use_appearance_embedding:
@@ -88,7 +91,8 @@ def __init__(
 
         # TODO: set this properly based on the aabb
         base_res: int = 16
-        per_level_scale = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1))
+        features_per_level: int = 2
+        growth_factor = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1))
 
         self.direction_encoding = tcnn.Encoding(
             n_input_dims=3,
@@ -104,10 +108,10 @@ def __init__(
             encoding_config={
                 "otype": "HashGrid",
                 "n_levels": num_levels,
-                "n_features_per_level": 2,
+                "n_features_per_level": features_per_level,
                 "log2_hashmap_size": log2_hashmap_size,
                 "base_resolution": base_res,
-                "per_level_scale": per_level_scale,
+                "per_level_scale": growth_factor,
             },
             network_config={
                 "otype": "FullyFusedMLP",
@@ -134,9 +138,16 @@ def __init__(
         )
 
     def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType, TensorType]:
-        positions = ray_samples.frustums.get_positions()
+        if self.spatial_distortion is not None:
+            positions = ray_samples.frustums.get_positions()
+            positions = self.spatial_distortion(positions)
+            positions = (positions + 2.0) / 4.0
+        else:
+            positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)
+        # Make sure the tcnn gets inputs between 0 and 1.
+        selector = ((positions > 0.0) & (positions < 1.0)).all(dim=-1)
+        positions = positions * selector[..., None]
         positions_flat = positions.view(-1, 3)
-        positions_flat = contract(x=positions_flat, roi=self.aabb, type=self.contraction_type)
 
         h = self.mlp_base(positions_flat).view(*ray_samples.frustums.shape, -1)
         density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1)
@@ -145,20 +156,18 @@ def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType, TensorType]:
         # softplus, because it enables high post-activation (float32) density outputs
         # from smaller internal (float16) parameters.
         density = trunc_exp(density_before_activation.to(positions))
+        density = density * selector[..., None]
         return density, base_mlp_out
 
     def get_outputs(
         self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None
     ) -> Dict[FieldHeadNames, TensorType]:
+        assert density_embedding is not None
         directions = shift_directions_for_tcnn(ray_samples.frustums.directions)
         directions_flat = directions.view(-1, 3)
-
         d = self.direction_encoding(directions_flat)
-        if density_embedding is None:
-            positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)
-            h = torch.cat([d, positions.view(-1, 3)], dim=-1)
-        else:
-            h = torch.cat([d, density_embedding.view(-1, self.geo_feat_dim)], dim=-1)
+
+        h = torch.cat([d, density_embedding.view(-1, self.geo_feat_dim)], dim=-1)
 
         if self.use_appearance_embedding:
             if ray_samples.camera_indices is None:
@@ -183,20 +192,5 @@ def get_opacity(self, positions: TensorType["bs":..., 3], step_size) -> TensorTy
             step_size: the step size to use for the opacity evaluation.
         """
         density = self.density_fn(positions)
-        ## TODO: We should scale step size based on the distortion. Currently it uses too much memory.
-        # aabb_min, aabb_max = self.aabb[0], self.aabb[1]
-        # if self.contraction_type is not ContractionType.AABB:
-        #     x = (positions - aabb_min) / (aabb_max - aabb_min)
-        #     x = x * 2 - 1  # aabb is at [-1, 1]
-        #     mag = x.norm(dim=-1, keepdim=True)
-        #     mask = mag.squeeze(-1) > 1
-
-        #     dev = (2 * mag - 1) / mag**2 + 2 * x**2 * (1 / mag**3 - (2 * mag - 1) / mag**4)
-        #     dev[~mask] = 1.0
-        #     dev = torch.clamp(dev, min=1e-6)
-        #     step_size = step_size / dev.norm(dim=-1, keepdim=True)
-        # else:
-        #     step_size = step_size * (aabb_max - aabb_min)
-
         opacity = density * step_size
         return opacity

nerfstudio/fields/nerfacto_field.py

@@ -103,7 +103,7 @@ def __init__(
         pass_semantic_gradients: bool = False,
         use_pred_normals: bool = False,
         use_average_appearance_embedding: bool = False,
-        spatial_distortion: SpatialDistortion = None,
+        spatial_distortion: Optional[SpatialDistortion] = None,
     ) -> None:
         super().__init__()
 

nerfstudio/fields/nerfplayer_ngp_field.py

@@ -20,7 +20,6 @@
 from typing import Dict, Optional, Tuple
 
 import torch
-from nerfacc import ContractionType, contract
 from torch.nn.parameter import Parameter
 from torchtyping import TensorType
 
@@ -29,6 +28,10 @@
 from nerfstudio.field_components.activations import trunc_exp
 from nerfstudio.field_components.embedding import Embedding
 from nerfstudio.field_components.field_heads import FieldHeadNames
+from nerfstudio.field_components.spatial_distortions import (
+    SceneContraction,
+    SpatialDistortion,
+)
 from nerfstudio.field_components.temporal_grid import TemporalGridEncoder
 from nerfstudio.fields.base_field import Field, shift_directions_for_tcnn
 
@@ -82,13 +85,13 @@ def __init__(
         disable_viewing_dependent: bool = False,
         num_images: Optional[int] = None,
         appearance_embedding_dim: int = 32,
-        contraction_type: ContractionType = ContractionType.UN_BOUNDED_SPHERE,
+        spatial_distortion: Optional[SpatialDistortion] = SceneContraction(),
     ) -> None:
         super().__init__()
 
         self.aabb = Parameter(aabb, requires_grad=False)
         self.geo_feat_dim = geo_feat_dim
-        self.contraction_type = contraction_type
+        self.spatial_distortion = spatial_distortion
 
         self.use_appearance_embedding = use_appearance_embedding
         if use_appearance_embedding:
@@ -144,9 +147,16 @@ def __init__(
         )
 
     def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType, TensorType]:
-        positions = ray_samples.frustums.get_positions()
+        if self.spatial_distortion is not None:
+            positions = ray_samples.frustums.get_positions()
+            positions = self.spatial_distortion(positions)
+            positions = (positions + 2.0) / 4.0
+        else:
+            positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)
+        # Make sure the tcnn gets inputs between 0 and 1.
+        selector = ((positions > 0.0) & (positions < 1.0)).all(dim=-1)
+        positions = positions * selector[..., None]
         positions_flat = positions.view(-1, 3)
-        positions_flat = contract(x=positions_flat, roi=self.aabb, type=self.contraction_type)
         assert ray_samples.times is not None, "Time should be included in the input for NeRFPlayer"
         times_flat = ray_samples.times.view(-1, 1)
 
@@ -158,28 +168,22 @@ def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType, TensorType]:
         # softplus, because it enables high post-activation (float32) density outputs
         # from smaller internal (float16) parameters.
         density = trunc_exp(density_before_activation.to(positions))
+        density = density * selector[..., None]
         return density, base_mlp_out
 
     def get_outputs(
         self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None
     ) -> Dict[FieldHeadNames, TensorType]:
+        assert density_embedding is not None
         directions = shift_directions_for_tcnn(ray_samples.frustums.directions)
         directions_flat = directions.view(-1, 3)
 
         if self.direction_encoding is not None:
             d = self.direction_encoding(directions_flat)
-            if density_embedding is None:
-                positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)
-                h = torch.cat([d, positions.view(-1, 3)], dim=-1)
-            else:
-                h = torch.cat([d, density_embedding.view(-1, self.geo_feat_dim)], dim=-1)
+            h = torch.cat([d, density_embedding.view(-1, self.geo_feat_dim)], dim=-1)
         else:
             # viewing direction is disabled
-            if density_embedding is None:
-                positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)
-                h = positions.view(-1, 3)
-            else:
-                h = density_embedding.view(-1, self.geo_feat_dim)
+            h = density_embedding.view(-1, self.geo_feat_dim)
 
         if self.use_appearance_embedding:
             if ray_samples.camera_indices is None: