feat(fields): Improve Type Hints (#1524)

* feat(fields/tensorf_fields): Improve Type Hints

* feat(fields/base_field): Improve Type Hints

* feat(fields/density_fields): Improve Type Hints

* feat(fields/instant_ngp_field): Improve Type Hints

* feat(fields/nerfacto_field): Improve Type Hints

* feat(fields/nerplayer_nerfacto_field): Improve Type Hints

* fix(fields/nerplayer_nerfacto_field): add typing imports

* feat(fields/nerfplayer_ngp_field): Improve Type Hints

* feat(fields/nerfw_field): Improve Type Hints

* feat(fields/semantic_nerf_field): Improve Type Hints

* feat(fields/vanilla_nerf_field): Improve Type Hints

* fix(fields/nerfplayer/nerfacto_field): reorder imports

* fix(feilds): str -> FieldHeadNames

* fix: remove Optionals

* fix: add missing comma

nerfstudio/fields/base_field.py

@@ -96,7 +96,7 @@ class Field(nn.Module):
             density_embedding: Density embeddings to condition on.
         """
 
-    def forward(self, ray_samples: RaySamples, compute_normals: bool = False):
+    def forward(self, ray_samples: RaySamples, compute_normals: bool = False) -> Dict[FieldHeadNames, TensorType]:
         """Evaluates the field at points along the ray.
 
         Args:

nerfstudio/fields/density_fields.py

@@ -17,7 +17,7 @@ Proposal network field.
 """
 
 
-from typing import Optional
+from typing import Optional, Tuple
 
 import numpy as np
 import torch
@@ -49,16 +49,16 @@ class HashMLPDensityField(Field):
 
     def __init__(
         self,
-        aabb,
+        aabb: TensorType,
         num_layers: int = 2,
         hidden_dim: int = 64,
         spatial_distortion: Optional[SpatialDistortion] = None,
-        use_linear=False,
-        num_levels=8,
-        max_res=1024,
-        base_res=16,
-        log2_hashmap_size=18,
-        features_per_level=2,
+        use_linear: bool = False,
+        num_levels: int = 8,
+        max_res: int = 1024,
+        base_res: int = 16,
+        log2_hashmap_size: int = 18,
+        features_per_level: int = 2,
     ) -> None:
         super().__init__()
         self.register_buffer("aabb", aabb)
@@ -99,7 +99,7 @@ class HashMLPDensityField(Field):
             self.encoding = tcnn.Encoding(n_input_dims=3, encoding_config=config["encoding"])
             self.linear = torch.nn.Linear(self.encoding.n_output_dims, 1)
 
-    def get_density(self, ray_samples: RaySamples):
+    def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType, None]:
         if self.spatial_distortion is not None:
             positions = self.spatial_distortion(ray_samples.frustums.get_positions())
             positions = (positions + 2.0) / 4.0
@@ -124,5 +124,5 @@ class HashMLPDensityField(Field):
         density = density * selector[..., None]
         return density, None
 
-    def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None):
+    def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None) -> dict:
         return {}

nerfstudio/fields/instant_ngp_field.py

@@ -17,7 +17,7 @@ Instant-NGP field implementations using tiny-cuda-nn, torch, ....
 """
 
 
-from typing import Optional
+from typing import Dict, Optional, Tuple
 
 import torch
 from nerfacc import ContractionType, contract
@@ -38,7 +38,7 @@ except ImportError:
     pass
 
 
-def get_normalized_directions(directions: TensorType["bs":..., 3]):
+def get_normalized_directions(directions: TensorType["bs":..., 3]) -> TensorType["bs":..., 3]:
     """SH encoding must be in the range [0, 1]
 
     Args:
@@ -67,13 +67,13 @@ class TCNNInstantNGPField(Field):
 
     def __init__(
         self,
-        aabb,
+        aabb: TensorType,
         num_layers: int = 2,
         hidden_dim: int = 64,
         geo_feat_dim: int = 15,
         num_layers_color: int = 3,
         hidden_dim_color: int = 64,
-        use_appearance_embedding: bool = False,
+        use_appearance_embedding: Optional[bool] = False,
         num_images: Optional[int] = None,
         appearance_embedding_dim: int = 32,
         contraction_type: ContractionType = ContractionType.UN_BOUNDED_SPHERE,
@@ -138,7 +138,7 @@ class TCNNInstantNGPField(Field):
             },
         )
 
-    def get_density(self, ray_samples: RaySamples):
+    def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType, TensorType]:
         positions = ray_samples.frustums.get_positions()
         positions_flat = positions.view(-1, 3)
         positions_flat = contract(x=positions_flat, roi=self.aabb, type=self.contraction_type)
@@ -152,7 +152,9 @@ class TCNNInstantNGPField(Field):
         density = trunc_exp(density_before_activation.to(positions))
         return density, base_mlp_out
 
-    def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None):
+    def get_outputs(
+        self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None
+    ) -> Dict[FieldHeadNames, TensorType]:
         directions = get_normalized_directions(ray_samples.frustums.directions)
         directions_flat = directions.view(-1, 3)
 

nerfstudio/fields/nerfacto_field.py

@@ -55,7 +55,7 @@ except ImportError:
     pass
 
 
-def get_normalized_directions(directions: TensorType["bs":..., 3]):
+def get_normalized_directions(directions: TensorType["bs":..., 3]) -> TensorType["bs":..., 3]:
     """SH encoding must be in the range [0, 1]
 
     Args:
@@ -92,7 +92,7 @@ class TCNNNerfactoField(Field):
 
     def __init__(
         self,
-        aabb,
+        aabb: TensorType,
         num_images: int,
         num_layers: int = 2,
         hidden_dim: int = 64,
@@ -112,7 +112,7 @@ class TCNNNerfactoField(Field):
         pass_semantic_gradients: bool = False,
         use_pred_normals: bool = False,
         use_average_appearance_embedding: bool = False,
-        spatial_distortion: Optional[SpatialDistortion] = None,
+        spatial_distortion: SpatialDistortion = None,
     ) -> None:
         super().__init__()
 
@@ -133,8 +133,8 @@ class TCNNNerfactoField(Field):
         self.use_pred_normals = use_pred_normals
         self.pass_semantic_gradients = pass_semantic_gradients
 
-        base_res = 16
-        features_per_level = 2
+        base_res: int = 16
+        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(
@@ -233,7 +233,7 @@ class TCNNNerfactoField(Field):
             },
         )
 
-    def get_density(self, ray_samples: RaySamples):
+    def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType, TensorType]:
         """Computes and returns the densities."""
         if self.spatial_distortion is not None:
             positions = ray_samples.frustums.get_positions()
@@ -259,7 +259,9 @@ class TCNNNerfactoField(Field):
         density = density * selector[..., None]
         return density, base_mlp_out
 
-    def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None):
+    def get_outputs(
+        self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None
+    ) -> Dict[FieldHeadNames, TensorType]:
         assert density_embedding is not None
         outputs = {}
         if ray_samples.camera_indices is None:
@@ -339,7 +341,7 @@ class TorchNerfactoField(Field):
 
     def __init__(
         self,
-        aabb,
+        aabb: TensorType,
         num_images: int,
         position_encoding: Encoding = HashEncoding(),
         direction_encoding: Encoding = SHEncoding(),
@@ -382,7 +384,7 @@ class TorchNerfactoField(Field):
         for field_head in self.field_heads:
             field_head.set_in_dim(self.mlp_head.get_out_dim())  # type: ignore
 
-    def get_density(self, ray_samples: RaySamples):
+    def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType, TensorType]:
         if self.spatial_distortion is not None:
             positions = ray_samples.frustums.get_positions()
             positions = self.spatial_distortion(positions)
@@ -396,7 +398,6 @@ class TorchNerfactoField(Field):
     def get_outputs(
         self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None
     ) -> Dict[FieldHeadNames, TensorType]:
-
         outputs_shape = ray_samples.frustums.directions.shape[:-1]
 
         if ray_samples.camera_indices is None:
@@ -427,4 +428,4 @@ class TorchNerfactoField(Field):
         return outputs
 
 
-field_implementation_to_class = {"tcnn": TCNNNerfactoField, "torch": TorchNerfactoField}
+field_implementation_to_class: Dict[str, Field] = {"tcnn": TCNNNerfactoField, "torch": TorchNerfactoField}

nerfstudio/fields/nerfplayer_nerfacto_field.py

@@ -17,7 +17,7 @@ Field implementations for NeRFPlayer (https://arxiv.org/abs/2210.15947) implemen
 """
 
 
-from typing import Optional
+from typing import Dict, Optional, Tuple
 
 import numpy as np
 import torch
@@ -47,7 +47,7 @@ except ImportError:
     pass
 
 
-def get_normalized_directions(directions: TensorType["bs":..., 3]):
+def get_normalized_directions(directions: TensorType["bs":..., 3]) -> TensorType["bs":..., 3]:
     """SH encoding must be in the range [0, 1]
 
     Args:
@@ -74,16 +74,16 @@ class TemporalHashMLPDensityField(Field):
 
     def __init__(
         self,
-        aabb,
+        aabb: TensorType,
         temporal_dim: int = 64,
         num_layers: int = 2,
         hidden_dim: int = 64,
         spatial_distortion: Optional[SpatialDistortion] = None,
-        num_levels=8,
-        max_res=1024,
-        base_res=16,
-        log2_hashmap_size=18,
-        features_per_level=2,
+        num_levels: int = 8,
+        max_res: int = 1024,
+        base_res: int = 16,
+        log2_hashmap_size: int = 18,
+        features_per_level: int = 2,
     ) -> None:
         super().__init__()
         # from .temporal_grid import test; test() # DEBUG
@@ -137,7 +137,7 @@ class TemporalHashMLPDensityField(Field):
         density, _ = self.get_density(ray_samples)
         return density
 
-    def get_density(self, ray_samples: RaySamples):
+    def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType, None]:
         if self.spatial_distortion is not None:
             positions = self.spatial_distortion(ray_samples.frustums.get_positions())
             positions = (positions + 2.0) / 4.0
@@ -154,7 +154,7 @@ class TemporalHashMLPDensityField(Field):
         density = trunc_exp(density_before_activation)
         return density, None
 
-    def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None):
+    def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None) -> dict:
         return {}
 
 
@@ -186,7 +186,7 @@ class NerfplayerNerfactoField(Field):
 
     def __init__(
         self,
-        aabb,
+        aabb: TensorType,
         num_images: int,
         num_layers: int = 2,
         hidden_dim: int = 64,
@@ -318,7 +318,7 @@ class NerfplayerNerfactoField(Field):
             },
         )
 
-    def get_density(self, ray_samples: RaySamples):
+    def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType, TensorType]:
         """Computes and returns the densities."""
         if self.spatial_distortion is not None:
             positions = ray_samples.frustums.get_positions()
@@ -339,7 +339,9 @@ class NerfplayerNerfactoField(Field):
         density = trunc_exp(density_before_activation.to(positions))
         return density, base_mlp_out
 
-    def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None):
+    def get_outputs(
+        self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None
+    ) -> Dict[FieldHeadNames, TensorType]:
         assert density_embedding is not None
         outputs = {}
         if ray_samples.camera_indices is None:

nerfstudio/fields/nerfplayer_ngp_field.py

@@ -17,7 +17,7 @@ NeRFPlayer (https://arxiv.org/abs/2210.15947) field implementations with Instant
 """
 
 
-from typing import Optional
+from typing import Dict, Optional, Tuple
 
 import torch
 from nerfacc import ContractionType, contract
@@ -68,7 +68,7 @@ class NerfplayerNGPField(Field):
 
     def __init__(
         self,
-        aabb,
+        aabb: TensorType,
         temporal_dim: int = 16,
         num_levels: int = 16,
         features_per_level: int = 2,
@@ -144,7 +144,7 @@ class NerfplayerNGPField(Field):
             },
         )
 
-    def get_density(self, ray_samples: RaySamples):
+    def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType, TensorType]:
         positions = ray_samples.frustums.get_positions()
         positions_flat = positions.view(-1, 3)
         positions_flat = contract(x=positions_flat, roi=self.aabb, type=self.contraction_type)
@@ -161,7 +161,9 @@ class NerfplayerNGPField(Field):
         density = trunc_exp(density_before_activation.to(positions))
         return density, base_mlp_out
 
-    def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None):
+    def get_outputs(
+        self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None
+    ) -> Dict[FieldHeadNames, TensorType]:
         directions = get_normalized_directions(ray_samples.frustums.directions)
         directions_flat = directions.view(-1, 3)
 

nerfstudio/fields/nerfw_field.py

@@ -109,7 +109,7 @@ class VanillaNerfWField(Field):
         self.field_head_transient_rgb = TransientRGBFieldHead(in_dim=self.mlp_transient.get_out_dim())
         self.field_head_transient_density = TransientDensityFieldHead(in_dim=self.mlp_transient.get_out_dim())
 
-    def get_density(self, ray_samples: RaySamples):
+    def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType, TensorType]:
         """Computes and returns the densities."""
         encoded_xyz = self.position_encoding(ray_samples.frustums.get_positions())
         encoded_xyz = self.position_encoding(ray_samples.frustums.get_positions())

nerfstudio/fields/semantic_nerf_field.py

@@ -88,7 +88,7 @@ class SemanticNerfField(Field):
             in_dim=self.mlp_semantic.get_out_dim(), num_classes=self.num_semantic_classes
         )
 
-    def get_density(self, ray_samples: RaySamples):
+    def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType, TensorType]:
         encoded_xyz = self.position_encoding(ray_samples.frustums.get_positions())
         base_mlp_out = self.mlp_base(encoded_xyz)
         density = self.field_head_density(base_mlp_out)

nerfstudio/fields/tensorf_field.py

@@ -15,7 +15,7 @@
 """TensoRF Field"""
 
 
-from typing import Optional
+from typing import Dict, Optional
 
 import torch
 from torch import nn
@@ -35,7 +35,7 @@ class TensoRFField(Field):
 
     def __init__(
         self,
-        aabb,
+        aabb: TensorType,
         # the aabb bounding box of the dataset
         feature_encoding: Encoding = Identity(in_dim=3),
         # the encoding method used for appearance encoding outputs
@@ -83,7 +83,7 @@ class TensoRFField(Field):
 
         self.field_output_rgb = RGBFieldHead(in_dim=self.mlp_head.get_out_dim(), activation=nn.Sigmoid())
 
-    def get_density(self, ray_samples: RaySamples):
+    def get_density(self, ray_samples: RaySamples) -> TensorType:
         positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb)
         positions = positions * 2 - 1
         density = self.density_encoding(positions)
@@ -118,7 +118,7 @@ class TensoRFField(Field):
         compute_normals: bool = False,
         mask: Optional[TensorType] = None,
         bg_color: Optional[TensorType] = None,
-    ):
+    ) -> Dict[FieldHeadNames, TensorType]:
         if compute_normals is True:
             raise ValueError("Surface normals are not currently supported with TensoRF")
         if mask is not None and bg_color is not None:

nerfstudio/fields/vanilla_nerf_field.py

@@ -58,7 +58,7 @@ class NeRFField(Field):
         head_mlp_num_layers: int = 2,
         head_mlp_layer_width: int = 128,
         skip_connections: Tuple[int] = (4,),
-        field_heads: Tuple[FieldHead] = (RGBFieldHead(),),
+        field_heads: Optional[Tuple[FieldHead]] = (RGBFieldHead(),),
         use_integrated_encoding: bool = False,
         spatial_distortion: Optional[SpatialDistortion] = None,
     ) -> None:
@@ -88,7 +88,7 @@ class NeRFField(Field):
         for field_head in self.field_heads:
             field_head.set_in_dim(self.mlp_head.get_out_dim())  # type: ignore
 
-    def get_density(self, ray_samples: RaySamples):
+    def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType, TensorType]:
         if self.use_integrated_encoding:
             gaussian_samples = ray_samples.frustums.get_gaussian_blob()
             if self.spatial_distortion is not None: