Fix some typing problems (#1940)

* Fix some typing problems

* nerfstudio.data typing error fixes

* Add tiny fixes

* Remove ABC from pipeline

* Prevent circular imports

* Tweak ModuleList typing

* TrainingCallbackAttributes cleanup

* Various minor nits

* More nits

* Run `ruff`

---------

Co-authored-by: Brent Yi <yibrenth@gmail.com>

nerfstudio/fields/base_field.py

@@ -45,12 +45,15 @@ class Field(nn.Module):
         self._sample_locations = None
         self._density_before_activation = None
 
-    def density_fn(self, positions: Shaped[Tensor, "*bs 3"]) -> Shaped[Tensor, "*bs 1"]:
+    def density_fn(
+        self, positions: Shaped[Tensor, "*bs 3"], times: Optional[Shaped[Tensor, "*bs 1"]] = None
+    ) -> Shaped[Tensor, "*bs 1"]:
         """Returns only the density. Used primarily with the density grid.
 
         Args:
             positions: the origin of the samples/frustums
         """
+        del times
         # Need to figure out a better way to describe positions with a ray.
         ray_samples = RaySamples(
             frustums=Frustums(

nerfstudio/fields/density_fields.py

@@ -47,6 +47,8 @@ class HashMLPDensityField(Field):
         use_linear: whether to skip the MLP and use a single linear layer instead
     """
 
+    aabb: Tensor
+
     def __init__(
         self,
         aabb: Tensor,

nerfstudio/fields/nerfacto_field.py

@@ -17,7 +17,7 @@ Field for compound nerf model, adds scene contraction and image embeddings to in
 """
 
 
-from typing import Dict, Optional, Tuple
+from typing import Dict, Optional, Tuple, Type
 
 import numpy as np
 import torch
@@ -80,6 +80,8 @@ class TCNNNerfactoField(Field):
         spatial_distortion: spatial distortion to apply to the scene
     """
 
+    aabb: Tensor
+
     def __init__(
         self,
         aabb: Tensor,
@@ -420,4 +422,4 @@ class TorchNerfactoField(Field):
         return outputs
 
 
-field_implementation_to_class: Dict[str, Field] = {"tcnn": TCNNNerfactoField, "torch": TorchNerfactoField}
+field_implementation_to_class: Dict[str, Type[Field]] = {"tcnn": TCNNNerfactoField, "torch": TorchNerfactoField}

nerfstudio/fields/nerfplayer_nerfacto_field.py

@@ -104,13 +104,16 @@ class TemporalHashMLPDensityField(Field):
             },
         )
 
-    def density_fn(self, positions: Float[Tensor, "*bs 3"], times: Float[Tensor, "bs 1"]) -> Float[Tensor, "*bs 1"]:
+    def density_fn(
+        self, positions: Float[Tensor, "*bs 3"], times: Optional[Float[Tensor, "bs 1"]] = None
+    ) -> Float[Tensor, "*bs 1"]:
         """Returns only the density. Used primarily with the density grid.
 
         Args:
             positions: the origin of the samples/frustums
             times: the time of rays
         """
+        assert times is not None, "TemporalHashMLPDensityField requires times to be specified"
         if len(positions.shape) == 3 and len(times.shape) == 2:
             # position is [ray, sample, 3]; times is [ray, 1]
             times = times[:, None]  # RaySamples can handle the shape

nerfstudio/fields/nerfplayer_ngp_field.py

@@ -115,7 +115,7 @@ class NerfplayerNGPField(Field):
             level_dim=features_per_level,
             base_resolution=base_resolution,
             log2_hashmap_size=log2_hashmap_size,
-            desired_resolution=1024 * (self.aabb.max() - self.aabb.min()),
+            desired_resolution=int(1024 * (self.aabb.max() - self.aabb.min())),
         )
         self.mlp_base_decode = tcnn.Network(
             n_input_dims=num_levels * features_per_level,
@@ -201,7 +201,9 @@ class NerfplayerNGPField(Field):
         rgb = self.mlp_head(h).view(*ray_samples.frustums.directions.shape[:-1], -1).to(directions)
         return {FieldHeadNames.RGB: rgb}
 
-    def density_fn(self, positions: Float[Tensor, "*bs 3"], times: Float[Tensor, "*bs 1"]) -> Float[Tensor, "*bs 1"]:
+    def density_fn(
+        self, positions: Float[Tensor, "*bs 3"], times: Optional[Float[Tensor, "*bs 1"]] = None
+    ) -> Float[Tensor, "*bs 1"]:
         """Returns only the density. Used primarily with the density grid.
         Overwrite this function since density is time dependent now.
 

nerfstudio/fields/sdf_field.py

@@ -48,6 +48,8 @@ class LearnedVariance(nn.Module):
         init_val: initial value in NeuS variance network
     """
 
+    variance: Tensor
+
     def __init__(self, init_val):
         super().__init__()
         self.register_parameter("variance", nn.Parameter(init_val * torch.ones(1), requires_grad=True))
@@ -205,6 +207,9 @@ class SDFField(Field):
 
         self._cos_anneal_ratio = 1.0
 
+        if self.use_grid_feature:
+            assert self.spatial_distortion is not None, "spatial distortion must be provided when using grid feature"
+
     def initialize_geo_layers(self) -> None:
         """
         Initialize layers for geometric network (sdf)
@@ -255,6 +260,7 @@ class SDFField(Field):
     def forward_geonetwork(self, inputs: Float[Tensor, "*batch 3"]) -> Float[Tensor, "*batch geo_features+1"]:
         """forward the geonetwork"""
         if self.use_grid_feature:
+            assert self.spatial_distortion is not None, "spatial distortion must be provided when using grid feature"
             positions = self.spatial_distortion(inputs)
             # map range [-2, 2] to [0, 1]
             positions = (positions + 2.0) / 4.0