From f481f5ebec847a9d6781f090602cc9bebf47cf23 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jon=C3=A1=C5=A1=20Kulh=C3=A1nek?= Date: Thu, 18 May 2023 02:36:56 +0200 Subject: [PATCH] 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 --- nerfstudio/cameras/rays.py | 6 ++--- nerfstudio/engine/callbacks.py | 18 ++++++++----- nerfstudio/engine/schedulers.py | 9 +------ nerfstudio/engine/trainer.py | 6 ++--- nerfstudio/exporter/marching_cubes.py | 2 +- nerfstudio/field_components/activations.py | 18 ++++++++++--- nerfstudio/field_components/encodings.py | 3 +++ nerfstudio/field_components/temporal_grid.py | 8 +++++- nerfstudio/fields/base_field.py | 5 +++- nerfstudio/fields/density_fields.py | 2 ++ nerfstudio/fields/nerfacto_field.py | 6 +++-- .../fields/nerfplayer_nerfacto_field.py | 5 +++- nerfstudio/fields/nerfplayer_ngp_field.py | 6 +++-- nerfstudio/fields/sdf_field.py | 6 +++++ nerfstudio/generative/stable_diffusion.py | 4 +-- nerfstudio/model_components/losses.py | 24 ++++++++++------- nerfstudio/model_components/ray_generators.py | 2 ++ nerfstudio/model_components/ray_samplers.py | 26 ++++++++++++++----- nerfstudio/models/base_surface_model.py | 16 +++++++++--- nerfstudio/models/depth_nerfacto.py | 8 +++--- nerfstudio/models/mipnerf.py | 2 ++ nerfstudio/models/nerfplayer_nerfacto.py | 12 +++++---- nerfstudio/models/tensorf.py | 15 ++++++++--- nerfstudio/pipelines/base_pipeline.py | 4 +-- nerfstudio/utils/misc.py | 11 +++++--- nerfstudio/utils/plotly_utils.py | 2 +- nerfstudio/utils/tensor_dataclass.py | 3 +-- 27 files changed, 154 insertions(+), 75 deletions(-) diff --git a/nerfstudio/cameras/rays.py b/nerfstudio/cameras/rays.py index 30a3057ba..33476f92e 100644 --- a/nerfstudio/cameras/rays.py +++ b/nerfstudio/cameras/rays.py @@ -152,14 +152,14 @@ class RaySamples(TensorDataclass): @overload @staticmethod def get_weights_and_transmittance_from_alphas( - alphas: Float[Tensor, "*batch num_samples 1"], weights_only: Literal[False] = False - ) -> Tuple[Float[Tensor, "*batch num_samples 1"], Float[Tensor, "*batch num_samples 1"]]: + alphas: Float[Tensor, "*batch num_samples 1"], weights_only: Literal[True] + ) -> Float[Tensor, "*batch num_samples 1"]: ... @overload @staticmethod def get_weights_and_transmittance_from_alphas( - alphas: Float[Tensor, "*batch num_samples 1"], weights_only: Literal[True] + alphas: Float[Tensor, "*batch num_samples 1"], weights_only: Literal[False] = False ) -> Tuple[Float[Tensor, "*batch num_samples 1"], Float[Tensor, "*batch num_samples 1"]]: ... diff --git a/nerfstudio/engine/callbacks.py b/nerfstudio/engine/callbacks.py index 3a1db68fb..591711ab5 100644 --- a/nerfstudio/engine/callbacks.py +++ b/nerfstudio/engine/callbacks.py @@ -17,10 +17,17 @@ Callback code used for training iterations """ from __future__ import annotations -from dataclasses import InitVar, dataclass +from dataclasses import dataclass from enum import Enum, auto from inspect import signature -from typing import Callable, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple + +from torch.cuda.amp.grad_scaler import GradScaler + +from nerfstudio.engine.optimizers import Optimizers + +if TYPE_CHECKING: + from nerfstudio.pipelines.base_pipeline import Pipeline @dataclass @@ -30,12 +37,11 @@ class TrainingCallbackAttributes: Instead of providing access to the entire Trainer object, we only provide these attributes. This should be least prone to errors and fairly clean from a user perspective.""" - # TODO(ethan): type this without circular imports - optimizers: Optional[InitVar] + optimizers: Optional[Optimizers] """optimizers for training""" - grad_scaler: Optional[InitVar] + grad_scaler: Optional[GradScaler] """gradient scalers""" - pipeline: Optional[InitVar] + pipeline: Optional["Pipeline"] # Prevent circular import. """reference to training pipeline""" diff --git a/nerfstudio/engine/schedulers.py b/nerfstudio/engine/schedulers.py index a922bc5df..a60e5c2f0 100644 --- a/nerfstudio/engine/schedulers.py +++ b/nerfstudio/engine/schedulers.py @@ -21,17 +21,10 @@ from typing import Literal, Optional, Tuple, Type import numpy as np from torch.optim import Optimizer, lr_scheduler +from torch.optim.lr_scheduler import LRScheduler from nerfstudio.configs.base_config import InstantiateConfig -try: - from torch.optim.lr_scheduler import ( # pylint: disable=ungrouped-imports - LRScheduler, - ) -except ImportError: - # Backward compatibility for PyTorch 1.x - from torch.optim.lr_scheduler import _LRScheduler as LRScheduler - @dataclass class SchedulerConfig(InstantiateConfig): diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index dc1dfb586..aa704d044 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -173,9 +173,9 @@ class Trainer: self.callbacks = self.pipeline.get_training_callbacks( TrainingCallbackAttributes( - optimizers=self.optimizers, # type: ignore - grad_scaler=self.grad_scaler, # type: ignore - pipeline=self.pipeline, # type: ignore + optimizers=self.optimizers, + grad_scaler=self.grad_scaler, + pipeline=self.pipeline, ) ) diff --git a/nerfstudio/exporter/marching_cubes.py b/nerfstudio/exporter/marching_cubes.py index 1a9292f18..a83307501 100644 --- a/nerfstudio/exporter/marching_cubes.py +++ b/nerfstudio/exporter/marching_cubes.py @@ -112,7 +112,7 @@ def evaluate_multiresolution_sdf( if pts_to_eval.shape[0] > 0: pts_sdf_eval = evaluate(pts_to_eval.contiguous()) assert pts_sdf is not None - pts_sdf[mask] = pts_sdf_eval # pylint: disable=unsupported-assignment-operation + pts_sdf[mask] = pts_sdf_eval if pid < 3: # Update mask diff --git a/nerfstudio/field_components/activations.py b/nerfstudio/field_components/activations.py index bc4cc58d0..e413965e2 100644 --- a/nerfstudio/field_components/activations.py +++ b/nerfstudio/field_components/activations.py @@ -16,7 +16,11 @@ Special activation functions. """ +from typing import TYPE_CHECKING + import torch +from jaxtyping import Float +from torch import Tensor from torch.autograd import Function from torch.cuda.amp import custom_bwd, custom_fwd @@ -37,6 +41,14 @@ class _TruncExp(Function): return g * torch.exp(x.clamp(-15, 15)) -trunc_exp = _TruncExp.apply -"""Same as torch.exp, but with the backward pass clipped to prevent vanishing/exploding -gradients.""" +if TYPE_CHECKING: + + def trunc_exp(_: Float[Tensor, "*bs"], /) -> Float[Tensor, "*bs"]: + """Same as torch.exp, but with the backward pass clipped to prevent vanishing/exploding + gradients.""" + raise NotImplementedError() + +else: + trunc_exp = _TruncExp.apply + """Same as torch.exp, but with the backward pass clipped to prevent vanishing/exploding + gradients.""" diff --git a/nerfstudio/field_components/encodings.py b/nerfstudio/field_components/encodings.py index 22408998b..94cd1f189 100644 --- a/nerfstudio/field_components/encodings.py +++ b/nerfstudio/field_components/encodings.py @@ -621,6 +621,7 @@ class KPlanesEncoding(Encoding): """Sample features from this encoder. Expects ``in_tensor`` to be in range [-1, 1]""" original_shape = in_tensor.shape + assert any(self.coo_combs) output = 1.0 if self.reduce == "product" else 0.0 # identity for corresponding op for ci, coo_comb in enumerate(self.coo_combs): grid = self.plane_coefs[ci].unsqueeze(0) # [1, feature_dim, reso1, reso2] @@ -634,6 +635,8 @@ class KPlanesEncoding(Encoding): else: output = output + interp + # Typing: output gets converted to a tensor after the first iteration of the loop + assert isinstance(output, Tensor) return output.reshape(*original_shape[:-1], self.num_components) diff --git a/nerfstudio/field_components/temporal_grid.py b/nerfstudio/field_components/temporal_grid.py index 71880b375..5c80c57a3 100644 --- a/nerfstudio/field_components/temporal_grid.py +++ b/nerfstudio/field_components/temporal_grid.py @@ -174,6 +174,11 @@ class TemporalGridEncoder(nn.Module): align_corners: same as other interpolation operators """ + sampling_index: Tensor + index_a_mask: Tensor + index_b_mask: Tensor + index_list: Tensor + def __init__( self, temporal_dim: int = 64, @@ -345,6 +350,7 @@ class TemporalGridEncoder(nn.Module): self.gridtype_id, self.align_corners, ) + assert isinstance(outputs, Tensor) return outputs def get_temporal_tv_loss(self) -> Float[Tensor, ""]: @@ -352,6 +358,6 @@ class TemporalGridEncoder(nn.Module): Sample a random channel combination (i.e., row for the combination table), and then compute loss on it. """ - row_idx = torch.randint(0, len(self.index_list), [1]).item() + row_idx = torch.randint(0, len(self.index_list), [1]) feat_idx = self.index_list[row_idx] return (self.embeddings[:, feat_idx[0]] - self.embeddings[:, feat_idx[1]]).abs().mean() diff --git a/nerfstudio/fields/base_field.py b/nerfstudio/fields/base_field.py index 2392ebb12..068d4e7f1 100644 --- a/nerfstudio/fields/base_field.py +++ b/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( diff --git a/nerfstudio/fields/density_fields.py b/nerfstudio/fields/density_fields.py index 40a965a67..47eec223a 100644 --- a/nerfstudio/fields/density_fields.py +++ b/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, diff --git a/nerfstudio/fields/nerfacto_field.py b/nerfstudio/fields/nerfacto_field.py index c458d6182..b1b1c2cfb 100644 --- a/nerfstudio/fields/nerfacto_field.py +++ b/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} diff --git a/nerfstudio/fields/nerfplayer_nerfacto_field.py b/nerfstudio/fields/nerfplayer_nerfacto_field.py index 0d8ecde48..0dc4ba4ce 100644 --- a/nerfstudio/fields/nerfplayer_nerfacto_field.py +++ b/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 diff --git a/nerfstudio/fields/nerfplayer_ngp_field.py b/nerfstudio/fields/nerfplayer_ngp_field.py index 229d5f6dc..803637729 100644 --- a/nerfstudio/fields/nerfplayer_ngp_field.py +++ b/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. diff --git a/nerfstudio/fields/sdf_field.py b/nerfstudio/fields/sdf_field.py index 358b2d7e5..031614f2f 100644 --- a/nerfstudio/fields/sdf_field.py +++ b/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 diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index 26a804b1e..5b222fa19 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/nerfstudio/generative/stable_diffusion.py @@ -19,7 +19,7 @@ import sys from dataclasses import dataclass from pathlib import Path -from typing import List, Optional, Union +from typing import List, Optional, Union, cast import appdirs import mediapy @@ -223,7 +223,7 @@ class StableDiffusion(nn.Module): if grad_scaler is not None: latents = grad_scaler.scale(latents) - loss = _SDSGradient.apply(latents, grad) + loss = cast(Tensor, _SDSGradient.apply(latents, grad)) return loss diff --git a/nerfstudio/model_components/losses.py b/nerfstudio/model_components/losses.py index ed766c10d..cd589cad3 100644 --- a/nerfstudio/model_components/losses.py +++ b/nerfstudio/model_components/losses.py @@ -16,13 +16,14 @@ Collection of Losses. """ from enum import Enum -from typing import Dict, Literal, Optional, TypeVar, cast +from typing import Dict, Literal, Optional, Tuple, cast import torch from jaxtyping import Bool, Float from torch import Tensor, nn from nerfstudio.cameras.rays import RaySamples +from nerfstudio.field_components.field_heads import FieldHeadNames from nerfstudio.utils.math import masked_reduction, normalized_depth_scale_and_shift L1Loss = nn.L1Loss @@ -103,7 +104,7 @@ def ray_samples_to_sdist(ray_samples): return sdist -def interlevel_loss(weights_list, ray_samples_list): +def interlevel_loss(weights_list, ray_samples_list) -> torch.Tensor: """Calculates the proposal loss in the MipNeRF-360 paper. https://github.com/kakaobrain/NeRF-Factory/blob/f61bb8744a5cb4820a4d968fb3bfbed777550f4a/src/model/mipnerf360/model.py#L515 @@ -111,12 +112,16 @@ def interlevel_loss(weights_list, ray_samples_list): """ c = ray_samples_to_sdist(ray_samples_list[-1]).detach() w = weights_list[-1][..., 0].detach() + assert len(ray_samples_list) > 0 + loss_interlevel = 0.0 for ray_samples, weights in zip(ray_samples_list[:-1], weights_list[:-1]): sdist = ray_samples_to_sdist(ray_samples) cp = sdist # (num_rays, num_samples + 1) wp = weights[..., 0] # (num_rays, num_samples) loss_interlevel += torch.mean(lossfun_outer(c, w, cp, wp)) + + assert isinstance(loss_interlevel, Tensor) return loss_interlevel @@ -398,7 +403,8 @@ class GradientLoss(nn.Module): Returns: gradient loss based on reduction function """ - total = 0 + assert self.__scales >= 1 + total = 0.0 for scale in range(self.__scales): step = pow(2, scale) @@ -410,7 +416,8 @@ class GradientLoss(nn.Module): ) total += grad_loss - return cast(Tensor, total) + assert isinstance(total, Tensor) + return total def gradient_loss( self, @@ -534,12 +541,9 @@ class _GradientScaler(torch.autograd.Function): # typing: ignore return output_grad * scaling, grad_scaling -K = TypeVar("K") - - def scale_gradients_by_distance_squared( - field_outputs: Dict[K, torch.Tensor], ray_samples: RaySamples -) -> Dict[K, torch.Tensor]: + field_outputs: Dict[FieldHeadNames, torch.Tensor], ray_samples: RaySamples +) -> Dict[FieldHeadNames, torch.Tensor]: """ Scale gradients by the ray distance to the pixel as suggested in `Radiance Field Gradient Scaling for Unbiased Near-Camera Training` paper @@ -554,5 +558,5 @@ def scale_gradients_by_distance_squared( ray_dist = (ray_samples.frustums.starts + ray_samples.frustums.ends) / 2 scaling = torch.square(ray_dist).clamp(0, 1) for key, value in field_outputs.items(): - out[key], _ = _GradientScaler.apply(value, scaling) + out[key], _ = cast(Tuple[Tensor, Tensor], _GradientScaler.apply(value, scaling)) return out diff --git a/nerfstudio/model_components/ray_generators.py b/nerfstudio/model_components/ray_generators.py index 89e64cf04..49754548d 100644 --- a/nerfstudio/model_components/ray_generators.py +++ b/nerfstudio/model_components/ray_generators.py @@ -32,6 +32,8 @@ class RayGenerator(nn.Module): pose_optimizer: pose optimization module, for optimizing noisy camera intrinsics/extrinsics. """ + image_coords: Tensor + def __init__(self, cameras: Cameras, pose_optimizer: CameraOptimizer) -> None: super().__init__() self.cameras = cameras diff --git a/nerfstudio/model_components/ray_samplers.py b/nerfstudio/model_components/ray_samplers.py index fa97905a3..59d6d7b82 100644 --- a/nerfstudio/model_components/ray_samplers.py +++ b/nerfstudio/model_components/ray_samplers.py @@ -17,7 +17,7 @@ Collection of sampling strategies """ from abc import abstractmethod -from typing import Callable, List, Optional, Tuple, Union +from typing import Callable, List, Optional, Protocol, Tuple, Union import torch from jaxtyping import Float @@ -277,7 +277,7 @@ class PDFSampler(Sampler): self, ray_bundle: Optional[RayBundle] = None, ray_samples: Optional[RaySamples] = None, - weights: Float[Tensor, "*batch num_samples 1"] = None, + weights: Optional[Float[Tensor, "*batch num_samples 1"]] = None, num_samples: Optional[int] = None, eps: float = 1e-5, ) -> RaySamples: @@ -296,6 +296,7 @@ class PDFSampler(Sampler): if ray_samples is None or ray_bundle is None: raise ValueError("ray_samples and ray_bundle must be provided") + assert weights is not None, "weights must be provided" num_samples = num_samples or self.num_samples assert num_samples is not None @@ -371,6 +372,17 @@ class PDFSampler(Sampler): return ray_samples +class DensityFn(Protocol): + """ + Function that evaluates density at a given point. + """ + + def __call__( + self, positions: Float[Tensor, "*batch 3"], times: Optional[Float[Tensor, "*batch 1"]] = None + ) -> Float[Tensor, "*batch 1"]: + ... + + class VolumetricSampler(Sampler): """Sampler inspired by the one proposed in the Instant-NGP paper. Generates samples along a ray by sampling the occupancy field. @@ -385,8 +397,8 @@ class VolumetricSampler(Sampler): def __init__( self, occupancy_grid: OccGridEstimator, - density_fn: Optional[Callable[[Float[Tensor, "*batch 3"]], Float[Tensor, "*batch 1"]]] = None, - ) -> None: + density_fn: Optional[DensityFn] = None, + ): super().__init__() assert occupancy_grid is not None self.density_fn = density_fn @@ -523,7 +535,7 @@ class ProposalNetworkSampler(Sampler): def __init__( self, - num_proposal_samples_per_ray: Tuple[int] = (64,), + num_proposal_samples_per_ray: Tuple[int, ...] = (64,), num_nerf_samples_per_ray: int = 32, num_proposal_network_iterations: int = 2, single_jitter: bool = False, @@ -634,7 +646,7 @@ class NeuSSampler(Sampler): def generate_ray_samples( self, ray_bundle: Optional[RayBundle] = None, - sdf_fn: Optional[Callable] = None, + sdf_fn: Optional[Callable[[RaySamples], torch.Tensor]] = None, ray_samples: Optional[RaySamples] = None, ) -> Union[Tuple[RaySamples, torch.Tensor], RaySamples]: assert ray_bundle is not None @@ -647,6 +659,7 @@ class NeuSSampler(Sampler): total_iters = 0 sorted_index = None + sdf: Optional[torch.Tensor] = None new_samples = ray_samples base_variance = self.base_variance @@ -657,6 +670,7 @@ class NeuSSampler(Sampler): # merge sdf predictions if sorted_index is not None: + assert sdf is not None sdf_merge = torch.cat([sdf.squeeze(-1), new_sdf.squeeze(-1)], -1) sdf = torch.gather(sdf_merge, 1, sorted_index).unsqueeze(-1) else: diff --git a/nerfstudio/models/base_surface_model.py b/nerfstudio/models/base_surface_model.py index 6076da403..cd0e745d4 100644 --- a/nerfstudio/models/base_surface_model.py +++ b/nerfstudio/models/base_surface_model.py @@ -20,7 +20,7 @@ from __future__ import annotations from abc import abstractmethod from dataclasses import dataclass, field -from typing import Any, Dict, List, Literal, Tuple, Type +from typing import Any, Dict, List, Literal, Tuple, Type, cast import torch import torch.nn.functional as F @@ -202,10 +202,16 @@ class SurfaceModel(Model): Returns: Outputs of model. (ie. rendered colors) """ + assert ( + ray_bundle.metadata is not None and "directions_norm" in ray_bundle.metadata + ), "directions_norm is required in ray_bundle.metadata" + samples_and_field_outputs = self.sample_and_forward_field(ray_bundle=ray_bundle) # shortcuts - field_outputs = samples_and_field_outputs["field_outputs"] + field_outputs: Dict[FieldHeadNames, torch.Tensor] = cast( + Dict[FieldHeadNames, torch.Tensor], samples_and_field_outputs["field_outputs"] + ) ray_samples = samples_and_field_outputs["ray_samples"] weights = samples_and_field_outputs["weights"] bg_transmittance = samples_and_field_outputs["bg_transmittance"] @@ -220,6 +226,8 @@ class SurfaceModel(Model): # background model if self.config.background_model != "none": + assert isinstance(self.field_background, torch.nn.Module), "field_background should be a module" + assert ray_bundle.fars is not None, "fars is required in ray_bundle" # sample inversely from far to 1000 and points and forward the bg model ray_bundle.nears = ray_bundle.fars assert ray_bundle.fars is not None @@ -264,8 +272,8 @@ class SurfaceModel(Model): outputs.update(samples_and_field_outputs) if "weights_list" in samples_and_field_outputs: - weights_list = samples_and_field_outputs["weights_list"] - ray_samples_list = samples_and_field_outputs["ray_samples_list"] + weights_list = cast(List[torch.Tensor], samples_and_field_outputs["weights_list"]) + ray_samples_list = cast(List[torch.Tensor], samples_and_field_outputs["ray_samples_list"]) for i in range(len(weights_list) - 1): outputs[f"prop_depth_{i}"] = self.renderer_depth( diff --git a/nerfstudio/models/depth_nerfacto.py b/nerfstudio/models/depth_nerfacto.py index 0bbde5cde..e2c7b6f05 100644 --- a/nerfstudio/models/depth_nerfacto.py +++ b/nerfstudio/models/depth_nerfacto.py @@ -115,13 +115,13 @@ class DepthNerfactoModel(NerfactoModel): predicted_depth_colormap = colormaps.apply_depth_colormap( outputs["depth"], accumulation=outputs["accumulation"], - near_plane=torch.min(ground_truth_depth), - far_plane=torch.max(ground_truth_depth), + near_plane=float(torch.min(ground_truth_depth).cpu()), + far_plane=float(torch.max(ground_truth_depth).cpu()), ) images["depth"] = torch.cat([ground_truth_depth_colormap, predicted_depth_colormap], dim=1) depth_mask = ground_truth_depth > 0 - metrics["depth_mse"] = torch.nn.functional.mse_loss( - outputs["depth"][depth_mask], ground_truth_depth[depth_mask] + metrics["depth_mse"] = float( + torch.nn.functional.mse_loss(outputs["depth"][depth_mask], ground_truth_depth[depth_mask]).cpu() ) return metrics, images diff --git a/nerfstudio/models/mipnerf.py b/nerfstudio/models/mipnerf.py index 776764ac1..f0deb9bc4 100644 --- a/nerfstudio/models/mipnerf.py +++ b/nerfstudio/models/mipnerf.py @@ -58,6 +58,7 @@ class MipNerfModel(Model): self.field = None assert config.collider_params is not None, "MipNeRF model requires bounding box collider parameters." super().__init__(config=config, **kwargs) + assert self.config.collider_params is not None, "mip-NeRF requires collider parameters to be set." def populate_modules(self): """Set the fields and modules""" @@ -150,6 +151,7 @@ class MipNerfModel(Model): def get_image_metrics_and_images( self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor] ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]: + assert self.config.collider_params is not None, "mip-NeRF requires collider parameters to be set." image = batch["image"].to(outputs["rgb_coarse"].device) rgb_coarse = outputs["rgb_coarse"] rgb_fine = outputs["rgb_fine"] diff --git a/nerfstudio/models/nerfplayer_nerfacto.py b/nerfstudio/models/nerfplayer_nerfacto.py index 38150905b..46f0fcd3f 100644 --- a/nerfstudio/models/nerfplayer_nerfacto.py +++ b/nerfstudio/models/nerfplayer_nerfacto.py @@ -20,7 +20,7 @@ from __future__ import annotations import functools from dataclasses import dataclass, field -from typing import Dict, List, Literal, Type +from typing import Dict, List, Literal, Sequence, Type, cast import numpy as np import torch @@ -118,15 +118,16 @@ class NerfplayerNerfactoModel(NerfactoModel): self.density_fns = [] num_prop_nets = self.config.num_proposal_iterations + # Build the proposal network(s) - self.proposal_networks = torch.nn.ModuleList() + proposal_networks: List[TemporalHashMLPDensityField] = [] if self.config.use_same_proposal_network: assert len(self.config.proposal_net_args_list) == 1, "Only one proposal network is allowed." prop_net_args = self.config.proposal_net_args_list[0] network = TemporalHashMLPDensityField( self.scene_box.aabb, spatial_distortion=scene_contraction, **prop_net_args ) - self.proposal_networks.append(network) + proposal_networks.append(network) self.density_fns.extend([network.density_fn for _ in range(num_prop_nets)]) else: for i in range(num_prop_nets): @@ -136,8 +137,9 @@ class NerfplayerNerfactoModel(NerfactoModel): spatial_distortion=scene_contraction, **prop_net_args, ) - self.proposal_networks.append(network) - self.density_fns.extend([network.density_fn for network in self.proposal_networks]) + proposal_networks.append(network) + self.density_fns.extend([network.density_fn for network in proposal_networks]) + self.proposal_networks = cast(Sequence[TemporalHashMLPDensityField], torch.nn.ModuleList(proposal_networks)) # Samplers def update_schedule(step): diff --git a/nerfstudio/models/tensorf.py b/nerfstudio/models/tensorf.py index 5cbb906b6..742818b64 100644 --- a/nerfstudio/models/tensorf.py +++ b/nerfstudio/models/tensorf.py @@ -19,7 +19,7 @@ TensorRF implementation. from __future__ import annotations from dataclasses import dataclass, field -from typing import Dict, List, Literal, Tuple, Type +from typing import Dict, List, Literal, Tuple, Type, cast import numpy as np import torch @@ -130,6 +130,8 @@ class TensoRFModel(Model): ) -> List[TrainingCallback]: # the callback that we want to run every X iterations after the training iteration def reinitialize_optimizer(self, training_callback_attributes: TrainingCallbackAttributes, step: int): + assert training_callback_attributes.optimizers is not None + assert training_callback_attributes.pipeline is not None index = self.upsampling_iters.index(step) resolution = self.upsampling_steps[index] @@ -313,8 +315,13 @@ class TensoRFModel(Model): l1_parameters.append(parameter.view(-1)) loss_dict["l1_reg"] = torch.abs(torch.cat(l1_parameters)).mean() elif self.config.regularization == "tv": - loss_dict["tv_reg_density"] = tv_loss(self.field.density_encoding.plane_coef) - loss_dict["tv_reg_color"] = tv_loss(self.field.color_encoding.plane_coef) + density_plane_coef = self.field.density_encoding.plane_coef + color_plane_coef = self.field.color_encoding.plane_coef + assert isinstance(color_plane_coef, torch.Tensor) and isinstance( + density_plane_coef, torch.Tensor + ), "TV reg only supported for TensoRF encoding types with plane_coef attribute" + loss_dict["tv_reg_density"] = tv_loss(density_plane_coef) + loss_dict["tv_reg_color"] = tv_loss(color_plane_coef) elif self.config.regularization == "none": pass else: @@ -344,7 +351,7 @@ class TensoRFModel(Model): rgb = torch.moveaxis(rgb, -1, 0)[None, ...] psnr = self.psnr(image, rgb) - ssim = self.ssim(image, rgb) + ssim = cast(torch.Tensor, self.ssim(image, rgb)) lpips = self.lpips(image, rgb) metrics_dict = { diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 9e80425cb..a17dfbb0c 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -21,7 +21,7 @@ import typing from abc import abstractmethod from dataclasses import dataclass, field from time import time -from typing import Any, Dict, List, Literal, Mapping, Optional, Type, Union, cast +from typing import Any, Dict, List, Literal, Mapping, Optional, Tuple, Type, Union, cast import torch import torch.distributed as dist @@ -300,7 +300,7 @@ class VanillaPipeline(Pipeline): raise NotImplementedError @profiler.time_function - def get_eval_loss_dict(self, step: int): + def get_eval_loss_dict(self, step: int) -> Tuple[Any, Dict[str, Any], Dict[str, Any]]: """This function gets your evaluation loss dict. It needs to get the data from the DataManager and feed it to the model's forward function diff --git a/nerfstudio/utils/misc.py b/nerfstudio/utils/misc.py index cfa5e73db..f05cd4d0c 100644 --- a/nerfstudio/utils/misc.py +++ b/nerfstudio/utils/misc.py @@ -17,12 +17,15 @@ Miscellaneous helper code. """ -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable, Dict, List, Optional, TypeVar, Union import torch +T = TypeVar("T") +TKey = TypeVar("TKey") -def get_dict_to_torch(stuff: Any, device: Union[torch.device, str] = "cpu", exclude: Optional[List[str]] = None): + +def get_dict_to_torch(stuff: T, device: Union[torch.device, str] = "cpu", exclude: Optional[List[str]] = None) -> T: """Set everything in the dict to the specified torch device. Args: @@ -42,7 +45,7 @@ def get_dict_to_torch(stuff: Any, device: Union[torch.device, str] = "cpu", excl return stuff -def get_dict_to_cpu(stuff: Any): +def get_dict_to_cpu(stuff: T) -> T: """Set everything in the dict to CPU. Args: @@ -57,7 +60,7 @@ def get_dict_to_cpu(stuff: Any): return stuff -def get_masked_dict(d, mask): +def get_masked_dict(d: Dict[TKey, torch.Tensor], mask) -> Dict[TKey, torch.Tensor]: """Return a masked dictionary. TODO(ethan): add more asserts/checks so this doesn't have unpredictable behavior. diff --git a/nerfstudio/utils/plotly_utils.py b/nerfstudio/utils/plotly_utils.py index 8eff73a25..9dbc7c851 100644 --- a/nerfstudio/utils/plotly_utils.py +++ b/nerfstudio/utils/plotly_utils.py @@ -148,7 +148,7 @@ def vis_dataset( ), ) fig = go.Figure(data=data, layout=layout) # type: ignore - return fig + return fig # type: ignore def get_random_color(colormap: Optional[List[str]] = None, idx: Optional[int] = None) -> str: diff --git a/nerfstudio/utils/tensor_dataclass.py b/nerfstudio/utils/tensor_dataclass.py index 734972428..0ebe7b5ce 100644 --- a/nerfstudio/utils/tensor_dataclass.py +++ b/nerfstudio/utils/tensor_dataclass.py @@ -127,9 +127,8 @@ class TensorDataclass: new_dict = {} for k, v in dict_.items(): if isinstance(v, torch.Tensor): - # If custom dimension key, then we need to + # Apply field-specific custom dimensions. if isinstance(self._field_custom_dimensions, dict) and k in self._field_custom_dimensions: - # pylint: disable=unsubscriptable-object new_dict[k] = v.broadcast_to( ( *batch_shape,