From 6ba2a18b2bc435ffe35de3d48433a2353c265210 Mon Sep 17 00:00:00 2001 From: Brent Yi Date: Tue, 16 May 2023 15:57:47 -0700 Subject: [PATCH] `torchtyping` => `jaxtyping` (#1838) * Run `ruff --fix` * Add jaxtyping migration script * Run migration scripts + ruff + black * Merge remote-tracking branch 'origin/main' into brent/jaxtyping * black * Fix lint * fix lint * Cleanup and more specific types * Fixed annotation errors * Fix more errors * Fix annotation * Delete _torchtyping_to_jaxtyping.py --------- Co-authored-by: Matt Tancik --- docs/conf.py | 1 - nerfstudio/cameras/camera_optimizers.py | 8 +- nerfstudio/cameras/camera_utils.py | 31 ++--- nerfstudio/cameras/cameras.py | 87 +++++++------- nerfstudio/cameras/lie_groups.py | 7 +- nerfstudio/cameras/rays.py | 61 +++++----- .../data/dataparsers/base_dataparser.py | 21 ++-- nerfstudio/data/datasets/base_dataset.py | 5 +- nerfstudio/data/pixel_samplers.py | 15 +-- nerfstudio/data/scene_box.py | 9 +- nerfstudio/data/utils/colmap_parsing_utils.py | 1 - nerfstudio/exporter/exporter_utils.py | 11 +- nerfstudio/exporter/marching_cubes.py | 17 ++- nerfstudio/exporter/texture_utils.py | 33 +++--- nerfstudio/exporter/tsdf_utils.py | 31 ++--- nerfstudio/field_components/__init__.py | 9 +- .../field_components/base_field_component.py | 6 +- nerfstudio/field_components/embedding.py | 5 +- nerfstudio/field_components/encodings.py | 44 +++---- nerfstudio/field_components/field_heads.py | 8 +- nerfstudio/field_components/mlp.py | 6 +- .../field_components/spatial_distortions.py | 18 +-- .../field_components/temporal_distortions.py | 6 +- nerfstudio/field_components/temporal_grid.py | 22 ++-- nerfstudio/fields/base_field.py | 20 ++-- nerfstudio/fields/density_fields.py | 8 +- nerfstudio/fields/instant_ngp_field.py | 13 +- nerfstudio/fields/nerfacto_field.py | 19 ++- .../fields/nerfplayer_nerfacto_field.py | 19 +-- nerfstudio/fields/nerfplayer_ngp_field.py | 15 +-- nerfstudio/fields/nerfw_field.py | 14 +-- nerfstudio/fields/sdf_field.py | 39 +++--- nerfstudio/fields/semantic_nerf_field.py | 9 +- nerfstudio/fields/tensorf_field.py | 15 ++- nerfstudio/fields/vanilla_nerf_field.py | 9 +- nerfstudio/generative/stable_diffusion.py | 20 ++-- nerfstudio/model_components/losses.py | 112 ++++++++++-------- nerfstudio/model_components/ray_generators.py | 6 +- nerfstudio/model_components/ray_samplers.py | 19 +-- nerfstudio/model_components/renderers.py | 68 +++++------ .../model_components/scene_colliders.py | 6 +- nerfstudio/model_components/shaders.py | 14 +-- nerfstudio/models/nerfacto.py | 16 ++- nerfstudio/models/nerfplayer_nerfacto.py | 17 ++- nerfstudio/models/neus_facto.py | 8 +- nerfstudio/models/semantic_nerfw.py | 5 +- .../datasets/process_nuscenes_masks.py | 1 - nerfstudio/scripts/process_data.py | 6 +- nerfstudio/scripts/render.py | 9 +- nerfstudio/utils/colormaps.py | 31 +++-- nerfstudio/utils/colors.py | 9 +- nerfstudio/utils/math.py | 49 ++++---- nerfstudio/utils/plotly_utils.py | 15 +-- nerfstudio/utils/poses.py | 11 +- nerfstudio/utils/tensor_dataclass.py | 16 ++- nerfstudio/utils/writer.py | 18 ++- nerfstudio/viewer/app/run_deploy.py | 1 - nerfstudio/viewer/server/state/node.py | 5 +- nerfstudio/viewer/server/state/state_node.py | 1 - nerfstudio/viewer/server/viewer_utils.py | 5 +- pyproject.toml | 4 +- .../dataparsers/test_nerfstudio_dataparser.py | 1 - tests/field_components/test_fields.py | 4 +- tests/field_components/test_temporal_grid.py | 5 +- tests/utils/test_tensor_dataclass.py | 4 +- 65 files changed, 597 insertions(+), 531 deletions(-) diff --git a/docs/conf.py b/docs/conf.py index 824a61d78..6c5028b9b 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -16,7 +16,6 @@ import sys sys.path.insert(0, os.path.abspath("..")) sys.path.append(os.path.abspath("./_pygments")) -import nerfstudio.configs.base_config # -- Project information ----------------------------------------------------- diff --git a/nerfstudio/cameras/camera_optimizers.py b/nerfstudio/cameras/camera_optimizers.py index 8cba97f26..cd6611aab 100644 --- a/nerfstudio/cameras/camera_optimizers.py +++ b/nerfstudio/cameras/camera_optimizers.py @@ -24,8 +24,8 @@ from typing import Literal, Type, Union import torch import tyro -from torch import nn -from torchtyping import TensorType +from jaxtyping import Float, Int +from torch import Tensor, nn from typing_extensions import assert_never from nerfstudio.cameras.lie_groups import exp_map_SE3, exp_map_SO3xR3 @@ -101,8 +101,8 @@ class CameraOptimizer(nn.Module): def forward( self, - indices: TensorType["num_cameras"], - ) -> TensorType["num_cameras", 3, 4]: + indices: Int[Tensor, "num_cameras"], + ) -> Float[Tensor, "num_cameras 3 4"]: """Indexing into camera adjustments. Args: indices: indices of Cameras to optimize. diff --git a/nerfstudio/cameras/camera_utils.py b/nerfstudio/cameras/camera_utils.py index 548e5d2dc..827f15d5a 100644 --- a/nerfstudio/cameras/camera_utils.py +++ b/nerfstudio/cameras/camera_utils.py @@ -21,8 +21,9 @@ from typing import List, Literal, Optional, Tuple import numpy as np import torch +from jaxtyping import Float from numpy.typing import ArrayLike -from torchtyping import TensorType +from torch import Tensor _EPS = np.finfo(float).eps * 4.0 @@ -184,7 +185,7 @@ def get_interpolated_poses(pose_a: ArrayLike, pose_b: ArrayLike, steps: int = 10 return poses_ab -def get_interpolated_k(k_a, k_b, steps: int = 10) -> TensorType[3, 4]: +def get_interpolated_k(k_a, k_b, steps: int = 10) -> Float[Tensor, "3 4"]: """ Returns interpolated path between two camera poses with specified number of steps. @@ -202,9 +203,9 @@ def get_interpolated_k(k_a, k_b, steps: int = 10) -> TensorType[3, 4]: def get_ordered_poses_and_k( - poses: TensorType["num_poses", 3, 4], - Ks: TensorType["num_poses", 3, 3], -) -> Tuple[TensorType["num_poses", 3, 4], TensorType["num_poses", 3, 3]]: + poses: Float[Tensor, "num_poses 3 4"], + Ks: Float[Tensor, "num_poses 3 3"], +) -> Tuple[Float[Tensor, "num_poses 3 4"], Float[Tensor, "num_poses 3 3"]]: """ Returns ordered poses and intrinsics by euclidian distance between poses. @@ -238,11 +239,11 @@ def get_ordered_poses_and_k( def get_interpolated_poses_many( - poses: TensorType["num_poses", 3, 4], - Ks: TensorType["num_poses", 3, 3], + poses: Float[Tensor, "num_poses 3 4"], + Ks: Float[Tensor, "num_poses 3 3"], steps_per_transition: int = 10, order_poses: bool = False, -) -> Tuple[TensorType["num_poses", 3, 4], TensorType["num_poses", 3, 3]]: +) -> Tuple[Float[Tensor, "num_poses 3 4"], Float[Tensor, "num_poses 3 3"]]: """Return interpolated poses for many camera poses. Args: @@ -273,7 +274,7 @@ def get_interpolated_poses_many( return torch.tensor(traj, dtype=torch.float32), torch.tensor(k_interp, dtype=torch.float32) -def normalize(x: torch.Tensor) -> TensorType[...]: +def normalize(x: torch.Tensor) -> Float[Tensor, "*batch"]: """Returns a normalized vector.""" return x / torch.linalg.norm(x) @@ -293,7 +294,7 @@ def normalize_with_norm(x: torch.Tensor, dim: int) -> Tuple[torch.Tensor, torch. return x / norm, norm -def viewmatrix(lookat: torch.Tensor, up: torch.Tensor, pos: torch.Tensor) -> TensorType[...]: +def viewmatrix(lookat: torch.Tensor, up: torch.Tensor, pos: torch.Tensor) -> Float[Tensor, "*batch"]: """Returns a camera transformation matrix. Args: @@ -319,7 +320,7 @@ def get_distortion_params( k4: float = 0.0, p1: float = 0.0, p2: float = 0.0, -) -> TensorType[...]: +) -> Float[Tensor, "*batch"]: """Returns a distortion parameters matrix. Args: @@ -442,7 +443,7 @@ def radial_and_tangential_undistort( return torch.stack([x, y], dim=-1) -def rotation_matrix(a: TensorType[3], b: TensorType[3]) -> TensorType[3, 3]: +def rotation_matrix(a: Float[Tensor, "3"], b: Float[Tensor, "3"]) -> Float[Tensor, "3 3"]: """Compute the rotation matrix that rotates vector a to vector b. Args: @@ -470,7 +471,7 @@ def rotation_matrix(a: TensorType[3], b: TensorType[3]) -> TensorType[3, 3]: return torch.eye(3) + skew_sym_mat + skew_sym_mat @ skew_sym_mat * ((1 - c) / (s**2 + 1e-8)) -def focus_of_attention(poses: TensorType["num_poses":..., 4, 4], initial_focus: TensorType[3]) -> TensorType[3]: +def focus_of_attention(poses: Float[Tensor, "*num_poses 4 4"], initial_focus: Float[Tensor, "3"]) -> Float[Tensor, "3"]: """Compute the focus of attention of a set of cameras. Only cameras that have the focus of attention in front of them are considered. @@ -509,10 +510,10 @@ def focus_of_attention(poses: TensorType["num_poses":..., 4, 4], initial_focus: def auto_orient_and_center_poses( - poses: TensorType["num_poses":..., 4, 4], + poses: Float[Tensor, "*num_poses 4 4"], method: Literal["pca", "up", "vertical", "none"] = "up", center_method: Literal["poses", "focus", "none"] = "poses", -) -> Tuple[TensorType["num_poses":..., 3, 4], TensorType[3, 4]]: +) -> Tuple[Float[Tensor, "*num_poses 3 4"], Float[Tensor, "3 4"]]: """Orients and centers the poses. We provide two methods for orientation: pca and up. pca: Orient the poses so that the principal directions of the camera centers are aligned diff --git a/nerfstudio/cameras/cameras.py b/nerfstudio/cameras/cameras.py index c0d114908..cb4e238a7 100644 --- a/nerfstudio/cameras/cameras.py +++ b/nerfstudio/cameras/cameras.py @@ -24,8 +24,9 @@ from typing import Dict, List, Optional, Tuple, Union import cv2 import torch import torchvision +from jaxtyping import Float, Int, Shaped +from torch import Tensor from torch.nn import Parameter -from torchtyping import TensorType import nerfstudio.utils.math import nerfstudio.utils.poses as pose_utils @@ -77,44 +78,44 @@ class Cameras(TensorDataclass): and will be broadcasted to the rays generated from any derivative RaySamples we create with this """ - camera_to_worlds: TensorType["num_cameras":..., 3, 4] - fx: TensorType["num_cameras":..., 1] - fy: TensorType["num_cameras":..., 1] - cx: TensorType["num_cameras":..., 1] - cy: TensorType["num_cameras":..., 1] - width: TensorType["num_cameras":..., 1] - height: TensorType["num_cameras":..., 1] - distortion_params: Optional[TensorType["num_cameras":..., 6]] - camera_type: TensorType["num_cameras":..., 1] - times: Optional[TensorType["num_cameras", 1]] + camera_to_worlds: Float[Tensor, "*num_cameras 3 4"] + fx: Float[Tensor, "*num_cameras 1"] + fy: Float[Tensor, "*num_cameras 1"] + cx: Float[Tensor, "*num_cameras 1"] + cy: Float[Tensor, "*num_cameras 1"] + width: Shaped[Tensor, "*num_cameras 1"] + height: Shaped[Tensor, "*num_cameras 1"] + distortion_params: Optional[Float[Tensor, "*num_cameras 6"]] + camera_type: Int[Tensor, "*num_cameras 1"] + times: Optional[Float[Tensor, "num_cameras 1"]] metadata: Optional[Dict] def __init__( self, - camera_to_worlds: TensorType["batch_c2ws":..., 3, 4], - fx: Union[TensorType["batch_fxs":..., 1], float], - fy: Union[TensorType["batch_fys":..., 1], float], - cx: Union[TensorType["batch_cxs":..., 1], float], - cy: Union[TensorType["batch_cys":..., 1], float], - width: Optional[Union[TensorType["batch_ws":..., 1], int]] = None, - height: Optional[Union[TensorType["batch_hs":..., 1], int]] = None, - distortion_params: Optional[TensorType["batch_dist_params":..., 6]] = None, + camera_to_worlds: Float[Tensor, "*batch_c2ws 3 4"], + fx: Union[Float[Tensor, "*batch_fxs 1"], float], + fy: Union[Float[Tensor, "*batch_fys 1"], float], + cx: Union[Float[Tensor, "*batch_cxs 1"], float], + cy: Union[Float[Tensor, "*batch_cys 1"], float], + width: Optional[Union[Shaped[Tensor, "*batch_ws 1"], int]] = None, + height: Optional[Union[Shaped[Tensor, "*batch_hs 1"], int]] = None, + distortion_params: Optional[Float[Tensor, "*batch_dist_params 6"]] = None, camera_type: Optional[ Union[ - TensorType["batch_cam_types":..., 1], + Int[Tensor, "*batch_cam_types 1"], int, List[CameraType], CameraType, ] ] = CameraType.PERSPECTIVE, - times: Optional[TensorType["num_cameras"]] = None, + times: Optional[Float[Tensor, "num_cameras"]] = None, metadata: Optional[Dict] = None, ) -> None: """Initializes the Cameras object. - Note on Input Tensor Dimensions: All of these tensors have items of dimensions TensorType[3, 4] - (in the case of the c2w matrices), TensorType[6] (in the case of distortion params), or - TensorType[1] (in the case of the rest of the elements). The dimensions before that are + Note on Input Tensor Dimensions: All of these tensors have items of dimensions Shaped[Tensor, "3 4"] + (in the case of the c2w matrices), Shaped[Tensor, "6"] (in the case of distortion params), or + Shaped[Tensor, "1"] (in the case of the rest of the elements). The dimensions before that are considered the batch dimension of that tensor (batch_c2ws, batch_fxs, etc.). We will broadcast all the tensors to be the same batch dimension. This means you can use any combination of the input types in the function signature and it won't break. Your batch size for all tensors @@ -172,9 +173,9 @@ class Cameras(TensorDataclass): def _init_get_camera_type( self, camera_type: Union[ - TensorType["batch_cam_types":..., 1], TensorType["batch_cam_types":...], int, List[CameraType], CameraType + Int[Tensor, "*batch_cam_types 1"], Int[Tensor, "*batch_cam_types"], int, List[CameraType], CameraType ], - ) -> TensorType["num_cameras":..., 1]: + ) -> Int[Tensor, "*num_cameras 1"]: """ Parses the __init__() argument camera_type @@ -213,9 +214,9 @@ class Cameras(TensorDataclass): def _init_get_height_width( self, - h_w: Union[TensorType["batch_hws":..., 1], TensorType["batch_hws":...], int, None], - c_x_y: TensorType["batch_cxys":...], - ) -> TensorType["num_cameras":..., 1]: + h_w: Union[Shaped[Tensor, "*batch_hws 1"], Shaped[Tensor, "*batch_hws"], int, None], + c_x_y: Shaped[Tensor, "*batch_cxys"], + ) -> Shaped[Tensor, "*num_cameras 1"]: """ Parses the __init__() argument for height or width @@ -260,12 +261,12 @@ class Cameras(TensorDataclass): return self.camera_to_worlds.device @property - def image_height(self) -> TensorType["num_cameras":..., 1]: + def image_height(self) -> Shaped[Tensor, "*num_cameras 1"]: """Returns the height of the images.""" return self.height @property - def image_width(self) -> TensorType["num_cameras":..., 1]: + def image_width(self) -> Shaped[Tensor, "*num_cameras 1"]: """Returns the height of the images.""" return self.width @@ -281,7 +282,7 @@ class Cameras(TensorDataclass): def get_image_coords( self, pixel_offset: float = 0.5, index: Optional[Tuple] = None - ) -> TensorType["height", "width", 2]: + ) -> Float[Tensor, "height width 2"]: """This gets the image coordinates of one of the cameras in this object. If no index is specified, it will return the maximum possible sized height / width image coordinate map, @@ -309,10 +310,10 @@ class Cameras(TensorDataclass): def generate_rays( # pylint: disable=too-many-statements self, - camera_indices: Union[TensorType["num_rays":..., "num_cameras_batch_dims"], int], - coords: Optional[TensorType["num_rays":..., 2]] = None, - camera_opt_to_camera: Optional[TensorType["num_rays":..., 3, 4]] = None, - distortion_params_delta: Optional[TensorType["num_rays":..., 6]] = None, + camera_indices: Union[Int[Tensor, "*num_rays num_cameras_batch_dims"], int], + coords: Optional[Float[Tensor, "*num_rays 2"]] = None, + camera_opt_to_camera: Optional[Float[Tensor, "*num_rays 3 4"]] = None, + distortion_params_delta: Optional[Float[Tensor, "*num_rays 6"]] = None, keep_shape: Optional[bool] = None, disable_distortion: bool = False, aabb_box: Optional[SceneBox] = None, @@ -487,10 +488,10 @@ class Cameras(TensorDataclass): # pylint: disable=too-many-statements def _generate_rays_from_coords( self, - camera_indices: TensorType["num_rays":..., "num_cameras_batch_dims"], - coords: TensorType["num_rays":..., 2], - camera_opt_to_camera: Optional[TensorType["num_rays":..., 3, 4]] = None, - distortion_params_delta: Optional[TensorType["num_rays":..., 6]] = None, + camera_indices: Int[Tensor, "*num_rays num_cameras_batch_dims"], + coords: Float[Tensor, "*num_rays 2"], + camera_opt_to_camera: Optional[Float[Tensor, "*num_rays 3 4"]] = None, + distortion_params_delta: Optional[Float[Tensor, "*num_rays 6"]] = None, disable_distortion: bool = False, ) -> RayBundle: """Generates rays for the given camera indices and coords where self isn't jagged @@ -734,7 +735,7 @@ class Cameras(TensorDataclass): ) def to_json( - self, camera_idx: int, image: Optional[TensorType["height", "width", 2]] = None, max_size: Optional[int] = None + self, camera_idx: int, image: Optional[Float[Tensor, "height width 2"]] = None, max_size: Optional[int] = None ) -> Dict: """Convert a camera to a json dictionary. @@ -768,7 +769,7 @@ class Cameras(TensorDataclass): json_["image"] = str("data:image/jpeg;base64," + base64.b64encode(data).decode("ascii")) return json_ - def get_intrinsics_matrices(self) -> TensorType["num_cameras":..., 3, 3]: + def get_intrinsics_matrices(self) -> Float[Tensor, "*num_cameras 3 3"]: """Returns the intrinsic matrices for each camera. Returns: @@ -783,7 +784,7 @@ class Cameras(TensorDataclass): return K def rescale_output_resolution( - self, scaling_factor: Union[TensorType["num_cameras":...], TensorType["num_cameras":..., 1], float, int] + self, scaling_factor: Union[Shaped[Tensor, "*num_cameras"], Shaped[Tensor, "*num_cameras 1"], float, int] ) -> None: """Rescale the output resolution of the cameras. diff --git a/nerfstudio/cameras/lie_groups.py b/nerfstudio/cameras/lie_groups.py index d6b033c32..c0fdfbef0 100644 --- a/nerfstudio/cameras/lie_groups.py +++ b/nerfstudio/cameras/lie_groups.py @@ -16,11 +16,12 @@ Helper for Lie group operations. Currently only used for pose optimization. """ import torch -from torchtyping import TensorType +from jaxtyping import Float +from torch import Tensor # We make an exception on snake case conventions because SO3 != so3. -def exp_map_SO3xR3(tangent_vector: TensorType["b", 6]) -> TensorType["b", 3, 4]: # pylint: disable=invalid-name +def exp_map_SO3xR3(tangent_vector: Float[Tensor, "b 6"]) -> Float[Tensor, "b 3 4"]: # pylint: disable=invalid-name """Compute the exponential map of the direct product group `SO(3) x R^3`. This can be used for learning pose deltas on SE(3), and is generally faster than `exp_map_SE3`. @@ -58,7 +59,7 @@ def exp_map_SO3xR3(tangent_vector: TensorType["b", 6]) -> TensorType["b", 3, 4]: return ret -def exp_map_SE3(tangent_vector: TensorType["b", 6]) -> TensorType["b", 3, 4]: # pylint: disable=invalid-name +def exp_map_SE3(tangent_vector: Float[Tensor, "b 6"]) -> Float[Tensor, "b 3 4"]: # pylint: disable=invalid-name """Compute the exponential map `se(3) -> SE(3)`. This can be used for learning pose deltas on `SE(3)`. diff --git a/nerfstudio/cameras/rays.py b/nerfstudio/cameras/rays.py index 6d8c76585..18bcabd14 100644 --- a/nerfstudio/cameras/rays.py +++ b/nerfstudio/cameras/rays.py @@ -20,7 +20,8 @@ from dataclasses import dataclass from typing import Callable, Dict, Optional, Tuple, Union import torch -from torchtyping import TensorType +from jaxtyping import Float, Int, Shaped +from torch import Tensor from nerfstudio.utils.math import Gaussians, conical_frustum_to_gaussian from nerfstudio.utils.tensor_dataclass import TensorDataclass @@ -32,20 +33,20 @@ TORCH_DEVICE = Union[str, torch.device] # pylint: disable=invalid-name class Frustums(TensorDataclass): """Describes region of space as a frustum.""" - origins: TensorType["bs":..., 3] + origins: Float[Tensor, "*bs 3"] """xyz coordinate for ray origin.""" - directions: TensorType["bs":..., 3] + directions: Float[Tensor, "*bs 3"] """Direction of ray.""" - starts: TensorType["bs":..., 1] + starts: Float[Tensor, "*bs 1"] """Where the frustum starts along a ray.""" - ends: TensorType["bs":..., 1] + ends: Float[Tensor, "*bs 1"] """Where the frustum ends along a ray.""" - pixel_area: TensorType["bs":..., 1] + pixel_area: Float[Tensor, "*bs 1"] """Projected area of pixel a distance 1 away from origin.""" - offsets: Optional[TensorType["bs":..., 3]] = None + offsets: Optional[Float[Tensor, "*bs 3"]] = None """Offsets for each sample position""" - def get_positions(self) -> TensorType[..., 3]: + def get_positions(self) -> Float[Tensor, "*batch 3"]: """Calculates "center" position of frustum. Not weighted by mass. Returns: @@ -56,7 +57,7 @@ class Frustums(TensorDataclass): pos = pos + self.offsets return pos - def get_start_positions(self) -> TensorType[..., 3]: + def get_start_positions(self) -> Float[Tensor, "*batch 3"]: """Calculates "start" position of frustum. Returns: @@ -108,23 +109,23 @@ class RaySamples(TensorDataclass): frustums: Frustums """Frustums along ray.""" - camera_indices: Optional[TensorType["bs":..., 1]] = None + camera_indices: Optional[Int[Tensor, "*bs 1"]] = None """Camera index.""" - deltas: Optional[TensorType["bs":..., 1]] = None + deltas: Optional[Float[Tensor, "*bs 1"]] = None """"width" of each sample.""" - spacing_starts: Optional[TensorType["bs":..., "num_samples", 1]] = None + spacing_starts: Optional[Float[Tensor, "*bs num_samples 1"]] = None """Start of normalized bin edges along ray [0,1], before warping is applied, ie. linear in disparity sampling.""" - spacing_ends: Optional[TensorType["bs":..., "num_samples", 1]] = None + spacing_ends: Optional[Float[Tensor, "*bs num_samples 1"]] = None """Start of normalized bin edges along ray [0,1], before warping is applied, ie. linear in disparity sampling.""" spacing_to_euclidean_fn: Optional[Callable] = None """Function to convert bins to euclidean distance.""" - metadata: Optional[Dict[str, TensorType["bs":..., "latent_dims"]]] = None + metadata: Optional[Dict[str, Shaped[Tensor, "*bs latent_dims"]]] = None """additional information relevant to generating ray samples""" - times: Optional[TensorType[..., 1]] = None + times: Optional[Float[Tensor, "*batch 1"]] = None """Times at which rays are sampled""" - def get_weights(self, densities: TensorType[..., "num_samples", 1]) -> TensorType[..., "num_samples", 1]: + def get_weights(self, densities: Float[Tensor, "*batch num_samples 1"]) -> Float[Tensor, "*batch num_samples 1"]: """Return weights based on predicted densities Args: @@ -150,8 +151,8 @@ class RaySamples(TensorDataclass): @staticmethod def get_weights_and_transmittance_from_alphas( - alphas: TensorType[..., "num_samples", 1], weights_only: bool = False - ) -> Tuple[TensorType[..., "num_samples", 1], TensorType[..., "num_samples", 1]]: + alphas: Float[Tensor, "*batch num_samples 1"], weights_only: bool = False + ) -> Tuple[Float[Tensor, "*batch num_samples 1"], Float[Tensor, "*batch num_samples 1"]]: """Return weights based on predicted alphas Args: alphas: Predicted alphas (maybe from sdf) for samples along ray @@ -175,21 +176,21 @@ class RayBundle(TensorDataclass): """A bundle of ray parameters.""" # TODO(ethan): make sure the sizes with ... are correct - origins: TensorType[..., 3] + origins: Float[Tensor, "*batch 3"] """Ray origins (XYZ)""" - directions: TensorType[..., 3] + directions: Float[Tensor, "*batch 3"] """Unit ray direction vector""" - pixel_area: TensorType[..., 1] + pixel_area: Float[Tensor, "*batch 1"] """Projected area of pixel a distance 1 away from origin""" - camera_indices: Optional[TensorType[..., 1]] = None + camera_indices: Optional[Int[Tensor, "*batch 1"]] = None """Camera indices""" - nears: Optional[TensorType[..., 1]] = None + nears: Optional[Float[Tensor, "*batch 1"]] = None """Distance along ray to start sampling""" - fars: Optional[TensorType[..., 1]] = None + fars: Optional[Float[Tensor, "*batch 1"]] = None """Rays Distance along ray to stop sampling""" - metadata: Optional[Dict[str, TensorType["num_rays", "latent_dims"]]] = None + metadata: Optional[Dict[str, Shaped[Tensor, "num_rays latent_dims"]]] = None """Additional metadata or data needed for interpolation, will mimic shape of rays""" - times: Optional[TensorType[..., 1]] = None + times: Optional[Float[Tensor, "*batch 1"]] = None """Times at which rays are sampled""" def set_camera_indices(self, camera_index: int) -> None: @@ -232,10 +233,10 @@ class RayBundle(TensorDataclass): def get_ray_samples( self, - bin_starts: TensorType["bs":..., "num_samples", 1], - bin_ends: TensorType["bs":..., "num_samples", 1], - spacing_starts: Optional[TensorType["bs":..., "num_samples", 1]] = None, - spacing_ends: Optional[TensorType["bs":..., "num_samples", 1]] = None, + bin_starts: Float[Tensor, "*bs num_samples 1"], + bin_ends: Float[Tensor, "*bs num_samples 1"], + spacing_starts: Optional[Float[Tensor, "*bs num_samples 1"]] = None, + spacing_ends: Optional[Float[Tensor, "*bs num_samples 1"]] = None, spacing_to_euclidean_fn: Optional[Callable] = None, ) -> RaySamples: """Produces samples for each ray by projection points along the ray direction. Currently samples uniformly. diff --git a/nerfstudio/data/dataparsers/base_dataparser.py b/nerfstudio/data/dataparsers/base_dataparser.py index 1d1c87a9a..1260833b9 100644 --- a/nerfstudio/data/dataparsers/base_dataparser.py +++ b/nerfstudio/data/dataparsers/base_dataparser.py @@ -23,7 +23,8 @@ from pathlib import Path from typing import Any, Dict, List, Literal, Optional, Type import torch -from torchtyping import TensorType +from jaxtyping import Float +from torch import Tensor import nerfstudio.configs.base_config as cfg from nerfstudio.cameras.cameras import Cameras @@ -39,7 +40,7 @@ class Semantics: """filenames to load semantic data""" classes: List[str] """class labels for data""" - colors: torch.Tensor + colors: Tensor """color mapping for classes""" mask_classes: List[str] = field(default_factory=lambda: []) """classes to mask out from training for all modalities""" @@ -54,9 +55,9 @@ class DataparserOutputs: """Filenames for the images.""" cameras: Cameras """Camera object storing collection of camera information in dataset.""" - alpha_color: Optional[TensorType[3]] = None + alpha_color: Optional[Float[Tensor, "3"]] = None """Color of dataset background.""" - scene_box: SceneBox = SceneBox() + scene_box: SceneBox = SceneBox(aabb=torch.tensor([[-1, -1, -1], [1, 1, 1]])) """Scene box of dataset. Used to bound the scene or provide the scene scale depending on model.""" mask_filenames: Optional[List[Path]] = None """Filenames for any masks that are required""" @@ -64,7 +65,7 @@ class DataparserOutputs: """Dictionary of any metadata that be required for the given experiment. Will be processed by the InputDataset to create any additional tensors that may be required. """ - dataparser_transform: TensorType[3, 4] = torch.eye(4)[:3, :] + dataparser_transform: Float[Tensor, "3 4"] = torch.eye(4)[:3, :] """Transform applied by the dataparser.""" dataparser_scale: float = 1.0 """Scale applied by the dataparser.""" @@ -91,9 +92,9 @@ class DataparserOutputs: def transform_poses_to_original_space( self, - poses: TensorType["num_poses", 3, 4], + poses: Float[Tensor, "num_poses 3 4"], camera_convention: Literal["opengl", "opencv"] = "opencv", - ) -> TensorType["num_poses", 3, 4]: + ) -> Float[Tensor, "num_poses 3 4"]: """ Transforms the poses in the transformed space back to the original world coordinate system. Args: @@ -164,11 +165,11 @@ class DataParser: def transform_poses_to_original_space( - poses: TensorType["num_poses", 3, 4], - applied_transform: TensorType[3, 4], + poses: Float[Tensor, "num_poses 3 4"], + applied_transform: Float[Tensor, "3 4"], applied_scale: float, camera_convention: Literal["opengl", "opencv"] = "opencv", -) -> TensorType["num_poses", 3, 4]: +) -> Float[Tensor, "num_poses 3 4"]: """ Transforms the poses in the transformed space back to the original world coordinate system. Args: diff --git a/nerfstudio/data/datasets/base_dataset.py b/nerfstudio/data/datasets/base_dataset.py index 5377459aa..f3245c675 100644 --- a/nerfstudio/data/datasets/base_dataset.py +++ b/nerfstudio/data/datasets/base_dataset.py @@ -24,9 +24,10 @@ from typing import Dict, List import numpy as np import numpy.typing as npt import torch +from jaxtyping import Float from PIL import Image +from torch import Tensor from torch.utils.data import Dataset -from torchtyping import TensorType from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs from nerfstudio.data.utils.data_utils import get_image_mask_tensor_from_path @@ -73,7 +74,7 @@ class InputDataset(Dataset): assert image.shape[2] in [3, 4], f"Image shape of {image.shape} is in correct." return image - def get_image(self, image_idx: int) -> TensorType["image_height", "image_width", "num_channels"]: + def get_image(self, image_idx: int) -> Float[Tensor, "image_height image_width num_channels"]: """Returns a 3 channel image. Args: diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 12e3b477a..59a7d18fe 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -20,7 +20,8 @@ import random from typing import Dict, Optional, Union import torch -from torchtyping import TensorType +from jaxtyping import Int +from torch import Tensor class PixelSampler: # pylint: disable=too-few-public-methods @@ -50,9 +51,9 @@ class PixelSampler: # pylint: disable=too-few-public-methods num_images: int, image_height: int, image_width: int, - mask: Optional[TensorType] = None, + mask: Optional[Tensor] = None, device: Union[torch.device, str] = "cpu", - ) -> TensorType["batch_size", 3]: + ) -> Int[Tensor, "batch_size 3"]: """ Naive pixel sampler, uniformly samples across all possible pixels of all possible images. @@ -218,9 +219,9 @@ class EquirectangularPixelSampler(PixelSampler): # pylint: disable=too-few-publ num_images: int, image_height: int, image_width: int, - mask: Optional[TensorType] = None, + mask: Optional[Tensor] = None, device: Union[torch.device, str] = "cpu", - ) -> TensorType["batch_size", 3]: + ) -> Int[Tensor, "batch_size 3"]: if isinstance(mask, torch.Tensor): # Note: if there is a mask, sampling reduces back to uniform sampling, which gives more # sampling weight to the poles of the image than the equators. @@ -274,9 +275,9 @@ class PatchPixelSampler(PixelSampler): # pylint: disable=too-few-public-methods num_images: int, image_height: int, image_width: int, - mask: Optional[TensorType] = None, + mask: Optional[Tensor] = None, device: Union[torch.device, str] = "cpu", - ) -> TensorType["batch_size", 3]: + ) -> Int[Tensor, "batch_size 3"]: if mask: # Note: if there is a mask, sampling reduces back to uniform sampling indices = super().sample_method(batch_size, num_images, image_height, image_width, mask=mask, device=device) diff --git a/nerfstudio/data/scene_box.py b/nerfstudio/data/scene_box.py index 8ac9488bb..aaa055645 100644 --- a/nerfstudio/data/scene_box.py +++ b/nerfstudio/data/scene_box.py @@ -20,14 +20,15 @@ from dataclasses import dataclass from typing import Union import torch -from torchtyping import TensorType +from jaxtyping import Float +from torch import Tensor @dataclass class SceneBox: """Data to represent the scene box.""" - aabb: TensorType[2, 3] = None + aabb: Float[Tensor, "2 3"] """aabb: axis-aligned bounding box. aabb[0] is the minimum (x,y,z) point. aabb[1] is the maximum (x,y,z) point.""" @@ -53,7 +54,7 @@ class SceneBox: return SceneBox(aabb=(self.aabb - self.get_center()) * scale_factor) @staticmethod - def get_normalized_positions(positions: TensorType[..., 3], aabb: TensorType[2, 3]): + def get_normalized_positions(positions: Float[Tensor, "*batch 3"], aabb: Float[Tensor, "2 3"]): """Return normalized positions in range [0, 1] based on the aabb axis-aligned bounding box. Args: @@ -65,7 +66,7 @@ class SceneBox: return normalized_positions @staticmethod - def from_camera_poses(poses: TensorType[..., 3, 4], scale_factor: float) -> "SceneBox": + def from_camera_poses(poses: Float[Tensor, "*batch 3 4"], scale_factor: float) -> "SceneBox": """Returns the instance of SceneBox that fully envelopes a set of poses Args: diff --git a/nerfstudio/data/utils/colmap_parsing_utils.py b/nerfstudio/data/utils/colmap_parsing_utils.py index 3ca853a27..2b7c48cc5 100644 --- a/nerfstudio/data/utils/colmap_parsing_utils.py +++ b/nerfstudio/data/utils/colmap_parsing_utils.py @@ -40,7 +40,6 @@ TODO(1480) Delete this file when moving to pycolmap. # # Author: Johannes L. Schoenberger (jsch-at-demuc-dot-de) -import argparse import collections import os import struct diff --git a/nerfstudio/exporter/exporter_utils.py b/nerfstudio/exporter/exporter_utils.py index 943331cd7..3e1088eab 100644 --- a/nerfstudio/exporter/exporter_utils.py +++ b/nerfstudio/exporter/exporter_utils.py @@ -28,6 +28,7 @@ import numpy as np import open3d as o3d import pymeshlab import torch +from jaxtyping import Float from rich.progress import ( BarColumn, Progress, @@ -35,7 +36,7 @@ from rich.progress import ( TextColumn, TimeRemainingColumn, ) -from torchtyping import TensorType +from torch import Tensor from nerfstudio.cameras.cameras import Cameras from nerfstudio.data.datasets.base_dataset import InputDataset @@ -47,13 +48,13 @@ from nerfstudio.utils.rich_utils import CONSOLE, ItersPerSecColumn class Mesh: """Class for a mesh.""" - vertices: TensorType["num_verts", 3] + vertices: Float[Tensor, "num_verts 3"] """Vertices of the mesh.""" - faces: TensorType["num_faces", 3] + faces: Float[Tensor, "num_faces 3"] """Faces of the mesh.""" - normals: TensorType["num_verts", 3] + normals: Float[Tensor, "num_verts 3"] """Normals of the mesh.""" - colors: Optional[TensorType["num_verts", 3]] = None + colors: Optional[Float[Tensor, "num_verts 3"]] = None """Colors of the mesh.""" diff --git a/nerfstudio/exporter/marching_cubes.py b/nerfstudio/exporter/marching_cubes.py index 9f0d3554c..7b38d619e 100644 --- a/nerfstudio/exporter/marching_cubes.py +++ b/nerfstudio/exporter/marching_cubes.py @@ -22,16 +22,15 @@ from typing import Callable, List, Tuple, Union import numpy as np import torch import trimesh +from jaxtyping import Bool, Float from skimage import measure -from torchtyping import TensorType +from torch import Tensor avg_pool_3d = torch.nn.AvgPool3d(2, stride=2) upsample = torch.nn.Upsample(scale_factor=2, mode="nearest") -def create_point_pyramid( - points: TensorType[3, "height", "width", "depth"] -) -> List[TensorType[3, "height", "width", "depth"]]: +def create_point_pyramid(points: Float[Tensor, "3 height width depth"]) -> List[Float[Tensor, "3 height width depth"]]: """ Create a point pyramid for multi-resolution evaluation. @@ -49,7 +48,7 @@ def create_point_pyramid( return points_pyramid -def evaluate_sdf(sdf: Callable, points: TensorType["batch", 3]) -> TensorType["batch"]: +def evaluate_sdf(sdf: Callable, points: Float[Tensor, "batch 3"]) -> Float[Tensor, "batch"]: """ Evaluate a signed distance function (SDF) for a batch of points. @@ -70,12 +69,12 @@ def evaluate_sdf(sdf: Callable, points: TensorType["batch", 3]) -> TensorType["b def evaluate_multiresolution_sdf( evaluate: Callable, - points_pyramid: List[TensorType[3, "height", "width", "depth"]], - coarse_mask: Union[TensorType[1, 1, "height", "width", "depth"], None], + points_pyramid: List[Float[Tensor, "3 height width depth"]], + coarse_mask: Union[Bool[Tensor, "1 1 height width depth"], None], x_max: float, x_min: float, crop_n: int, -) -> TensorType["batch"]: +) -> Float[Tensor, "batch"]: """ Evaluate SDF values using a multi-resolution approach with a given point pyramid. @@ -135,7 +134,7 @@ def generate_mesh_with_multires_marching_cubes( bounding_box_min: Tuple[float, float, float] = (-1.0, -1.0, -1.0), bounding_box_max: Tuple[float, float, float] = (1.0, 1.0, 1.0), isosurface_threshold: float = 0.0, - coarse_mask: Union[None, TensorType["height", "width", "depth"]] = None, + coarse_mask: Union[None, Bool[Tensor, "height width depth"]] = None, ) -> trimesh.Trimesh: """ Computes the isosurface of a signed distance function (SDF) defined by the diff --git a/nerfstudio/exporter/texture_utils.py b/nerfstudio/exporter/texture_utils.py index ac76ff76e..3fde25985 100644 --- a/nerfstudio/exporter/texture_utils.py +++ b/nerfstudio/exporter/texture_utils.py @@ -28,7 +28,8 @@ import mediapy as media import numpy as np import torch import xatlas -from torchtyping import TensorType +from jaxtyping import Float +from torch import Tensor from nerfstudio.cameras.rays import RayBundle from nerfstudio.exporter.exporter_utils import Mesh @@ -39,8 +40,8 @@ TORCH_DEVICE = Union[torch.device, str] # pylint: disable=invalid-name def get_parallelogram_area( - p: TensorType["bs":..., 2], v0: TensorType["bs":..., 2], v1: TensorType["bs":..., 2] -) -> TensorType["bs":...]: + p: Float[Tensor, "*bs 2"], v0: Float[Tensor, "*bs 2"], v1: Float[Tensor, "*bs 2"] +) -> Float[Tensor, "*bs"]: """Given three 2D points, return the area defined by the parallelogram. I.e., 2x the triangle area. Args: @@ -56,7 +57,7 @@ def get_parallelogram_area( def get_texture_image( num_pixels_w: int, num_pixels_h: int, device: TORCH_DEVICE -) -> Tuple[TensorType["num_pixels_h", "num_pixels_w", 2], TensorType["num_pixels_h", "num_pixels_w", 2]]: +) -> Tuple[Float[Tensor, "num_pixels_h num_pixels_w 2"], Float[Tensor, "num_pixels_h num_pixels_w 2"]]: """Get a texture image.""" px_w = 1.0 / num_pixels_w px_h = 1.0 / num_pixels_h @@ -76,14 +77,14 @@ def get_texture_image( def unwrap_mesh_per_uv_triangle( - vertices: TensorType["num_verts", 3], - faces: TensorType["num_faces", 3], - vertex_normals: TensorType["num_verts", 3], + vertices: Float[Tensor, "num_verts 3"], + faces: Float[Tensor, "num_faces 3"], + vertex_normals: Float[Tensor, "num_verts 3"], px_per_uv_triangle: int, ) -> Tuple[ - TensorType["num_faces", 3, 2], - TensorType["num_pixels", "num_pixels", 3], - TensorType["num_pixels", "num_pixels", "num_pixels"], + Float[Tensor, "num_faces 3 2"], + Float[Tensor, "num_pixels num_pixels 3"], + Float[Tensor, "num_pixels num_pixels num_pixels"], ]: """Unwrap a mesh to a UV texture. This is done by making a grid of rectangles in the UV texture map and then having two triangles per rectangle. Then the texture image is rasterized and uses barycentric @@ -212,15 +213,15 @@ def unwrap_mesh_per_uv_triangle( def unwrap_mesh_with_xatlas( - vertices: TensorType["num_verts", 3], - faces: TensorType["num_faces", 3, torch.long], - vertex_normals: TensorType["num_verts", 3], + vertices: Float[Tensor, "num_verts 3"], + faces: Float[Tensor, "num_faces 3 torch.long"], + vertex_normals: Float[Tensor, "num_verts 3"], num_pixels_per_side=1024, num_faces_per_barycentric_chunk=10, ) -> Tuple[ - TensorType["num_faces", 3, 2], - TensorType["num_pixels", "num_pixels", 3], - TensorType["num_pixels", "num_pixels", "num_pixels"], + Float[Tensor, "num_faces 3 2"], + Float[Tensor, "num_pixels num_pixels 3"], + Float[Tensor, "num_pixels num_pixels num_pixels"], ]: """Unwrap a mesh using xatlas. We use xatlas to unwrap the mesh with UV coordinates. Then we rasterize the mesh with a square pattern. We interpolate the XYZ and normal diff --git a/nerfstudio/exporter/tsdf_utils.py b/nerfstudio/exporter/tsdf_utils.py index 3b0a752ca..ca99e80ad 100644 --- a/nerfstudio/exporter/tsdf_utils.py +++ b/nerfstudio/exporter/tsdf_utils.py @@ -28,8 +28,9 @@ import numpy as np import pymeshlab import torch import torch.nn.functional as F +from jaxtyping import Bool, Float from skimage import measure -from torchtyping import TensorType +from torch import Tensor from nerfstudio.exporter.exporter_utils import Mesh, render_trajectory from nerfstudio.pipelines.base_pipeline import Pipeline @@ -44,17 +45,17 @@ class TSDF: Class for creating TSDFs. """ - voxel_coords: TensorType[3, "xdim", "ydim", "zdim"] + voxel_coords: Float[Tensor, "3 xdim ydim zdim"] """Coordinates of each voxel in the TSDF.""" - values: TensorType["xdim", "ydim", "zdim"] + values: Float[Tensor, "xdim ydim zdim"] """TSDF values for each voxel.""" - weights: TensorType["xdim", "ydim", "zdim"] + weights: Float[Tensor, "xdim ydim zdim"] """TSDF weights for each voxel.""" - colors: TensorType["xdim", "ydim", "zdim", 3] + colors: Float[Tensor, "xdim ydim zdim 3"] """TSDF colors for each voxel.""" - voxel_size: TensorType[3] + voxel_size: Float[Tensor, "3"] """Size of each voxel in the TSDF. [x, y, z] size.""" - origin: TensorType[3] + origin: Float[Tensor, "3"] """Origin of the TSDF [xmin, ymin, zmin].""" truncation_margin: float = 5.0 """Margin for truncation.""" @@ -86,7 +87,7 @@ class TSDF: return truncation @staticmethod - def from_aabb(aabb: TensorType[2, 3], volume_dims: TensorType[3]): + def from_aabb(aabb: Float[Tensor, "2 3"], volume_dims: Float[Tensor, "3"]): """Returns an instance of TSDF from an axis-aligned bounding box and volume dimensions. Args: @@ -167,11 +168,11 @@ class TSDF: def integrate_tsdf( self, - c2w: TensorType["batch", 4, 4], - K: TensorType["batch", 3, 3], - depth_images: TensorType["batch", 1, "height", "width"], - color_images: Optional[TensorType["batch", 3, "height", "width"]] = None, - mask_images: Optional[TensorType["batch", 1, "height", "width"]] = None, + c2w: Float[Tensor, "batch 4 4"], + K: Float[Tensor, "batch 3 3"], + depth_images: Float[Tensor, "batch 1 height width"], + color_images: Optional[Float[Tensor, "batch 3 height width"]] = None, + mask_images: Optional[Bool[Tensor, "batch 1 height width"]] = None, ) -> None: """Integrates a batch of depth images into the TSDF. @@ -327,11 +328,11 @@ def export_tsdf_mesh( ) # camera extrinsics and intrinsics - c2w: TensorType["N", 3, 4] = cameras.camera_to_worlds.to(device) + c2w: Float[Tensor, "N 3 4"] = cameras.camera_to_worlds.to(device) # make c2w homogeneous c2w = torch.cat([c2w, torch.zeros(c2w.shape[0], 1, 4, device=device)], dim=1) c2w[:, 3, 3] = 1 - K: TensorType["N", 3, 3] = cameras.get_intrinsics_matrices().to(device) + K: Float[Tensor, "N 3 3"] = cameras.get_intrinsics_matrices().to(device) color_images = torch.tensor(np.array(color_images), device=device).permute(0, 3, 1, 2) # shape (N, 3, H, W) depth_images = torch.tensor(np.array(depth_images), device=device).permute(0, 3, 1, 2) # shape (N, 1, H, W) diff --git a/nerfstudio/field_components/__init__.py b/nerfstudio/field_components/__init__.py index f6b2c4f3e..2bd039624 100644 --- a/nerfstudio/field_components/__init__.py +++ b/nerfstudio/field_components/__init__.py @@ -12,7 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +# pylint:disable=useless-import-alias + """init field modules""" -from .base_field_component import FieldComponent -from .encodings import Encoding, ScalingAndOffset -from .mlp import MLP +from .base_field_component import FieldComponent as FieldComponent +from .encodings import Encoding as Encoding +from .encodings import ScalingAndOffset as ScalingAndOffset +from .mlp import MLP as MLP diff --git a/nerfstudio/field_components/base_field_component.py b/nerfstudio/field_components/base_field_component.py index e9dcaed7f..8d84b782b 100644 --- a/nerfstudio/field_components/base_field_component.py +++ b/nerfstudio/field_components/base_field_component.py @@ -18,8 +18,8 @@ The field module baseclass. from abc import abstractmethod from typing import Optional -from torch import nn -from torchtyping import TensorType +from jaxtyping import Shaped +from torch import Tensor, nn class FieldComponent(nn.Module): @@ -56,7 +56,7 @@ class FieldComponent(nn.Module): return self.out_dim @abstractmethod - def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]: + def forward(self, in_tensor: Shaped[Tensor, "*bs input_dim"]) -> Shaped[Tensor, "*bs output_dim"]: """ Returns processed tensor diff --git a/nerfstudio/field_components/embedding.py b/nerfstudio/field_components/embedding.py index 4a10b4694..cc934b099 100644 --- a/nerfstudio/field_components/embedding.py +++ b/nerfstudio/field_components/embedding.py @@ -18,7 +18,8 @@ Code for embeddings. import torch -from torchtyping import TensorType +from jaxtyping import Shaped +from torch import Tensor from nerfstudio.field_components.base_field_component import FieldComponent @@ -45,7 +46,7 @@ class Embedding(FieldComponent): """Return the mean of the embedding weights along a dim.""" return self.embedding.weight.mean(dim) - def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: + def forward(self, in_tensor: Shaped[Tensor, "*batch input_dim"]) -> Shaped[Tensor, "*batch output_dim"]: """Call forward Args: diff --git a/nerfstudio/field_components/encodings.py b/nerfstudio/field_components/encodings.py index b455468f2..22408998b 100644 --- a/nerfstudio/field_components/encodings.py +++ b/nerfstudio/field_components/encodings.py @@ -23,8 +23,8 @@ from typing import Literal, Optional, Sequence import numpy as np import torch import torch.nn.functional as F -from torch import nn -from torchtyping import TensorType +from jaxtyping import Float, Int, Shaped +from torch import Tensor, nn from nerfstudio.field_components.base_field_component import FieldComponent from nerfstudio.utils.math import components_from_spherical_harmonics, expected_sin @@ -51,7 +51,7 @@ class Encoding(FieldComponent): super().__init__(in_dim=in_dim) @abstractmethod - def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]: + def forward(self, in_tensor: Shaped[Tensor, "*bs input_dim"]) -> Shaped[Tensor, "*bs output_dim"]: """Call forward and returns and processed tensor Args: @@ -68,7 +68,7 @@ class Identity(Encoding): raise ValueError("Input dimension has not been set") return self.in_dim - def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]: + def forward(self, in_tensor: Shaped[Tensor, "*bs input_dim"]) -> Shaped[Tensor, "*bs output_dim"]: return in_tensor @@ -92,7 +92,7 @@ class ScalingAndOffset(Encoding): raise ValueError("Input dimension has not been set") return self.in_dim - def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]: + def forward(self, in_tensor: Float[Tensor, "*bs input_dim"]) -> Float[Tensor, "*bs output_dim"]: return self.scaling * in_tensor + self.offset @@ -128,9 +128,9 @@ class NeRFEncoding(Encoding): def forward( self, - in_tensor: TensorType["bs":..., "input_dim"], - covs: Optional[TensorType["bs":..., "input_dim", "input_dim"]] = None, - ) -> TensorType["bs":..., "output_dim"]: + in_tensor: Float[Tensor, "*bs input_dim"], + covs: Optional[Float[Tensor, "*bs input_dim input_dim"]] = None, + ) -> Float[Tensor, "*bs output_dim"]: """Calculates NeRF encoding. If covariances are provided the encodings will be integrated as proposed in mip-NeRF. @@ -187,9 +187,9 @@ class RFFEncoding(Encoding): def forward( self, - in_tensor: TensorType["bs":..., "input_dim"], - covs: Optional[TensorType["bs":..., "input_dim", "input_dim"]] = None, - ) -> TensorType["bs":..., "output_dim"]: + in_tensor: Float[Tensor, "*bs input_dim"], + covs: Optional[Float[Tensor, "*bs input_dim input_dim"]] = None, + ) -> Float[Tensor, "*bs output_dim"]: """Calculates RFF encoding. If covariances are provided the encodings will be integrated as proposed in mip-NeRF. @@ -285,7 +285,7 @@ class HashEncoding(Encoding): def get_out_dim(self) -> int: return self.num_levels * self.features_per_level - def hash_fn(self, in_tensor: TensorType["bs":..., "num_levels", 3]) -> TensorType["bs":..., "num_levels"]: + def hash_fn(self, in_tensor: Int[Tensor, "*bs num_levels 3"]) -> Shaped[Tensor, "*bs num_levels"]: """Returns hash tensor using method described in Instant-NGP Args: @@ -304,7 +304,7 @@ class HashEncoding(Encoding): x += self.hash_offset.to(x.device) return x - def pytorch_fwd(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]: + def pytorch_fwd(self, in_tensor: Float[Tensor, "*bs input_dim"]) -> Float[Tensor, "*bs output_dim"]: """Forward pass using pytorch. Significantly slower than TCNN implementation.""" assert in_tensor.shape[-1] == 3 @@ -347,7 +347,7 @@ class HashEncoding(Encoding): return torch.flatten(encoded_value, start_dim=-2, end_dim=-1) # [..., num_levels * features_per_level] - def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]: + def forward(self, in_tensor: Float[Tensor, "*bs input_dim"]) -> Float[Tensor, "*bs output_dim"]: if TCNN_EXISTS and self.tcnn_encoding is not None: return self.tcnn_encoding(in_tensor) return self.pytorch_fwd(in_tensor) @@ -374,7 +374,7 @@ class TensorCPEncoding(Encoding): def get_out_dim(self) -> int: return self.num_components - def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]: + def forward(self, in_tensor: Float[Tensor, "*bs input_dim"]) -> Float[Tensor, "*bs output_dim"]: line_coord = torch.stack([in_tensor[..., 2], in_tensor[..., 1], in_tensor[..., 0]]) # [3, ...] line_coord = torch.stack([torch.zeros_like(line_coord), line_coord], dim=-1) # [3, ...., 2] @@ -412,8 +412,8 @@ class TensorVMEncoding(Encoding): init_scale: Initialization scale. """ - plane_coef: TensorType[3, "num_components", "resolution", "resolution"] - line_coef: TensorType[3, "num_components", "resolution", 1] + plane_coef: Float[Tensor, "3 num_components resolution resolution"] + line_coef: Float[Tensor, "3 num_components resolution 1"] def __init__( self, @@ -432,7 +432,7 @@ class TensorVMEncoding(Encoding): def get_out_dim(self) -> int: return self.num_components * 3 - def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]: + def forward(self, in_tensor: Float[Tensor, "*bs input_dim"]) -> Float[Tensor, "*bs output_dim"]: """Compute encoding for each position in in_positions Args: @@ -497,7 +497,7 @@ class TriplaneEncoding(Encoding): product: Whether to use the element-wise product of the planes or the sum """ - plane_coef: TensorType[3, "num_components", "resolution", "resolution"] + plane_coef: Float[Tensor, "3 num_components resolution resolution"] def __init__( self, @@ -520,7 +520,7 @@ class TriplaneEncoding(Encoding): def get_out_dim(self) -> int: return self.num_components - def forward(self, in_tensor: TensorType["bs":..., 3]) -> TensorType["bs":..., "num_components", "featuresize"]: + def forward(self, in_tensor: Float[Tensor, "*bs 3"]) -> Float[Tensor, "*bs num_components featuresize"]: """Sample features from this encoder. Expects in_tensor to be in range [0, resolution]""" original_shape = in_tensor.shape @@ -617,7 +617,7 @@ class KPlanesEncoding(Encoding): def get_out_dim(self) -> int: return self.num_components - def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]: + def forward(self, in_tensor: Float[Tensor, "*bs input_dim"]) -> Float[Tensor, "*bs output_dim"]: """Sample features from this encoder. Expects ``in_tensor`` to be in range [-1, 1]""" original_shape = in_tensor.shape @@ -656,5 +656,5 @@ class SHEncoding(Encoding): return self.levels**2 @torch.no_grad() - def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]: + def forward(self, in_tensor: Float[Tensor, "*bs input_dim"]) -> Float[Tensor, "*bs output_dim"]: return components_from_spherical_harmonics(levels=self.levels, directions=in_tensor) diff --git a/nerfstudio/field_components/field_heads.py b/nerfstudio/field_components/field_heads.py index 0b668c732..ed199b4b9 100644 --- a/nerfstudio/field_components/field_heads.py +++ b/nerfstudio/field_components/field_heads.py @@ -19,8 +19,8 @@ from enum import Enum from typing import Callable, Optional, Union import torch -from torch import nn -from torchtyping import TensorType +from jaxtyping import Float, Shaped +from torch import Tensor, nn from nerfstudio.field_components.base_field_component import FieldComponent @@ -76,7 +76,7 @@ class FieldHead(FieldComponent): def _construct_net(self): self.net = nn.Linear(self.in_dim, self.out_dim) - def forward(self, in_tensor: TensorType["bs":..., "in_dim"]) -> TensorType["bs":..., "out_dim"]: + def forward(self, in_tensor: Shaped[Tensor, "*bs in_dim"]) -> Shaped[Tensor, "*bs out_dim"]: """Process network output for renderer Args: @@ -196,7 +196,7 @@ class PredNormalsFieldHead(FieldHead): def __init__(self, in_dim: Optional[int] = None, activation: Optional[nn.Module] = nn.Tanh()) -> None: super().__init__(in_dim=in_dim, out_dim=3, field_head_name=FieldHeadNames.PRED_NORMALS, activation=activation) - def forward(self, in_tensor: TensorType["bs":..., "in_dim"]) -> TensorType["bs":..., "out_dim"]: + def forward(self, in_tensor: Float[Tensor, "*bs in_dim"]) -> Float[Tensor, "*bs out_dim"]: """Needed to normalize the output into valid normals.""" out_tensor = super().forward(in_tensor) out_tensor = torch.nn.functional.normalize(out_tensor, dim=-1) diff --git a/nerfstudio/field_components/mlp.py b/nerfstudio/field_components/mlp.py index 3ce3cdad9..c51ed5f3b 100644 --- a/nerfstudio/field_components/mlp.py +++ b/nerfstudio/field_components/mlp.py @@ -18,8 +18,8 @@ Multi Layer Perceptron from typing import Optional, Set, Tuple import torch -from torch import nn -from torchtyping import TensorType +from jaxtyping import Float +from torch import Tensor, nn from nerfstudio.field_components.base_field_component import FieldComponent @@ -76,7 +76,7 @@ class MLP(FieldComponent): layers.append(nn.Linear(self.layer_width, self.out_dim)) self.layers = nn.ModuleList(layers) - def forward(self, in_tensor: TensorType["bs":..., "in_dim"]) -> TensorType["bs":..., "out_dim"]: + def forward(self, in_tensor: Float[Tensor, "*bs in_dim"]) -> Float[Tensor, "*bs out_dim"]: """Process input with a multilayer perceptron. Args: diff --git a/nerfstudio/field_components/spatial_distortions.py b/nerfstudio/field_components/spatial_distortions.py index 3a1f29731..ec9fbac56 100644 --- a/nerfstudio/field_components/spatial_distortions.py +++ b/nerfstudio/field_components/spatial_distortions.py @@ -18,8 +18,8 @@ from typing import Optional, Union import torch from functorch import jacrev, vmap -from torch import nn -from torchtyping import TensorType +from jaxtyping import Float +from torch import Tensor, nn from nerfstudio.utils.math import Gaussians @@ -27,9 +27,7 @@ from nerfstudio.utils.math import Gaussians class SpatialDistortion(nn.Module): """Apply spatial distortions""" - def forward( - self, positions: Union[TensorType["bs":..., 3], Gaussians] - ) -> Union[TensorType["bs":..., 3], Gaussians]: + def forward(self, positions: Union[Float[Tensor, "*bs 3"], Gaussians]) -> Union[Float[Tensor, "*bs 3"], Gaussians]: """ Args: positions: Sample to distort @@ -71,10 +69,12 @@ class SceneContraction(SpatialDistortion): if isinstance(positions, Gaussians): means = contract(positions.mean.clone()) - contract = lambda x: (2 - (1 / torch.linalg.norm(x, ord=self.order, dim=-1, keepdim=True))) * ( - x / torch.linalg.norm(x, ord=self.order, dim=-1, keepdim=True) - ) - jc_means = vmap(jacrev(contract))(positions.mean.view(-1, positions.mean.shape[-1])) + def contract_gauss(x): + return (2 - 1 / torch.linalg.norm(x, ord=self.order, dim=-1, keepdim=True)) * ( + x / torch.linalg.norm(x, ord=self.order, dim=-1, keepdim=True) + ) + + jc_means = vmap(jacrev(contract_gauss))(positions.mean.view(-1, positions.mean.shape[-1])) jc_means = jc_means.view(list(positions.mean.shape) + [positions.mean.shape[-1]]) # Only update covariances on positions outside the unit sphere diff --git a/nerfstudio/field_components/temporal_distortions.py b/nerfstudio/field_components/temporal_distortions.py index 36585646a..7f3e4c3dd 100644 --- a/nerfstudio/field_components/temporal_distortions.py +++ b/nerfstudio/field_components/temporal_distortions.py @@ -18,8 +18,8 @@ from enum import Enum from typing import Any, Dict, Optional, Tuple import torch -from torch import nn -from torchtyping import TensorType +from jaxtyping import Float +from torch import Tensor, nn from nerfstudio.field_components.encodings import Encoding, NeRFEncoding from nerfstudio.field_components.mlp import MLP @@ -28,7 +28,7 @@ from nerfstudio.field_components.mlp import MLP class TemporalDistortion(nn.Module): """Apply spatial distortions as a function of time""" - def forward(self, positions: TensorType["bs":..., 3], times: Optional[TensorType[1]]) -> TensorType["bs":..., 3]: + def forward(self, positions: Float[Tensor, "*bs 3"], times: Optional[Float[Tensor, "1"]]) -> Float[Tensor, "*bs 3"]: """ Args: positions: Samples to translate as a function of time diff --git a/nerfstudio/field_components/temporal_grid.py b/nerfstudio/field_components/temporal_grid.py index 5186cbb26..9d9bccc70 100644 --- a/nerfstudio/field_components/temporal_grid.py +++ b/nerfstudio/field_components/temporal_grid.py @@ -21,10 +21,10 @@ from typing import Optional import numpy as np import torch -from torch import nn +from jaxtyping import Float, Int +from torch import Tensor, nn from torch.autograd import Function from torch.cuda.amp import custom_bwd, custom_fwd -from torchtyping import TensorType import nerfstudio.field_components.cuda as _C @@ -37,16 +37,16 @@ class TemporalGridEncodeFunc(Function): @custom_fwd def forward( ctx, - inputs: TensorType["bs", "input_dim"], - temporal_row_index: TensorType["bs", "temporal_index_dim"], - embeddings: TensorType["table_size", "embed_dim"], - offsets: TensorType["num_levels+1"], + inputs: Float[Tensor, "bs input_dim"], + temporal_row_index: Float[Tensor, "bs temporal_index_dim"], + embeddings: Float[Tensor, "table_size embed_dim"], + offsets: Int[Tensor, "num_levels_plus_1"], per_level_scale: float, base_resolution: int, calc_grad_inputs: bool = False, gridtype: int = 0, align_corners: bool = False, - ) -> TensorType["bs", "output_dim"]: + ) -> Float[Tensor, "bs output_dim"]: """Call forward and interpolate the feature from embeddings Args: @@ -314,7 +314,7 @@ class TemporalGridEncoder(nn.Module): f"gridtype={self.gridtype} align_corners={self.align_corners}" ) - def get_temporal_index(self, time: TensorType["bs"]) -> TensorType["bs", "temporal_index_dim"]: + def get_temporal_index(self, time: Float[Tensor, "bs"]) -> Float[Tensor, "bs temporal_index_dim"]: """Convert the time into sampling index lists.""" row_idx_value = time * (len(self.sampling_index) - 1) row_idx = row_idx_value.long() @@ -326,7 +326,9 @@ class TemporalGridEncoder(nn.Module): temporal_row_index[mask_b] = row_idx_value - row_idx return temporal_row_index - def forward(self, xyz: TensorType["bs", "input_dim"], time: TensorType["bs", 1]) -> TensorType["bs", "output_dim"]: + def forward( + self, xyz: Float[Tensor, "bs input_dim"], time: Float[Tensor, "bs 1"] + ) -> Float[Tensor, "bs output_dim"]: """Forward and sampling feature vectors from the embedding. Args: @@ -346,7 +348,7 @@ class TemporalGridEncoder(nn.Module): ) return outputs - def get_temporal_tv_loss(self) -> TensorType[()]: + def get_temporal_tv_loss(self) -> Float[Tensor, ""]: """Apply TV loss on the temporal channels. Sample a random channel combination (i.e., row for the combination table), and then compute loss on it. diff --git a/nerfstudio/fields/base_field.py b/nerfstudio/fields/base_field.py index 9fae32c75..2392ebb12 100644 --- a/nerfstudio/fields/base_field.py +++ b/nerfstudio/fields/base_field.py @@ -21,8 +21,8 @@ from dataclasses import dataclass, field from typing import Dict, Optional, Tuple, Type import torch -from torch import nn -from torchtyping import TensorType +from jaxtyping import Float, Shaped +from torch import Tensor, nn from nerfstudio.cameras.rays import Frustums, RaySamples from nerfstudio.configs.base_config import InstantiateConfig @@ -45,7 +45,7 @@ class Field(nn.Module): self._sample_locations = None self._density_before_activation = None - def density_fn(self, positions: TensorType["bs":..., 3]) -> TensorType["bs":..., 1]: + def density_fn(self, positions: Shaped[Tensor, "*bs 3"]) -> Shaped[Tensor, "*bs 1"]: """Returns only the density. Used primarily with the density grid. Args: @@ -65,14 +65,16 @@ class Field(nn.Module): return density @abstractmethod - def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType[..., 1], TensorType[..., "num_features"]]: + def get_density( + self, ray_samples: RaySamples + ) -> Tuple[Shaped[Tensor, "*batch 1"], Float[Tensor, "*batch num_features"]]: """Computes and returns the densities. Returns a tensor of densities and a tensor of features. Args: ray_samples: Samples locations to compute density. """ - def get_normals(self) -> TensorType[..., 3]: + def get_normals(self) -> Float[Tensor, "*batch 3"]: """Computes and returns a tensor of normals. Args: @@ -97,8 +99,8 @@ class Field(nn.Module): @abstractmethod def get_outputs( - self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None - ) -> Dict[FieldHeadNames, TensorType]: + self, ray_samples: RaySamples, density_embedding: Optional[Tensor] = None + ) -> Dict[FieldHeadNames, Tensor]: """Computes and returns the colors. Returns output field values. Args: @@ -106,7 +108,7 @@ class Field(nn.Module): density_embedding: Density embeddings to condition on. """ - def forward(self, ray_samples: RaySamples, compute_normals: bool = False) -> Dict[FieldHeadNames, TensorType]: + def forward(self, ray_samples: RaySamples, compute_normals: bool = False) -> Dict[FieldHeadNames, Tensor]: """Evaluates the field at points along the ray. Args: @@ -128,7 +130,7 @@ class Field(nn.Module): return field_outputs -def shift_directions_for_tcnn(directions: TensorType["bs":..., 3]) -> TensorType["bs":..., 3]: +def shift_directions_for_tcnn(directions: Float[Tensor, "*bs 3"]) -> Float[Tensor, "*bs 3"]: """Shift directions from [-1, 1] to [0, 1] Args: diff --git a/nerfstudio/fields/density_fields.py b/nerfstudio/fields/density_fields.py index 0364db00d..40a965a67 100644 --- a/nerfstudio/fields/density_fields.py +++ b/nerfstudio/fields/density_fields.py @@ -21,7 +21,7 @@ from typing import Optional, Tuple import numpy as np import torch -from torchtyping import TensorType +from torch import Tensor from nerfstudio.cameras.rays import RaySamples from nerfstudio.data.scene_box import SceneBox @@ -49,7 +49,7 @@ class HashMLPDensityField(Field): def __init__( self, - aabb: TensorType, + aabb: Tensor, num_layers: int = 2, hidden_dim: int = 64, spatial_distortion: Optional[SpatialDistortion] = None, @@ -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) -> Tuple[TensorType, None]: + def get_density(self, ray_samples: RaySamples) -> Tuple[Tensor, 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) -> dict: + def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[Tensor] = None) -> dict: return {} diff --git a/nerfstudio/fields/instant_ngp_field.py b/nerfstudio/fields/instant_ngp_field.py index 87bfb778a..b75c539ed 100644 --- a/nerfstudio/fields/instant_ngp_field.py +++ b/nerfstudio/fields/instant_ngp_field.py @@ -21,8 +21,9 @@ from typing import Dict, Optional, Tuple import numpy as np import torch +from jaxtyping import Float +from torch import Tensor from torch.nn.parameter import Parameter -from torchtyping import TensorType from nerfstudio.cameras.rays import RaySamples from nerfstudio.data.scene_box import SceneBox @@ -63,7 +64,7 @@ class TCNNInstantNGPField(Field): def __init__( self, - aabb: TensorType, + aabb: Tensor, num_layers: int = 2, hidden_dim: int = 64, geo_feat_dim: int = 15, @@ -137,7 +138,7 @@ class TCNNInstantNGPField(Field): }, ) - def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType, TensorType]: + def get_density(self, ray_samples: RaySamples) -> Tuple[Tensor, Tensor]: if self.spatial_distortion is not None: positions = ray_samples.frustums.get_positions() positions = self.spatial_distortion(positions) @@ -160,8 +161,8 @@ class TCNNInstantNGPField(Field): return density, base_mlp_out def get_outputs( - self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None - ) -> Dict[FieldHeadNames, TensorType]: + self, ray_samples: RaySamples, density_embedding: Optional[Tensor] = None + ) -> Dict[FieldHeadNames, Tensor]: assert density_embedding is not None directions = shift_directions_for_tcnn(ray_samples.frustums.directions) directions_flat = directions.view(-1, 3) @@ -184,7 +185,7 @@ class TCNNInstantNGPField(Field): rgb = self.mlp_head(h).view(*ray_samples.frustums.directions.shape[:-1], -1).to(directions) return {FieldHeadNames.RGB: rgb} - def get_opacity(self, positions: TensorType["bs":..., 3], step_size) -> TensorType["bs":..., 1]: + def get_opacity(self, positions: Float[Tensor, "*bs 3"], step_size) -> Float[Tensor, "*bs 1"]: """Returns the opacity for a position. Used primarily by the occupancy grid. Args: diff --git a/nerfstudio/fields/nerfacto_field.py b/nerfstudio/fields/nerfacto_field.py index 840268bfc..c458d6182 100644 --- a/nerfstudio/fields/nerfacto_field.py +++ b/nerfstudio/fields/nerfacto_field.py @@ -21,9 +21,8 @@ from typing import Dict, Optional, Tuple import numpy as np import torch -from torch import nn +from torch import Tensor, nn from torch.nn.parameter import Parameter -from torchtyping import TensorType from nerfstudio.cameras.rays import RaySamples from nerfstudio.data.scene_box import SceneBox @@ -83,7 +82,7 @@ class TCNNNerfactoField(Field): def __init__( self, - aabb: TensorType, + aabb: Tensor, num_images: int, num_layers: int = 2, hidden_dim: int = 64, @@ -224,7 +223,7 @@ class TCNNNerfactoField(Field): }, ) - def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType, TensorType]: + def get_density(self, ray_samples: RaySamples) -> Tuple[Tensor, Tensor]: """Computes and returns the densities.""" if self.spatial_distortion is not None: positions = ray_samples.frustums.get_positions() @@ -251,8 +250,8 @@ class TCNNNerfactoField(Field): return density, base_mlp_out def get_outputs( - self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None - ) -> Dict[FieldHeadNames, TensorType]: + self, ray_samples: RaySamples, density_embedding: Optional[Tensor] = None + ) -> Dict[FieldHeadNames, Tensor]: assert density_embedding is not None outputs = {} if ray_samples.camera_indices is None: @@ -332,7 +331,7 @@ class TorchNerfactoField(Field): def __init__( self, - aabb: TensorType, + aabb: Tensor, num_images: int, position_encoding: Encoding = HashEncoding(), direction_encoding: Encoding = SHEncoding(), @@ -375,7 +374,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) -> Tuple[TensorType, TensorType]: + def get_density(self, ray_samples: RaySamples) -> Tuple[Tensor, Tensor]: if self.spatial_distortion is not None: positions = ray_samples.frustums.get_positions() positions = self.spatial_distortion(positions) @@ -387,8 +386,8 @@ class TorchNerfactoField(Field): return density, base_mlp_out def get_outputs( - self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None - ) -> Dict[FieldHeadNames, TensorType]: + self, ray_samples: RaySamples, density_embedding: Optional[Tensor] = None + ) -> Dict[FieldHeadNames, Tensor]: outputs_shape = ray_samples.frustums.directions.shape[:-1] if ray_samples.camera_indices is None: diff --git a/nerfstudio/fields/nerfplayer_nerfacto_field.py b/nerfstudio/fields/nerfplayer_nerfacto_field.py index df40f1d5b..9ce2a02d2 100644 --- a/nerfstudio/fields/nerfplayer_nerfacto_field.py +++ b/nerfstudio/fields/nerfplayer_nerfacto_field.py @@ -21,8 +21,9 @@ from typing import Dict, Optional, Tuple import numpy as np import torch +from jaxtyping import Float +from torch import Tensor from torch.nn.parameter import Parameter -from torchtyping import TensorType from nerfstudio.cameras.rays import Frustums, RaySamples from nerfstudio.data.scene_box import SceneBox @@ -65,7 +66,7 @@ class TemporalHashMLPDensityField(Field): def __init__( self, - aabb: TensorType, + aabb: Tensor, temporal_dim: int = 64, num_layers: int = 2, hidden_dim: int = 64, @@ -104,7 +105,7 @@ class TemporalHashMLPDensityField(Field): ) # pylint: disable=arguments-differ - def density_fn(self, positions: TensorType["bs":..., 3], times: TensorType["bs", 1]) -> TensorType["bs":..., 1]: + def density_fn(self, positions: Float[Tensor, "*bs 3"], times: Float[Tensor, "bs 1"]) -> Float[Tensor, "*bs 1"]: """Returns only the density. Used primarily with the density grid. Args: @@ -128,7 +129,7 @@ class TemporalHashMLPDensityField(Field): density, _ = self.get_density(ray_samples) return density - def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType, None]: + def get_density(self, ray_samples: RaySamples) -> Tuple[Tensor, None]: if self.spatial_distortion is not None: positions = self.spatial_distortion(ray_samples.frustums.get_positions()) positions = (positions + 2.0) / 4.0 @@ -145,7 +146,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) -> dict: + def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[Tensor] = None) -> dict: return {} @@ -177,7 +178,7 @@ class NerfplayerNerfactoField(Field): def __init__( self, - aabb: TensorType, + aabb: Tensor, num_images: int, num_layers: int = 2, hidden_dim: int = 64, @@ -309,7 +310,7 @@ class NerfplayerNerfactoField(Field): }, ) - def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType, TensorType]: + def get_density(self, ray_samples: RaySamples) -> Tuple[Tensor, Tensor]: """Computes and returns the densities.""" if self.spatial_distortion is not None: positions = ray_samples.frustums.get_positions() @@ -331,8 +332,8 @@ class NerfplayerNerfactoField(Field): return density, base_mlp_out def get_outputs( - self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None - ) -> Dict[FieldHeadNames, TensorType]: + self, ray_samples: RaySamples, density_embedding: Optional[Tensor] = None + ) -> Dict[FieldHeadNames, Tensor]: 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 a66576acd..19b0ba3af 100644 --- a/nerfstudio/fields/nerfplayer_ngp_field.py +++ b/nerfstudio/fields/nerfplayer_ngp_field.py @@ -20,8 +20,9 @@ NeRFPlayer (https://arxiv.org/abs/2210.15947) field implementations with Instant from typing import Dict, Optional, Tuple import torch +from jaxtyping import Float +from torch import Tensor from torch.nn.parameter import Parameter -from torchtyping import TensorType from nerfstudio.cameras.rays import Frustums, RaySamples from nerfstudio.data.scene_box import SceneBox @@ -70,7 +71,7 @@ class NerfplayerNGPField(Field): def __init__( self, - aabb: TensorType, + aabb: Tensor, temporal_dim: int = 16, num_levels: int = 16, features_per_level: int = 2, @@ -146,7 +147,7 @@ class NerfplayerNGPField(Field): }, ) - def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType, TensorType]: + def get_density(self, ray_samples: RaySamples) -> Tuple[Tensor, Tensor]: if self.spatial_distortion is not None: positions = ray_samples.frustums.get_positions() positions = self.spatial_distortion(positions) @@ -172,8 +173,8 @@ class NerfplayerNGPField(Field): return density, base_mlp_out def get_outputs( - self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None - ) -> Dict[FieldHeadNames, TensorType]: + self, ray_samples: RaySamples, density_embedding: Optional[Tensor] = None + ) -> Dict[FieldHeadNames, Tensor]: assert density_embedding is not None directions = shift_directions_for_tcnn(ray_samples.frustums.directions) directions_flat = directions.view(-1, 3) @@ -201,7 +202,7 @@ class NerfplayerNGPField(Field): return {FieldHeadNames.RGB: rgb} # pylint: disable=arguments-differ - def density_fn(self, positions: TensorType["bs":..., 3], times: TensorType["bs":..., 1]) -> TensorType["bs":..., 1]: + def density_fn(self, positions: Float[Tensor, "*bs 3"], times: Float[Tensor, "*bs 1"]) -> Float[Tensor, "*bs 1"]: """Returns only the density. Used primarily with the density grid. Overwrite this function since density is time dependent now. @@ -222,7 +223,7 @@ class NerfplayerNGPField(Field): density, _ = self.get_density(ray_samples) return density - def get_opacity(self, positions: TensorType["bs":..., 3], step_size, time_intervals=10) -> TensorType["bs":..., 1]: + def get_opacity(self, positions: Float[Tensor, "*bs 3"], step_size, time_intervals=10) -> Float[Tensor, "*bs 1"]: """Returns the opacity for a position and time. Used primarily by the occupancy grid. This will return the maximum opacity of the points in the space in a dynamic sequence. diff --git a/nerfstudio/fields/nerfw_field.py b/nerfstudio/fields/nerfw_field.py index fbca8d9e9..be71f15bb 100644 --- a/nerfstudio/fields/nerfw_field.py +++ b/nerfstudio/fields/nerfw_field.py @@ -17,8 +17,8 @@ from typing import Dict, Optional, Tuple import torch -from torch import nn -from torchtyping import TensorType +from jaxtyping import Float +from torch import Tensor, nn from nerfstudio.cameras.rays import RaySamples from nerfstudio.field_components.embedding import Embedding @@ -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) -> Tuple[TensorType, TensorType]: + def get_density(self, ray_samples: RaySamples) -> Tuple[Tensor, Tensor]: """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()) @@ -118,16 +118,16 @@ class VanillaNerfWField(Field): return density, base_mlp_out def get_outputs( - self, ray_samples: RaySamples, density_embedding: Optional[TensorType[..., "embedding_size"]] = None - ) -> Dict[FieldHeadNames, TensorType]: + self, ray_samples: RaySamples, density_embedding: Optional[Float[Tensor, "*batch embedding_size"]] = None + ) -> Dict[FieldHeadNames, Tensor]: """Returns the outputs of the NeRF-W field. Args: ray_samples (RaySamples): Ray samples. - density_embedding (TensorType[..., "embedding_size"], optional): Density embedding. + density_embedding (Float[Tensor, "*batch embedding_size"], optional): Density embedding. Returns: - Dict[FieldHeadNames, TensorType]: Outputs of the NeRF-W field. + Dict[FieldHeadNames, Tensor]: Outputs of the NeRF-W field. """ outputs = {} encoded_dir = self.direction_encoding(ray_samples.frustums.directions) diff --git a/nerfstudio/fields/sdf_field.py b/nerfstudio/fields/sdf_field.py index 694d67389..775dac5ad 100644 --- a/nerfstudio/fields/sdf_field.py +++ b/nerfstudio/fields/sdf_field.py @@ -23,9 +23,9 @@ from typing import Dict, Literal, Optional, Type import numpy as np import torch import torch.nn.functional as F -from torch import nn +from jaxtyping import Float +from torch import Tensor, nn from torch.nn.parameter import Parameter -from torchtyping import TensorType from nerfstudio.cameras.rays import RaySamples from nerfstudio.field_components.embedding import Embedding @@ -52,11 +52,11 @@ class LearnedVariance(nn.Module): super().__init__() self.register_parameter("variance", nn.Parameter(init_val * torch.ones(1), requires_grad=True)) - def forward(self, x: TensorType[1]) -> TensorType[1]: + def forward(self, x: Float[Tensor, "1"]) -> Float[Tensor, "1"]: """Returns current variance value""" return torch.ones([len(x), 1], device=x.device) * torch.exp(self.variance * 10.0) - def get_variance(self) -> TensorType[1]: + def get_variance(self) -> Float[Tensor, "1"]: """return current variance value""" return torch.exp(self.variance * 10.0).clip(1e-6, 1e6) @@ -128,7 +128,7 @@ class SDFField(Field): def __init__( self, config: SDFFieldConfig, - aabb: TensorType[2, 3], + aabb: Float[Tensor, "2 3"], num_images: int, use_average_appearance_embedding: bool = False, spatial_distortion: Optional[SpatialDistortion] = None, @@ -252,7 +252,7 @@ class SDFField(Field): """Set the anneal value for the proposal network.""" self._cos_anneal_ratio = anneal - def forward_geonetwork(self, inputs: TensorType[..., 3]) -> TensorType[..., "geo-features+1"]: + def forward_geonetwork(self, inputs: Float[Tensor, "*batch 3"]) -> Float[Tensor, "*batch geo_features+1"]: """forward the geonetwork""" if self.use_grid_feature: positions = self.spatial_distortion(inputs) @@ -281,7 +281,8 @@ class SDFField(Field): outputs = self.softplus(outputs) return outputs - def get_sdf(self, ray_samples: RaySamples) -> TensorType["num_samples", -1, 1]: + # TODO: fix ... in shape annotations. + def get_sdf(self, ray_samples: RaySamples) -> Float[Tensor, "num_samples ... 1"]: """predict the sdf value for ray samples""" positions = ray_samples.frustums.get_start_positions() positions_flat = positions.view(-1, 3) @@ -292,9 +293,9 @@ class SDFField(Field): def get_alpha( self, ray_samples: RaySamples, - sdf: Optional[TensorType["num_samples", -1, 1]] = None, - gradients: Optional[TensorType["num_samples", -1, 1]] = None, - ) -> TensorType["num_samples", -1, 1]: + sdf: Optional[Float[Tensor, "num_samples ... 1"]] = None, + gradients: Optional[Float[Tensor, "num_samples ... 1"]] = None, + ) -> Float[Tensor, "num_samples ... 1"]: """compute alpha from sdf as in NeuS""" if sdf is None or gradients is None: inputs = ray_samples.frustums.get_start_positions() @@ -344,12 +345,12 @@ class SDFField(Field): def get_colors( self, - points: TensorType[..., 3], - directions: TensorType[..., 3], - normals: TensorType[..., 3], - geo_features: TensorType[..., "geo-feat-dim"], - camera_indices: TensorType, - ) -> TensorType[..., 3]: + points: Float[Tensor, "*batch 3"], + directions: Float[Tensor, "*batch 3"], + normals: Float[Tensor, "*batch 3"], + geo_features: Float[Tensor, "*batch geo_feat_dim"], + camera_indices: Tensor, + ) -> Float[Tensor, "*batch 3"]: """compute colors""" d = self.direction_encoding(directions) @@ -395,9 +396,9 @@ class SDFField(Field): def get_outputs( self, ray_samples: RaySamples, - density_embedding: Optional[TensorType] = None, + density_embedding: Optional[Tensor] = None, return_alphas: bool = False, - ) -> Dict[FieldHeadNames, TensorType]: + ) -> Dict[FieldHeadNames, Tensor]: """compute output of ray samples""" if ray_samples.camera_indices is None: raise AttributeError("Camera indices are not provided.") @@ -445,7 +446,7 @@ class SDFField(Field): def forward( self, ray_samples: RaySamples, compute_normals: bool = False, return_alphas: bool = False - ) -> Dict[FieldHeadNames, TensorType]: + ) -> Dict[FieldHeadNames, Tensor]: """Evaluates the field at points along the ray. Args: diff --git a/nerfstudio/fields/semantic_nerf_field.py b/nerfstudio/fields/semantic_nerf_field.py index 6d4163dc2..5681dd862 100644 --- a/nerfstudio/fields/semantic_nerf_field.py +++ b/nerfstudio/fields/semantic_nerf_field.py @@ -18,8 +18,7 @@ Semantic NeRF field implementation. from typing import Dict, Optional, Tuple import torch -from torch import nn -from torchtyping import TensorType +from torch import Tensor, nn from nerfstudio.cameras.rays import RaySamples from nerfstudio.field_components.encodings import Encoding, Identity @@ -88,15 +87,15 @@ class SemanticNerfField(Field): in_dim=self.mlp_semantic.get_out_dim(), num_classes=self.num_semantic_classes ) - def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType, TensorType]: + def get_density(self, ray_samples: RaySamples) -> Tuple[Tensor, Tensor]: 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) return density, base_mlp_out def get_outputs( - self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None - ) -> Dict[FieldHeadNames, TensorType]: + self, ray_samples: RaySamples, density_embedding: Optional[Tensor] = None + ) -> Dict[FieldHeadNames, Tensor]: encoded_dir = self.direction_encoding(ray_samples.frustums.directions) mlp_out = self.mlp_head(torch.cat([encoded_dir, density_embedding], dim=-1)) # type: ignore outputs = {} diff --git a/nerfstudio/fields/tensorf_field.py b/nerfstudio/fields/tensorf_field.py index f5077dad9..27d5f3cd4 100644 --- a/nerfstudio/fields/tensorf_field.py +++ b/nerfstudio/fields/tensorf_field.py @@ -18,9 +18,8 @@ from typing import Dict, Optional import torch -from torch import nn +from torch import Tensor, nn from torch.nn.parameter import Parameter -from torchtyping import TensorType from nerfstudio.cameras.rays import RaySamples from nerfstudio.data.scene_box import SceneBox @@ -35,7 +34,7 @@ class TensoRFField(Field): def __init__( self, - aabb: TensorType, + aabb: Tensor, # the aabb bounding box of the dataset feature_encoding: Encoding = Identity(in_dim=3), # the encoding method used for appearance encoding outputs @@ -83,7 +82,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) -> TensorType: + def get_density(self, ray_samples: RaySamples) -> Tensor: positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) positions = positions * 2 - 1 density = self.density_encoding(positions) @@ -92,7 +91,7 @@ class TensoRFField(Field): density_enc = relu(density_enc) return density_enc - def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None) -> TensorType: + def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[Tensor] = None) -> Tensor: d = ray_samples.frustums.directions positions = SceneBox.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) positions = positions * 2 - 1 @@ -116,9 +115,9 @@ class TensoRFField(Field): self, ray_samples: RaySamples, compute_normals: bool = False, - mask: Optional[TensorType] = None, - bg_color: Optional[TensorType] = None, - ) -> Dict[FieldHeadNames, TensorType]: + mask: Optional[Tensor] = None, + bg_color: Optional[Tensor] = None, + ) -> Dict[FieldHeadNames, Tensor]: 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 9e3473521..96ef0dc56 100644 --- a/nerfstudio/fields/vanilla_nerf_field.py +++ b/nerfstudio/fields/vanilla_nerf_field.py @@ -18,8 +18,7 @@ from typing import Dict, Optional, Tuple import torch -from torch import nn -from torchtyping import TensorType +from torch import Tensor, nn from nerfstudio.cameras.rays import RaySamples from nerfstudio.field_components.encodings import Encoding, Identity @@ -88,7 +87,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) -> Tuple[TensorType, TensorType]: + def get_density(self, ray_samples: RaySamples) -> Tuple[Tensor, Tensor]: if self.use_integrated_encoding: gaussian_samples = ray_samples.frustums.get_gaussian_blob() if self.spatial_distortion is not None: @@ -104,8 +103,8 @@ class NeRFField(Field): return density, base_mlp_out def get_outputs( - self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None - ) -> Dict[FieldHeadNames, TensorType]: + self, ray_samples: RaySamples, density_embedding: Optional[Tensor] = None + ) -> Dict[FieldHeadNames, Tensor]: outputs = {} for field_head in self.field_heads: encoded_dir = self.direction_encoding(ray_samples.frustums.directions) diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index 4c3f07a4a..178c5ae46 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/nerfstudio/generative/stable_diffusion.py @@ -27,10 +27,10 @@ import numpy as np import torch import torch.nn.functional as F import tyro -from torch import nn +from jaxtyping import Float +from torch import Tensor, nn from torch.cuda.amp import custom_bwd, custom_fwd from torch.cuda.amp.grad_scaler import GradScaler -from torchtyping import TensorType from nerfstudio.utils.rich_utils import CONSOLE @@ -148,7 +148,7 @@ class StableDiffusion(nn.Module): def get_text_embeds( self, prompt: Union[str, List[str]], negative_prompt: Union[str, List[str]] - ) -> TensorType[2, "max_length", "embed_dim"]: + ) -> Float[Tensor, "2 max_length embed_dim"]: """Get text embeddings for prompt and negative prompt Args: prompt: Prompt text @@ -184,8 +184,8 @@ class StableDiffusion(nn.Module): def sds_loss( self, - text_embeddings: TensorType["N", "max_length", "embed_dim"], - image: TensorType["BS", 3, "H", "W"], + text_embeddings: Float[Tensor, "N max_length embed_dim"], + image: Float[Tensor, "BS 3 H W"], guidance_scale: float = 100.0, grad_scaler: Optional[GradScaler] = None, ) -> torch.Tensor: @@ -229,13 +229,13 @@ class StableDiffusion(nn.Module): def produce_latents( self, - text_embeddings: TensorType["N", "max_length", "embed_dim"], + text_embeddings: Float[Tensor, "N max_length embed_dim"], height: int = IMG_DIM, width: int = IMG_DIM, num_inference_steps: int = 50, guidance_scale: float = 7.5, - latents: Optional[TensorType["BS", 4, "H", "W"]] = None, - ) -> TensorType["BS", 4, "H", "W"]: + latents: Optional[Float[Tensor, "BS 4 H W"]] = None, + ) -> Float[Tensor, "BS 4 H W"]: """Produce latents for a given text embedding Args: text_embeddings: Text embeddings @@ -275,7 +275,7 @@ class StableDiffusion(nn.Module): latents = self.scheduler.step(noise_pred, t, latents)["prev_sample"] # type: ignore return latents - def latents_to_img(self, latents: TensorType["BS", 4, "H", "W"]) -> TensorType["BS", 3, "H", "W"]: + def latents_to_img(self, latents: Float[Tensor, "BS 4 H W"]) -> Float[Tensor, "BS 3 H W"]: """Convert latents to images Args: latents: Latents to convert @@ -292,7 +292,7 @@ class StableDiffusion(nn.Module): return imgs - def imgs_to_latent(self, imgs: TensorType["BS", 3, "H", "W"]) -> TensorType["BS", 4, "H", "W"]: + def imgs_to_latent(self, imgs: Float[Tensor, "BS 3 H W"]) -> Float[Tensor, "BS 4 H W"]: """Convert images to latents Args: imgs: Images to convert diff --git a/nerfstudio/model_components/losses.py b/nerfstudio/model_components/losses.py index d7d71d15d..a0cc6f22d 100644 --- a/nerfstudio/model_components/losses.py +++ b/nerfstudio/model_components/losses.py @@ -19,8 +19,8 @@ from enum import Enum from typing import Dict, Literal import torch -from torch import nn -from torchtyping import TensorType +from jaxtyping import Bool, Float +from torch import Tensor, nn from nerfstudio.cameras.rays import RaySamples from nerfstudio.utils.math import masked_reduction, normalized_depth_scale_and_shift @@ -44,12 +44,12 @@ class DepthLossType(Enum): def outer( - t0_starts: TensorType[..., "num_samples_0"], - t0_ends: TensorType[..., "num_samples_0"], - t1_starts: TensorType[..., "num_samples_1"], - t1_ends: TensorType[..., "num_samples_1"], - y1: TensorType[..., "num_samples_1"], -) -> TensorType[..., "num_samples_0"]: + t0_starts: Float[Tensor, "*batch num_samples_0"], + t0_ends: Float[Tensor, "*batch num_samples_0"], + t1_starts: Float[Tensor, "*batch num_samples_1"], + t1_ends: Float[Tensor, "*batch num_samples_1"], + y1: Float[Tensor, "*batch num_samples_1"], +) -> Float[Tensor, "*batch num_samples_0"]: """Faster version of https://github.com/kakaobrain/NeRF-Factory/blob/f61bb8744a5cb4820a4d968fb3bfbed777550f4a/src/model/mipnerf360/helper.py#L117 @@ -76,10 +76,10 @@ def outer( def lossfun_outer( - t: TensorType[..., "num_samples+1"], - w: TensorType[..., "num_samples"], - t_env: TensorType[..., "num_samples+1"], - w_env: TensorType[..., "num_samples"], + t: Float[Tensor, "*batch num_samples+1"], + w: Float[Tensor, "*batch num_samples"], + t_env: Float[Tensor, "*batch num_samples+1"], + w_env: Float[Tensor, "*batch num_samples"], ): """ https://github.com/kakaobrain/NeRF-Factory/blob/f61bb8744a5cb4820a4d968fb3bfbed777550f4a/src/model/mipnerf360/helper.py#L136 @@ -145,9 +145,9 @@ def distortion_loss(weights_list, ray_samples_list): def nerfstudio_distortion_loss( ray_samples: RaySamples, - densities: TensorType["bs":..., "num_samples", 1] = None, - weights: TensorType["bs":..., "num_samples", 1] = None, -) -> TensorType["bs":..., 1]: + densities: Float[Tensor, "*bs num_samples 1"] = None, + weights: Float[Tensor, "*bs num_samples 1"] = None, +) -> Float[Tensor, "*bs 1"]: """Ray based distortion loss proposed in MipNeRF-360. Returns distortion Loss. .. math:: @@ -186,9 +186,9 @@ def nerfstudio_distortion_loss( def orientation_loss( - weights: TensorType["bs":..., "num_samples", 1], - normals: TensorType["bs":..., "num_samples", 3], - viewdirs: TensorType["bs":..., 3], + weights: Float[Tensor, "*bs num_samples 1"], + normals: Float[Tensor, "*bs num_samples 3"], + viewdirs: Float[Tensor, "*bs 3"], ): """Orientation loss proposed in Ref-NeRF. Loss that encourages that all visible normals are facing towards the camera. @@ -201,21 +201,21 @@ def orientation_loss( def pred_normal_loss( - weights: TensorType["bs":..., "num_samples", 1], - normals: TensorType["bs":..., "num_samples", 3], - pred_normals: TensorType["bs":..., "num_samples", 3], + weights: Float[Tensor, "*bs num_samples 1"], + normals: Float[Tensor, "*bs num_samples 3"], + pred_normals: Float[Tensor, "*bs num_samples 3"], ): """Loss between normals calculated from density and normals from prediction network.""" return (weights[..., 0] * (1.0 - torch.sum(normals * pred_normals, dim=-1))).sum(dim=-1) def ds_nerf_depth_loss( - weights: TensorType[..., "num_samples", 1], - termination_depth: TensorType[..., 1], - steps: TensorType[..., "num_samples", 1], - lengths: TensorType[..., "num_samples", 1], - sigma: TensorType[0], -) -> TensorType[..., 1]: + weights: Float[Tensor, "*batch num_samples 1"], + termination_depth: Float[Tensor, "*batch 1"], + steps: Float[Tensor, "*batch num_samples 1"], + lengths: Float[Tensor, "*batch num_samples 1"], + sigma: Float[Tensor, "0"], +) -> Float[Tensor, "*batch 1"]: """Depth loss from Depth-supervised NeRF (Deng et al., 2022). Args: @@ -235,12 +235,12 @@ def ds_nerf_depth_loss( def urban_radiance_field_depth_loss( - weights: TensorType[..., "num_samples", 1], - termination_depth: TensorType[..., 1], - predicted_depth: TensorType[..., 1], - steps: TensorType[..., "num_samples", 1], - sigma: TensorType[0], -) -> TensorType[..., 1]: + weights: Float[Tensor, "*batch num_samples 1"], + termination_depth: Float[Tensor, "*batch 1"], + predicted_depth: Float[Tensor, "*batch 1"], + steps: Float[Tensor, "*batch num_samples 1"], + sigma: Float[Tensor, "0"], +) -> Float[Tensor, "*batch 1"]: """Lidar losses from Urban Radiance Fields (Rematas et al., 2022). Args: @@ -274,15 +274,15 @@ def urban_radiance_field_depth_loss( def depth_loss( - weights: TensorType[..., "num_samples", 1], + weights: Float[Tensor, "*batch num_samples 1"], ray_samples: RaySamples, - termination_depth: TensorType[..., 1], - predicted_depth: TensorType[..., 1], - sigma: TensorType[0], - directions_norm: TensorType[..., 1], + termination_depth: Float[Tensor, "*batch 1"], + predicted_depth: Float[Tensor, "*batch 1"], + sigma: Float[Tensor, "0"], + directions_norm: Float[Tensor, "*batch 1"], is_euclidean: bool, depth_loss_type: DepthLossType, -) -> TensorType[0]: +) -> Float[Tensor, "0"]: """Implementation of depth losses. Args: @@ -313,8 +313,8 @@ def depth_loss( def monosdf_normal_loss( - normal_pred: TensorType["num_samples", 3], normal_gt: TensorType["num_samples", 3] -) -> TensorType[0]: + normal_pred: Float[Tensor, "num_samples 3"], normal_gt: Float[Tensor, "num_samples 3"] +) -> Float[Tensor, "0"]: """ Normal consistency loss proposed in monosdf - https://niujinshuchong.github.io/monosdf/ Enforces consistency between the volume rendered normal and the predicted monocular normal. @@ -343,8 +343,11 @@ class MiDaSMSELoss(nn.Module): self.mse_loss = MSELoss(reduction="none") def forward( - self, prediction: TensorType[1, 32, "mult"], target: TensorType[1, 32, "mult"], mask: TensorType[1, 32, "mult"] - ) -> TensorType[0]: + self, + prediction: Float[Tensor, "1 32 mult"], + target: Float[Tensor, "1 32 mult"], + mask: Bool[Tensor, "1 32 mult"], + ) -> Float[Tensor, "0"]: """ Args: prediction: predicted depth map @@ -380,8 +383,11 @@ class GradientLoss(nn.Module): self.__scales = scales def forward( - self, prediction: TensorType[1, 32, "mult"], target: TensorType[1, 32, "mult"], mask: TensorType[1, 32, "mult"] - ) -> TensorType[0]: + self, + prediction: Float[Tensor, "1 32 mult"], + target: Float[Tensor, "1 32 mult"], + mask: Bool[Tensor, "1 32 mult"], + ) -> Float[Tensor, "0"]: """ Args: prediction: predicted depth map @@ -405,8 +411,11 @@ class GradientLoss(nn.Module): return total def gradient_loss( - self, prediction: TensorType[1, 32, "mult"], target: TensorType[1, 32, "mult"], mask: TensorType[1, 32, "mult"] - ) -> TensorType[0]: + self, + prediction: Float[Tensor, "1 32 mult"], + target: Float[Tensor, "1 32 mult"], + mask: Bool[Tensor, "1 32 mult"], + ) -> Float[Tensor, "0"]: """ multiscale, scale-invariant gradient matching term to the disparity space. This term biases discontinuities to be sharp and to coincide with discontinuities in the ground truth @@ -458,8 +467,11 @@ class ScaleAndShiftInvariantLoss(nn.Module): self.__prediction_ssi = None def forward( - self, prediction: TensorType[1, 32, "mult"], target: TensorType[1, 32, "mult"], mask: TensorType[1, 32, "mult"] - ) -> TensorType[0]: + self, + prediction: Float[Tensor, "1 32 mult"], + target: Float[Tensor, "1 32 mult"], + mask: Bool[Tensor, "1 32 mult"], + ) -> Float[Tensor, "0"]: """ Args: prediction: predicted depth map (unnormalized) @@ -487,7 +499,7 @@ class ScaleAndShiftInvariantLoss(nn.Module): prediction_ssi = property(__get_prediction_ssi) -def tv_loss(grids: TensorType["grids", "feature_dim", "row", "column"]) -> TensorType[()]: +def tv_loss(grids: Float[Tensor, "grids feature_dim row column"]) -> Float[Tensor, ""]: """ https://github.com/apchenstu/TensoRF/blob/4ec894dc1341a2201fe13ae428631b58458f105d/utils.py#L139 diff --git a/nerfstudio/model_components/ray_generators.py b/nerfstudio/model_components/ray_generators.py index d540bc405..89e64cf04 100644 --- a/nerfstudio/model_components/ray_generators.py +++ b/nerfstudio/model_components/ray_generators.py @@ -15,8 +15,8 @@ """ Ray generator. """ -from torch import nn -from torchtyping import TensorType +from jaxtyping import Int +from torch import Tensor, nn from nerfstudio.cameras.camera_optimizers import CameraOptimizer from nerfstudio.cameras.cameras import Cameras @@ -38,7 +38,7 @@ class RayGenerator(nn.Module): self.pose_optimizer = pose_optimizer self.register_buffer("image_coords", cameras.get_image_coords(), persistent=False) - def forward(self, ray_indices: TensorType["num_rays", 3]) -> RayBundle: + def forward(self, ray_indices: Int[Tensor, "num_rays 3"]) -> RayBundle: """Index into the cameras to generate the rays. Args: diff --git a/nerfstudio/model_components/ray_samplers.py b/nerfstudio/model_components/ray_samplers.py index 2874707f0..388ddae5a 100644 --- a/nerfstudio/model_components/ray_samplers.py +++ b/nerfstudio/model_components/ray_samplers.py @@ -20,9 +20,9 @@ from abc import abstractmethod from typing import Callable, List, Optional, Tuple, Union import torch +from jaxtyping import Float from nerfacc import OccGridEstimator -from torch import nn -from torchtyping import TensorType +from torch import Tensor, nn from nerfstudio.cameras.rays import Frustums, RayBundle, RaySamples @@ -111,7 +111,10 @@ class SpacedSampler(Sampler): bins = bin_lower + (bin_upper - bin_lower) * t_rand s_near, s_far = (self.spacing_fn(x) for x in (ray_bundle.nears, ray_bundle.fars)) - spacing_to_euclidean_fn = lambda x: self.spacing_fn_inv(x * s_far + (1 - x) * s_near) + + def spacing_to_euclidean_fn(x): + return self.spacing_fn_inv(x * s_far + (1 - x) * s_near) + euclidean_bins = spacing_to_euclidean_fn(bins) # [num_rays, num_samples+1] ray_samples = ray_bundle.get_ray_samples( @@ -274,7 +277,7 @@ class PDFSampler(Sampler): self, ray_bundle: Optional[RayBundle] = None, ray_samples: Optional[RaySamples] = None, - weights: TensorType[..., "num_samples", 1] = None, + weights: Float[Tensor, "*batch num_samples 1"] = None, num_samples: Optional[int] = None, eps: float = 1e-5, ) -> RaySamples: @@ -382,7 +385,7 @@ class VolumetricSampler(Sampler): def __init__( self, occupancy_grid: OccGridEstimator, - density_fn: Optional[Callable[[TensorType[..., 3]], TensorType[..., 1]]] = None, + density_fn: Optional[Callable[[Float[Tensor, "*batch 3"]], Float[Tensor, "*batch 1"]]] = None, ) -> None: super().__init__() assert occupancy_grid is not None @@ -429,7 +432,7 @@ class VolumetricSampler(Sampler): far_plane: Optional[float] = None, alpha_thre: float = 0.01, cone_angle: float = 0.0, - ) -> Tuple[RaySamples, TensorType["total_samples",]]: + ) -> Tuple[RaySamples, Float[Tensor, "total_samples "]]: """Generate ray samples in a bounding box. Args: @@ -683,8 +686,8 @@ class NeuSSampler(Sampler): @staticmethod def rendering_sdf_with_fixed_inv_s( - ray_samples: RaySamples, sdf: TensorType["num_samples", -1], inv_s: int - ) -> TensorType["num_samples", -1]: + ray_samples: RaySamples, sdf: Float[Tensor, "num_samples 1"], inv_s: int + ) -> Float[Tensor, "num_samples 1"]: """ rendering given a fixed inv_s as NeuS diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index 1e0ee26e9..12cc45a57 100644 --- a/nerfstudio/model_components/renderers.py +++ b/nerfstudio/model_components/renderers.py @@ -32,18 +32,18 @@ from typing import Generator, Literal, Optional, Union import nerfacc import torch -from torch import nn -from torchtyping import TensorType +from jaxtyping import Float, Int +from torch import Tensor, nn from nerfstudio.cameras.rays import RaySamples from nerfstudio.utils import colors from nerfstudio.utils.math import components_from_spherical_harmonics, safe_normalize -BACKGROUND_COLOR_OVERRIDE: Optional[TensorType[3]] = None +BACKGROUND_COLOR_OVERRIDE: Optional[Float[Tensor, "3"]] = None @contextlib.contextmanager -def background_color_override_context(mode: TensorType[3]) -> Generator[None, None, None]: +def background_color_override_context(mode: Float[Tensor, "3"]) -> Generator[None, None, None]: """Context manager for setting background mode.""" global BACKGROUND_COLOR_OVERRIDE # pylint: disable=global-statement old_background_color = BACKGROUND_COLOR_OVERRIDE @@ -61,19 +61,21 @@ class RGBRenderer(nn.Module): background_color: Background color as RGB. Uses random colors if None. """ - def __init__(self, background_color: Union[Literal["random", "last_sample"], TensorType[3]] = "random") -> None: + def __init__( + self, background_color: Union[Literal["random", "last_sample"], Float[Tensor, "3"]] = "random" + ) -> None: super().__init__() self.background_color = background_color @classmethod def combine_rgb( cls, - rgb: TensorType["bs":..., "num_samples", 3], - weights: TensorType["bs":..., "num_samples", 1], - background_color: Union[Literal["random", "white", "black", "last_sample"], TensorType[3]] = "random", - ray_indices: Optional[TensorType["num_samples"]] = None, + rgb: Float[Tensor, "*bs num_samples 3"], + weights: Float[Tensor, "*bs num_samples 1"], + background_color: Union[Literal["random", "white", "black", "last_sample"], Float[Tensor, "3"]] = "random", + ray_indices: Optional[Int[Tensor, "num_samples"]] = None, num_rays: Optional[int] = None, - ) -> TensorType["bs":..., 3]: + ) -> Float[Tensor, "*bs 3"]: """Composite samples along ray and render color image Args: @@ -116,11 +118,11 @@ class RGBRenderer(nn.Module): def forward( self, - rgb: TensorType["bs":..., "num_samples", 3], - weights: TensorType["bs":..., "num_samples", 1], - ray_indices: Optional[TensorType["num_samples"]] = None, + rgb: Float[Tensor, "*bs num_samples 3"], + weights: Float[Tensor, "*bs num_samples 1"], + ray_indices: Optional[Int[Tensor, "num_samples"]] = None, num_rays: Optional[int] = None, - ) -> TensorType["bs":..., 3]: + ) -> Float[Tensor, "*bs 3"]: """Composite samples along ray and render color image Args: @@ -153,7 +155,7 @@ class SHRenderer(nn.Module): def __init__( self, - background_color: Union[Literal["random", "last_sample"], TensorType[3]] = "random", + background_color: Union[Literal["random", "last_sample"], Float[Tensor, "3"]] = "random", activation: Optional[nn.Module] = nn.Sigmoid(), ) -> None: super().__init__() @@ -162,10 +164,10 @@ class SHRenderer(nn.Module): def forward( self, - sh: TensorType[..., "num_samples", "coeffs"], - directions: TensorType[..., "num_samples", 3], - weights: TensorType[..., "num_samples", 1], - ) -> TensorType[..., 3]: + sh: Float[Tensor, "*batch num_samples coeffs"], + directions: Float[Tensor, "*batch num_samples 3"], + weights: Float[Tensor, "*batch num_samples 1"], + ) -> Float[Tensor, "*batch 3"]: """Composite samples along ray and render color image Args: @@ -203,10 +205,10 @@ class AccumulationRenderer(nn.Module): @classmethod def forward( cls, - weights: TensorType["bs":..., "num_samples", 1], - ray_indices: Optional[TensorType["num_samples"]] = None, + weights: Float[Tensor, "*bs num_samples 1"], + ray_indices: Optional[Int[Tensor, "num_samples"]] = None, num_rays: Optional[int] = None, - ) -> TensorType["bs":..., 1]: + ) -> Float[Tensor, "*bs 1"]: """Composite samples along ray and calculate accumulation. Args: @@ -245,11 +247,11 @@ class DepthRenderer(nn.Module): def forward( self, - weights: TensorType[..., "num_samples", 1], + weights: Float[Tensor, "*batch num_samples 1"], ray_samples: RaySamples, - ray_indices: Optional[TensorType["num_samples"]] = None, + ray_indices: Optional[Int[Tensor, "num_samples"]] = None, num_rays: Optional[int] = None, - ) -> TensorType[..., 1]: + ) -> Float[Tensor, "*batch 1"]: """Composite samples along ray and calculate depths. Args: @@ -301,8 +303,8 @@ class UncertaintyRenderer(nn.Module): @classmethod def forward( - cls, betas: TensorType["bs":..., "num_samples", 1], weights: TensorType["bs":..., "num_samples", 1] - ) -> TensorType["bs":..., 1]: + cls, betas: Float[Tensor, "*bs num_samples 1"], weights: Float[Tensor, "*bs num_samples 1"] + ) -> Float[Tensor, "*bs 1"]: """Calculate uncertainty along the ray. Args: @@ -322,9 +324,9 @@ class SemanticRenderer(nn.Module): @classmethod def forward( cls, - semantics: TensorType["bs":..., "num_samples", "num_classes"], - weights: TensorType["bs":..., "num_samples", 1], - ) -> TensorType["bs":..., "num_classes"]: + semantics: Float[Tensor, "*bs num_samples num_classes"], + weights: Float[Tensor, "*bs num_samples 1"], + ) -> Float[Tensor, "*bs num_classes"]: """Calculate semantics along the ray.""" sem = torch.sum(weights * semantics, dim=-2) return sem @@ -336,10 +338,10 @@ class NormalsRenderer(nn.Module): @classmethod def forward( cls, - normals: TensorType["bs":..., "num_samples", 3], - weights: TensorType["bs":..., "num_samples", 1], + normals: Float[Tensor, "*bs num_samples 3"], + weights: Float[Tensor, "*bs num_samples 1"], normalize: bool = True, - ) -> TensorType["bs":..., 3]: + ) -> Float[Tensor, "*bs 3"]: """Calculate normals along the ray. Args: diff --git a/nerfstudio/model_components/scene_colliders.py b/nerfstudio/model_components/scene_colliders.py index a55ebd911..97f53d17b 100644 --- a/nerfstudio/model_components/scene_colliders.py +++ b/nerfstudio/model_components/scene_colliders.py @@ -19,8 +19,8 @@ Scene Colliders from __future__ import annotations import torch -from torch import nn -from torchtyping import TensorType +from jaxtyping import Float +from torch import Tensor, nn from nerfstudio.cameras.rays import RayBundle from nerfstudio.data.scene_box import SceneBox @@ -57,7 +57,7 @@ class AABBBoxCollider(SceneCollider): self.near_plane = near_plane def _intersect_with_aabb( - self, rays_o: TensorType["num_rays", 3], rays_d: TensorType["num_rays", 3], aabb: TensorType[2, 3] + self, rays_o: Float[Tensor, "num_rays 3"], rays_d: Float[Tensor, "num_rays 3"], aabb: Float[Tensor, "2 3"] ): """Returns collection of valid rays within a specified near/far bounding box along with a mask specifying which rays are valid diff --git a/nerfstudio/model_components/shaders.py b/nerfstudio/model_components/shaders.py index 096503e1a..f6cc22721 100644 --- a/nerfstudio/model_components/shaders.py +++ b/nerfstudio/model_components/shaders.py @@ -15,8 +15,8 @@ """Shaders for rendering.""" from typing import Optional -from torch import nn -from torchtyping import TensorType +from jaxtyping import Float +from torch import Tensor, nn class LambertianShader(nn.Module): @@ -25,9 +25,9 @@ class LambertianShader(nn.Module): @classmethod def forward( cls, - rgb: TensorType["bs":..., 3], - normals: TensorType["bs":..., 3], - light_direction: TensorType["bs":..., 3], + rgb: Float[Tensor, "*bs 3"], + normals: Float[Tensor, "*bs 3"], + light_direction: Float[Tensor, "*bs 3"], shading_weight: float = 1.0, detach_normals=True, ): @@ -59,8 +59,8 @@ class NormalsShader(nn.Module): @classmethod def forward( cls, - normals: TensorType["bs":..., 3], - weights: Optional[TensorType["bs":..., 1]] = None, + normals: Float[Tensor, "*bs 3"], + weights: Optional[Float[Tensor, "*bs 1"]] = None, ): """Applies a rainbow colormap to the normals. diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 3af9059d6..233234e89 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -187,11 +187,12 @@ class NerfactoModel(Model): self.density_fns.extend([network.density_fn for network in self.proposal_networks]) # Samplers - update_schedule = lambda step: np.clip( - np.interp(step, [0, self.config.proposal_warmup], [0, self.config.proposal_update_every]), - 1, - self.config.proposal_update_every, - ) + def update_schedule(step): + return np.clip( + np.interp(step, [0, self.config.proposal_warmup], [0, self.config.proposal_update_every]), + 1, + self.config.proposal_update_every, + ) # Change proposal network initial sampler if uniform initial_sampler = None # None is for piecewise as default (see ProposalNetworkSampler) @@ -244,7 +245,10 @@ class NerfactoModel(Model): def set_anneal(step): # https://arxiv.org/pdf/2111.12077.pdf eq. 18 train_frac = np.clip(step / N, 0, 1) - bias = lambda x, b: (b * x) / ((b - 1) * x + 1) + + def bias(x, b): + return b * x / ((b - 1) * x + 1) + anneal = bias(train_frac, self.config.proposal_weights_anneal_slope) self.proposal_sampler.set_anneal(anneal) diff --git a/nerfstudio/models/nerfplayer_nerfacto.py b/nerfstudio/models/nerfplayer_nerfacto.py index 5f837ca4e..9facfe580 100644 --- a/nerfstudio/models/nerfplayer_nerfacto.py +++ b/nerfstudio/models/nerfplayer_nerfacto.py @@ -140,11 +140,13 @@ class NerfplayerNerfactoModel(NerfactoModel): self.density_fns.extend([network.density_fn for network in self.proposal_networks]) # Samplers - update_schedule = lambda step: np.clip( - np.interp(step, [0, self.config.proposal_warmup], [0, self.config.proposal_update_every]), - 1, - self.config.proposal_update_every, - ) + def update_schedule(step): + return np.clip( + np.interp(step, [0, self.config.proposal_warmup], [0, self.config.proposal_update_every]), + 1, + self.config.proposal_update_every, + ) + self.proposal_sampler = ProposalNetworkSampler( num_nerf_samples_per_ray=self.config.num_nerf_samples_per_ray, num_proposal_samples_per_ray=self.config.num_proposal_samples_per_ray, @@ -246,7 +248,10 @@ class NerfplayerNerfactoModel(NerfactoModel): if "depth_image" in batch.keys() and self.config.depth_weight > 0: mask = (batch["depth_image"] != 0).view([-1]) loss_dict["depth_loss"] = 0 - l = lambda x: self.config.depth_weight * (x - batch["depth_image"][mask]).pow(2).mean() + + def l(x): + return self.config.depth_weight * (x - batch["depth_image"][mask]).pow(2).mean() + loss_dict["depth_loss"] = l(outputs["depth"][mask]) for i in range(self.config.num_proposal_iterations): loss_dict["depth_loss"] += l(outputs[f"prop_depth_{i}"][mask]) diff --git a/nerfstudio/models/neus_facto.py b/nerfstudio/models/neus_facto.py index c37faa3e8..1f7104620 100644 --- a/nerfstudio/models/neus_facto.py +++ b/nerfstudio/models/neus_facto.py @@ -119,7 +119,8 @@ class NeuSFactoModel(NeuSModel): self.density_fns.extend([network.density_fn for network in self.proposal_networks]) # update proposal network every iterations - update_schedule = lambda step: -1 + def update_schedule(_): + return -1 initial_sampler = UniformSampler(single_jitter=self.config.use_single_jitter) self.proposal_sampler = ProposalNetworkSampler( @@ -149,7 +150,10 @@ class NeuSFactoModel(NeuSModel): def set_anneal(step: int): # https://arxiv.org/pdf/2111.12077.pdf eq. 18 train_frac = np.clip(step / N, 0, 1) - bias = lambda x, b: (b * x) / ((b - 1) * x + 1) + + def bias(x, b): + return b * x / ((b - 1) * x + 1) + anneal = bias(train_frac, self.config.proposal_weights_anneal_slope) self.proposal_sampler.set_anneal(anneal) diff --git a/nerfstudio/models/semantic_nerfw.py b/nerfstudio/models/semantic_nerfw.py index 35623cc99..a006ca3b2 100644 --- a/nerfstudio/models/semantic_nerfw.py +++ b/nerfstudio/models/semantic_nerfw.py @@ -160,7 +160,10 @@ class SemanticNerfWModel(Model): def set_anneal(step): # https://arxiv.org/pdf/2111.12077.pdf eq. 18 train_frac = np.clip(step / N, 0, 1) - bias = lambda x, b: (b * x) / ((b - 1) * x + 1) + + def bias(x, b): + return b * x / ((b - 1) * x + 1) + anneal = bias(train_frac, self.config.proposal_weights_anneal_slope) self.proposal_sampler.set_anneal(anneal) diff --git a/nerfstudio/scripts/datasets/process_nuscenes_masks.py b/nerfstudio/scripts/datasets/process_nuscenes_masks.py index cb92a32c6..f88588595 100644 --- a/nerfstudio/scripts/datasets/process_nuscenes_masks.py +++ b/nerfstudio/scripts/datasets/process_nuscenes_masks.py @@ -21,7 +21,6 @@ import cv2 import numpy as np import tyro from nuscenes.nuscenes import NuScenes as NuScenesDatabase -from nuscenes.utils.data_classes import Box from nuscenes.utils.geometry_utils import BoxVisibility, view_points from tqdm import tqdm diff --git a/nerfstudio/scripts/process_data.py b/nerfstudio/scripts/process_data.py index ef17d5343..50a10bb5a 100644 --- a/nerfstudio/scripts/process_data.py +++ b/nerfstudio/scripts/process_data.py @@ -413,5 +413,7 @@ def entrypoint(): if __name__ == "__main__": entrypoint() -# For sphinx docs -get_parser_fn = lambda: tyro.extras.get_parser(Commands) # type: ignore + +def get_parser_fn(): + """Get the parser function for the sphinx docs.""" + return tyro.extras.get_parser(Commands) # type: ignore diff --git a/nerfstudio/scripts/render.py b/nerfstudio/scripts/render.py index c07f48c59..afe33ccd9 100644 --- a/nerfstudio/scripts/render.py +++ b/nerfstudio/scripts/render.py @@ -31,6 +31,7 @@ import mediapy as media import numpy as np import torch import tyro +from jaxtyping import Float from rich import box, style from rich.panel import Panel from rich.progress import ( @@ -41,7 +42,7 @@ from rich.progress import ( TimeRemainingColumn, ) from rich.table import Table -from torchtyping import TensorType +from torch import Tensor from typing_extensions import Annotated from nerfstudio.cameras.camera_paths import ( @@ -253,11 +254,11 @@ xmlns:GSpherical='http://ns.google.com/videos/1.0/spherical/'> class CropData: """Data for cropping an image.""" - background_color: TensorType[3] = torch.Tensor([0.0, 0.0, 0.0]) + background_color: Float[Tensor, "3"] = torch.Tensor([0.0, 0.0, 0.0]) """background color""" - center: TensorType[3] = torch.Tensor([0.0, 0.0, 0.0]) + center: Float[Tensor, "3"] = torch.Tensor([0.0, 0.0, 0.0]) """center of the crop""" - scale: TensorType[3] = torch.Tensor([2.0, 2.0, 2.0]) + scale: Float[Tensor, "3"] = torch.Tensor([2.0, 2.0, 2.0]) """scale of the crop""" diff --git a/nerfstudio/utils/colormaps.py b/nerfstudio/utils/colormaps.py index 7346ca379..27b7e4e21 100644 --- a/nerfstudio/utils/colormaps.py +++ b/nerfstudio/utils/colormaps.py @@ -19,7 +19,8 @@ from typing import Literal, Optional import matplotlib import torch -from torchtyping import TensorType +from jaxtyping import Bool, Float +from torch import Tensor from nerfstudio.utils import colors @@ -43,10 +44,10 @@ class ColormapOptions: def apply_colormap( - image: TensorType["bs":..., 1], + image: Float[Tensor, "*bs channels"], colormap_options: ColormapOptions = ColormapOptions(), eps: float = 1e-9, -) -> TensorType["bs":..., "rgb":3]: +) -> Float[Tensor, "*bs rgb=3"]: """ Applies a colormap to a tensor image. If single channel, applies a colormap to the image. @@ -89,9 +90,7 @@ def apply_colormap( raise NotImplementedError -def apply_float_colormap( - image: TensorType["bs":..., 1], colormap: Colormaps = "viridis" -) -> TensorType["bs":..., "rgb":3]: +def apply_float_colormap(image: Float[Tensor, "*bs 1"], colormap: Colormaps = "viridis") -> Float[Tensor, "*bs rgb=3"]: """Convert single channel to a color image. Args: @@ -99,7 +98,7 @@ def apply_float_colormap( colormap: Colormap for image. Returns: - TensorType: Colored image with colors in [0, 1] + Tensor: Colored image with colors in [0, 1] """ if colormap == "default": colormap = "turbo" @@ -116,12 +115,12 @@ def apply_float_colormap( def apply_depth_colormap( - depth: TensorType["bs":..., 1], - accumulation: Optional[TensorType["bs":..., 1]] = None, + depth: Float[Tensor, "*bs 1"], + accumulation: Optional[Float[Tensor, "*bs 1"]] = None, near_plane: Optional[float] = None, far_plane: Optional[float] = None, colormap_options: ColormapOptions = ColormapOptions(), -) -> TensorType["bs":..., "rgb":3]: +) -> Float[Tensor, "*bs rgb=3"]: """Converts a depth image to color for easier analysis. Args: @@ -151,10 +150,10 @@ def apply_depth_colormap( def apply_boolean_colormap( - image: TensorType["bs":..., 1, bool], - true_color: TensorType["bs":..., "rgb":3] = colors.WHITE, - false_color: TensorType["bs":..., "rgb":3] = colors.BLACK, -) -> TensorType["bs":..., "rgb":3]: + image: Bool[Tensor, "*bs 1"], + true_color: Float[Tensor, "*bs rgb=3"] = colors.WHITE, + false_color: Float[Tensor, "*bs rgb=3"] = colors.BLACK, +) -> Float[Tensor, "*bs rgb=3"]: """Converts a depth image to color for easier analysis. Args: @@ -172,7 +171,7 @@ def apply_boolean_colormap( return colored_image -def apply_pca_colormap(image: TensorType["bs":..., "dim"]) -> TensorType["bs":..., "rgb":3]: +def apply_pca_colormap(image: Float[Tensor, "*bs dim"]) -> Float[Tensor, "*bs rgb=3"]: """Convert feature image to 3-channel RGB via PCA. The first three principle components are used for the color channels, with outlier rejection per-channel @@ -180,7 +179,7 @@ def apply_pca_colormap(image: TensorType["bs":..., "dim"]) -> TensorType["bs":.. image: image of arbitrary vectors Returns: - TensorType: Colored image + Tensor: Colored image """ original_shape = image.shape image = image.view(-1, image.shape[-1]) diff --git a/nerfstudio/utils/colors.py b/nerfstudio/utils/colors.py index 5ec1b80b8..66ac8d243 100644 --- a/nerfstudio/utils/colors.py +++ b/nerfstudio/utils/colors.py @@ -16,7 +16,8 @@ from typing import Union import torch -from torchtyping import TensorType +from jaxtyping import Float +from torch import Tensor WHITE = torch.tensor([1.0, 1.0, 1.0]) BLACK = torch.tensor([0.0, 0.0, 0.0]) @@ -33,13 +34,13 @@ COLORS_DICT = { } -def get_color(color: Union[str, list]) -> TensorType[3]: +def get_color(color: Union[str, list]) -> Float[Tensor, "3"]: """ Args: - color (Union[str, list]): Color as a string or a rgb list + Color as a string or a rgb list Returns: - TensorType[3]: Parsed color + Parsed color """ if isinstance(color, str): color = color.lower() diff --git a/nerfstudio/utils/math.py b/nerfstudio/utils/math.py index 57a1e3421..7f46f31e9 100644 --- a/nerfstudio/utils/math.py +++ b/nerfstudio/utils/math.py @@ -18,10 +18,13 @@ from dataclasses import dataclass from typing import Literal, Tuple import torch -from torchtyping import TensorType +from jaxtyping import Bool, Float +from torch import Tensor -def components_from_spherical_harmonics(levels: int, directions: TensorType[..., 3]) -> TensorType[..., "components"]: +def components_from_spherical_harmonics( + levels: int, directions: Float[Tensor, "*batch 3"] +) -> Float[Tensor, "*batch components"]: """ Returns value for each component of spherical harmonics. @@ -94,15 +97,15 @@ class Gaussians: cov: Covariance of multivariate Gaussian. """ - mean: TensorType[..., "dim"] - cov: TensorType[..., "dim", "dim"] + mean: Float[Tensor, "*batch dim"] + cov: Float[Tensor, "*batch dim dim"] def compute_3d_gaussian( - directions: TensorType[..., 3], - means: TensorType[..., 3], - dir_variance: TensorType[..., 1], - radius_variance: TensorType[..., 1], + directions: Float[Tensor, "*batch 3"], + means: Float[Tensor, "*batch 3"], + dir_variance: Float[Tensor, "*batch 1"], + radius_variance: Float[Tensor, "*batch 1"], ) -> Gaussians: """Compute gaussian along ray. @@ -127,11 +130,11 @@ def compute_3d_gaussian( def cylinder_to_gaussian( - origins: TensorType[..., 3], - directions: TensorType[..., 3], - starts: TensorType[..., 1], - ends: TensorType[..., 1], - radius: TensorType[..., 1], + origins: Float[Tensor, "*batch 3"], + directions: Float[Tensor, "*batch 3"], + starts: Float[Tensor, "*batch 1"], + ends: Float[Tensor, "*batch 1"], + radius: Float[Tensor, "*batch 1"], ) -> Gaussians: """Approximates cylinders with a Gaussian distributions. @@ -152,11 +155,11 @@ def cylinder_to_gaussian( def conical_frustum_to_gaussian( - origins: TensorType[..., 3], - directions: TensorType[..., 3], - starts: TensorType[..., 1], - ends: TensorType[..., 1], - radius: TensorType[..., 1], + origins: Float[Tensor, "*batch 3"], + directions: Float[Tensor, "*batch 3"], + starts: Float[Tensor, "*batch 1"], + ends: Float[Tensor, "*batch 1"], + radius: Float[Tensor, "*batch 1"], ) -> Gaussians: """Approximates conical frustums with a Gaussian distributions. @@ -236,9 +239,9 @@ def intersect_aabb( def safe_normalize( - vectors: TensorType["batch_dim":..., "N"], + vectors: Float[Tensor, "*batch_dim N"], eps: float = 1e-10, -) -> TensorType["batch_dim":..., "N"]: +) -> Float[Tensor, "*batch_dim N"]: """Normalizes vectors. Args: @@ -252,7 +255,9 @@ def safe_normalize( def masked_reduction( - input_tensor: TensorType[1, 32, "mult"], mask: TensorType[1, 32, "mult"], reduction_type: Literal["image", "batch"] + input_tensor: Float[Tensor, "1 32 mult"], + mask: Bool[Tensor, "1 32 mult"], + reduction_type: Literal["image", "batch"], ): """ Whether to consolidate the input_tensor across the batch or across the image @@ -277,7 +282,7 @@ def masked_reduction( def normalized_depth_scale_and_shift( - prediction: TensorType[1, 32, "mult"], target: TensorType[1, 32, "mult"], mask: TensorType[1, 32, "mult"] + prediction: Float[Tensor, "1 32 mult"], target: Float[Tensor, "1 32 mult"], mask: Bool[Tensor, "1 32 mult"] ): """ More info here: https://arxiv.org/pdf/2206.00665.pdf supplementary section A2 Depth Consistency Loss diff --git a/nerfstudio/utils/plotly_utils.py b/nerfstudio/utils/plotly_utils.py index 967fb3db9..aa08a6008 100644 --- a/nerfstudio/utils/plotly_utils.py +++ b/nerfstudio/utils/plotly_utils.py @@ -27,8 +27,9 @@ import numpy as np import plotly.graph_objects as go import torch import torch.nn.functional as F +from jaxtyping import Float from plotly import express as ex -from torchtyping import TensorType +from torch import Tensor from nerfstudio.cameras.cameras import Cameras from nerfstudio.cameras.rays import Frustums, RayBundle @@ -49,7 +50,7 @@ def color_str(color): def get_line_segments_from_lines( - lines: TensorType["num_rays", 2, 3], + lines: Float[Tensor, "num_rays 2 3"], color: str = color_str((1, 0, 0)), marker_color: str = color_str((1, 0, 0)), colors: Optional[List[str]] = None, @@ -99,7 +100,7 @@ def get_line_segments_from_lines( return data -def vis_dataset(camera_origins: TensorType["num_cameras", 3], ray_bundle: RayBundle) -> go.FigureWidget: # type: ignore +def vis_dataset(camera_origins: Float[Tensor, "num_cameras 3"], ray_bundle: RayBundle) -> go.FigureWidget: """Visualize a dataset with plotly using our cameras and generated rays. Args: @@ -165,7 +166,7 @@ def get_random_color(colormap: Optional[List[str]] = None, idx: Optional[int] = def get_sphere( - radius: float, center: TensorType[3] = None, color: str = "black", opacity: float = 1.0, resolution: int = 32 + radius: float, center: Float[Tensor, "3"] = None, color: str = "black", opacity: float = 1.0, resolution: int = 32 ) -> go.Mesh3d: # type: ignore """Returns a sphere object for plotting with plotly. @@ -206,7 +207,7 @@ def get_sphere( def get_cube( side_length: Union[float, torch.Tensor], - center: TensorType[3] = None, + center: Float[Tensor, "3"] = None, color: str = "black", opacity: float = 1.0, ) -> go.Mesh3d: # type: ignore @@ -253,8 +254,8 @@ def get_cube( def get_gaussian_ellipsiod( - mean: TensorType[3], - cov: TensorType[3, 3], + mean: Float[Tensor, "3"], + cov: Float[Tensor, "3 3"], n_std: int = 2, color="lightblue", opacity: float = 0.5, diff --git a/nerfstudio/utils/poses.py b/nerfstudio/utils/poses.py index 1b026404f..e0f495747 100644 --- a/nerfstudio/utils/poses.py +++ b/nerfstudio/utils/poses.py @@ -17,10 +17,11 @@ Common 3D pose methods """ import torch -from torchtyping import TensorType +from jaxtyping import Float +from torch import Tensor -def to4x4(pose: TensorType[..., 3, 4]) -> TensorType[..., 4, 4]: +def to4x4(pose: Float[Tensor, "*batch 3 4"]) -> Float[Tensor, "*batch 4 4"]: """Convert 3x4 pose matrices to a 4x4 with the addition of a homogeneous coordinate. Args: @@ -34,7 +35,7 @@ def to4x4(pose: TensorType[..., 3, 4]) -> TensorType[..., 4, 4]: return torch.cat([pose, constants], dim=-2) -def inverse(pose: TensorType[..., 3, 4]) -> TensorType[..., 3, 4]: +def inverse(pose: Float[Tensor, "*batch 3 4"]) -> Float[Tensor, "*batch 3 4"]: """Invert provided pose matrix. Args: @@ -50,7 +51,7 @@ def inverse(pose: TensorType[..., 3, 4]) -> TensorType[..., 3, 4]: return torch.cat([R_inverse, t_inverse], dim=-1) -def multiply(pose_a: TensorType[..., 3, 4], pose_b: TensorType[..., 3, 4]) -> TensorType[..., 3, 4]: +def multiply(pose_a: Float[Tensor, "*batch 3 4"], pose_b: Float[Tensor, "*batch 3 4"]) -> Float[Tensor, "*batch 3 4"]: """Multiply two pose matrices, A @ B. Args: @@ -67,7 +68,7 @@ def multiply(pose_a: TensorType[..., 3, 4], pose_b: TensorType[..., 3, 4]) -> Te return torch.cat([R, t], dim=-1) -def normalize(poses: TensorType[..., 3, 4]) -> TensorType[..., 3, 4]: +def normalize(poses: Float[Tensor, "*batch 3 4"]) -> Float[Tensor, "*batch 3 4"]: """Normalize the XYZs of poses to fit within a unit cube ([-1, 1]). Note: This operation is not in-place. Args: diff --git a/nerfstudio/utils/tensor_dataclass.py b/nerfstudio/utils/tensor_dataclass.py index cf0b1a1ed..7eca91610 100644 --- a/nerfstudio/utils/tensor_dataclass.py +++ b/nerfstudio/utils/tensor_dataclass.py @@ -152,8 +152,12 @@ class TensorDataclass: if isinstance(indices, (int, slice, type(Ellipsis))): indices = (indices,) assert isinstance(indices, tuple) - tensor_fn = lambda x: x[indices + (slice(None),)] - dataclass_fn = lambda x: x[indices] + + def tensor_fn(x): + return x[indices + (slice(None),)] + + def dataclass_fn(x): + return x[indices] def custom_tensor_dims_fn(k, v): custom_dims = self._field_custom_dimensions[k] # pylint: disable=unsubscriptable-object @@ -207,8 +211,12 @@ class TensorDataclass: """ if isinstance(shape, int): shape = (shape,) - tensor_fn = lambda x: x.reshape((*shape, x.shape[-1])) - dataclass_fn = lambda x: x.reshape(shape) + + def tensor_fn(x): + return x.reshape((*shape, x.shape[-1])) + + def dataclass_fn(x): + return x.reshape(shape) def custom_tensor_dims_fn(k, v): custom_dims = self._field_custom_dimensions[k] # pylint: disable=unsubscriptable-object diff --git a/nerfstudio/utils/writer.py b/nerfstudio/utils/writer.py index 03cbb36c3..d6fcef3cc 100644 --- a/nerfstudio/utils/writer.py +++ b/nerfstudio/utils/writer.py @@ -26,15 +26,21 @@ from typing import Any, Dict, List, Optional, Union import torch import wandb +from jaxtyping import Float +from torch import Tensor from torch.utils.tensorboard import SummaryWriter -from torchtyping import TensorType from nerfstudio.configs import base_config as cfg from nerfstudio.utils.decorators import check_main_thread, decorate_all from nerfstudio.utils.printing import human_format from nerfstudio.utils.rich_utils import CONSOLE -to8b = lambda x: (255 * torch.clamp(x, min=0, max=1)).to(torch.uint8) + +def to8b(x): + """Converts a torch tensor to 8 bit""" + return (255 * torch.clamp(x, min=0, max=1)).to(torch.uint8) + + EVENT_WRITERS = [] EVENT_STORAGE = [] GLOBAL_BUFFER = {} @@ -63,7 +69,7 @@ class EventType(enum.Enum): @check_main_thread -def put_image(name, image: TensorType["H", "W", "C"], step: int): +def put_image(name, image: Float[Tensor, "H W C"], step: int): """Setter function to place images into the queue to be written out Args: @@ -227,7 +233,7 @@ class Writer: """Writer class""" @abstractmethod - def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None: + def write_image(self, name: str, image: Float[Tensor, "H W C"], step: int) -> None: """method to write out image Args: @@ -301,7 +307,7 @@ class WandbWriter(Writer): reinit=True, ) - def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None: + def write_image(self, name: str, image: Float[Tensor, "H W C"], step: int) -> None: image = torch.permute(image, (2, 0, 1)) wandb.log({name: wandb.Image(image)}, step=step) @@ -326,7 +332,7 @@ class TensorboardWriter(Writer): def __init__(self, log_dir: Path): self.tb_writer = SummaryWriter(log_dir=log_dir) - def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None: + def write_image(self, name: str, image: Float[Tensor, "H W C"], step: int) -> None: image = to8b(image) self.tb_writer.add_image(name, image, step, dataformats="HWC") diff --git a/nerfstudio/viewer/app/run_deploy.py b/nerfstudio/viewer/app/run_deploy.py index deff54ca9..8a5fce02d 100644 --- a/nerfstudio/viewer/app/run_deploy.py +++ b/nerfstudio/viewer/app/run_deploy.py @@ -17,7 +17,6 @@ Code for deploying the built viewer folder to a server and handing versioning. We use the library sshconf (https://github.com/sorend/sshconf) for working with the ssh config file. """ import json -import os import subprocess from os.path import expanduser from pathlib import Path diff --git a/nerfstudio/viewer/server/state/node.py b/nerfstudio/viewer/server/state/node.py index b161c0bfb..ca08b5cd7 100644 --- a/nerfstudio/viewer/server/state/node.py +++ b/nerfstudio/viewer/server/state/node.py @@ -37,7 +37,10 @@ def get_tree(node_class: Callable) -> Callable: or the current node if it does. """ assert isinstance(node_class(), Node) - tree = lambda: node_class(tree) + + def tree(): + return node_class(tree) + return tree() diff --git a/nerfstudio/viewer/server/state/state_node.py b/nerfstudio/viewer/server/state/state_node.py index 0e8ce0853..3aebfeff2 100644 --- a/nerfstudio/viewer/server/state/state_node.py +++ b/nerfstudio/viewer/server/state/state_node.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import List from nerfstudio.viewer.server.state.node import Node diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index 07a25b026..6b18391a9 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/nerfstudio/viewer/server/viewer_utils.py @@ -21,7 +21,7 @@ import os import socket import sys from pathlib import Path -from typing import TYPE_CHECKING, Optional, Tuple +from typing import Optional, Tuple import torch @@ -29,9 +29,6 @@ from nerfstudio.data.scene_box import SceneBox from nerfstudio.models.base_model import Model from nerfstudio.utils.io import load_from_json -if TYPE_CHECKING: - from nerfstudio.engine.trainer import Trainer - def get_viewer_version() -> str: """Return the version of the viewer.""" diff --git a/pyproject.toml b/pyproject.toml index abe400912..9e9deb9c4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,7 @@ dependencies = [ "imageio>=2.21.1", 'importlib-metadata>=6.0.0; python_version < "3.10"', "ipywidgets>=7.6", + "jaxtyping>=0.2.15", "jupyterlab>=3.3.4", "matplotlib>=3.5.3", "mediapy>=1.1.0", @@ -49,7 +50,6 @@ dependencies = [ "tensorboard>=2.13.0", "torch>=1.12.1,<2.0.0", "torchmetrics[image]>=0.9.3", - "torchtyping>=0.1.4", "torchvision>=0.13.0", "typing_extensions>=4.4.0", "viser>=0.0.5", @@ -149,7 +149,7 @@ disable = [ #pytest [tool.pytest.ini_options] -addopts = "-n=4 --typeguard-packages=nerfstudio --torchtyping-patch-typeguard --disable-warnings" +addopts = "-n=4 --typeguard-packages=nerfstudio --disable-warnings" testpaths = [ "tests", ] diff --git a/tests/dataparsers/test_nerfstudio_dataparser.py b/tests/dataparsers/test_nerfstudio_dataparser.py index c9dce5284..f3f6f1223 100644 --- a/tests/dataparsers/test_nerfstudio_dataparser.py +++ b/tests/dataparsers/test_nerfstudio_dataparser.py @@ -72,7 +72,6 @@ def test_nerfstudio_dataparser_split_filelist(mocked_dataset): json.dump(data, f) from nerfstudio.data.dataparsers.nerfstudio_dataparser import ( - DataparserOutputs, Nerfstudio, NerfstudioDataParserConfig, ) diff --git a/tests/field_components/test_fields.py b/tests/field_components/test_fields.py index c379a567e..81437645a 100644 --- a/tests/field_components/test_fields.py +++ b/tests/field_components/test_fields.py @@ -9,10 +9,8 @@ from nerfstudio.fields.instant_ngp_field import TCNNInstantNGPField def test_tcnn_instant_ngp_field(): """Test the tiny-cuda-nn field""" - # pylint: disable=import-outside-toplevel - # pylint: disable=unused-import try: - import tinycudann as tcnn + import tinycudann as tcnn # pylint:disable=import-outside-toplevel, unused-import except ImportError as e: # tinycudann module doesn't exist print(e) diff --git a/tests/field_components/test_temporal_grid.py b/tests/field_components/test_temporal_grid.py index a83450c38..e539b3030 100644 --- a/tests/field_components/test_temporal_grid.py +++ b/tests/field_components/test_temporal_grid.py @@ -12,7 +12,7 @@ def test_temporal_grid(): print("Unable to test temporal grid without GPU, since CUDA kernels involved.") return - params = dict( + model = temporal_grid.TemporalGridEncoder( temporal_dim=2, input_dim=1, num_levels=1, @@ -23,8 +23,7 @@ def test_temporal_grid(): desired_resolution=None, gridtype="tiled", align_corners=False, - ) - model = temporal_grid.TemporalGridEncoder(**params).cuda() + ).cuda() random_embedding = torch.rand_like(model.embeddings) random_embedding[:, 0] = torch.arange(8).to(model.embeddings) model.embeddings = torch.nn.Parameter(random_embedding, requires_grad=True) diff --git a/tests/utils/test_tensor_dataclass.py b/tests/utils/test_tensor_dataclass.py index 838655ed4..1ade936a4 100644 --- a/tests/utils/test_tensor_dataclass.py +++ b/tests/utils/test_tensor_dataclass.py @@ -36,9 +36,9 @@ def test_init(): dummy_vals: torch.Tensor = None - dummy = Dummy(dummy_vals=torch.ones(1)) + Dummy(dummy_vals=torch.ones(1)) with pytest.raises(ValueError): - dummy = Dummy() + Dummy() def test_broadcasting():