diff --git a/nerfstudio/fields/base_field.py b/nerfstudio/fields/base_field.py index 5bf49bf5e..e20774d88 100644 --- a/nerfstudio/fields/base_field.py +++ b/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: diff --git a/nerfstudio/fields/density_fields.py b/nerfstudio/fields/density_fields.py index 443ca92c3..dcab3865d 100644 --- a/nerfstudio/fields/density_fields.py +++ b/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 {} diff --git a/nerfstudio/fields/instant_ngp_field.py b/nerfstudio/fields/instant_ngp_field.py index 846ea121e..102912ba7 100644 --- a/nerfstudio/fields/instant_ngp_field.py +++ b/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) diff --git a/nerfstudio/fields/nerfacto_field.py b/nerfstudio/fields/nerfacto_field.py index 7e4908e0b..830962edf 100644 --- a/nerfstudio/fields/nerfacto_field.py +++ b/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} diff --git a/nerfstudio/fields/nerfplayer_nerfacto_field.py b/nerfstudio/fields/nerfplayer_nerfacto_field.py index e87df3b03..0b186464f 100644 --- a/nerfstudio/fields/nerfplayer_nerfacto_field.py +++ b/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: diff --git a/nerfstudio/fields/nerfplayer_ngp_field.py b/nerfstudio/fields/nerfplayer_ngp_field.py index ed1cf250e..829ff0d33 100644 --- a/nerfstudio/fields/nerfplayer_ngp_field.py +++ b/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) diff --git a/nerfstudio/fields/nerfw_field.py b/nerfstudio/fields/nerfw_field.py index 483f5b502..377c900ea 100644 --- a/nerfstudio/fields/nerfw_field.py +++ b/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()) diff --git a/nerfstudio/fields/semantic_nerf_field.py b/nerfstudio/fields/semantic_nerf_field.py index 11fdb142a..08026b0a2 100644 --- a/nerfstudio/fields/semantic_nerf_field.py +++ b/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) diff --git a/nerfstudio/fields/tensorf_field.py b/nerfstudio/fields/tensorf_field.py index ac0c242f3..bbadc54f0 100644 --- a/nerfstudio/fields/tensorf_field.py +++ b/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: diff --git a/nerfstudio/fields/vanilla_nerf_field.py b/nerfstudio/fields/vanilla_nerf_field.py index 45511c1d1..03e496d48 100644 --- a/nerfstudio/fields/vanilla_nerf_field.py +++ b/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: