diff --git a/nerfstudio/cameras/camera_utils.py b/nerfstudio/cameras/camera_utils.py index 5e7de6179..b1dd5a267 100644 --- a/nerfstudio/cameras/camera_utils.py +++ b/nerfstudio/cameras/camera_utils.py @@ -21,13 +21,14 @@ from typing import List, Optional, Tuple import numpy as np import torch +from numpy.typing import ArrayLike from torchtyping import TensorType from typing_extensions import Literal _EPS = np.finfo(float).eps * 4.0 -def unit_vector(data, axis: Optional[int] = None) -> np.ndarray: +def unit_vector(data: ArrayLike, axis: Optional[int] = None) -> np.ndarray: """Return ndarray normalized by length, i.e. Euclidean norm, along axis. Args: @@ -46,7 +47,7 @@ def unit_vector(data, axis: Optional[int] = None) -> np.ndarray: return data -def quaternion_from_matrix(matrix, isprecise: bool = False) -> np.ndarray: +def quaternion_from_matrix(matrix: ArrayLike, isprecise: bool = False) -> np.ndarray: """Return quaternion from rotation matrix. Args: @@ -102,7 +103,9 @@ def quaternion_from_matrix(matrix, isprecise: bool = False) -> np.ndarray: return q -def quaternion_slerp(quat0, quat1, fraction: float, spin: int = 0, shortestpath: bool = True) -> np.ndarray: +def quaternion_slerp( + quat0: ArrayLike, quat1: ArrayLike, fraction: float, spin: int = 0, shortestpath: bool = True +) -> np.ndarray: """Return spherical linear interpolation between two quaternions. Args: quat0: first quaternion @@ -136,7 +139,7 @@ def quaternion_slerp(quat0, quat1, fraction: float, spin: int = 0, shortestpath: return q0 -def quaternion_matrix(quaternion) -> np.ndarray: +def quaternion_matrix(quaternion: ArrayLike) -> np.ndarray: """Return homogeneous rotation matrix from quaternion. Args: @@ -158,7 +161,7 @@ def quaternion_matrix(quaternion) -> np.ndarray: ) -def get_interpolated_poses(pose_a, pose_b, steps: int = 10) -> List[float]: +def get_interpolated_poses(pose_a: ArrayLike, pose_b: ArrayLike, steps: int = 10) -> List[float]: """Return interpolation of poses with specified number of steps. Args: poseA: first pose @@ -229,7 +232,7 @@ def get_interpolated_poses_many( return torch.tensor(traj), torch.tensor(k_interp) -def normalize(x) -> TensorType[...]: +def normalize(x: torch.Tensor) -> TensorType[...]: """Returns a normalized vector.""" return x / torch.linalg.norm(x) @@ -249,7 +252,7 @@ def normalize_with_norm(x: torch.Tensor, dim: int) -> Tuple[torch.Tensor, torch. return x / norm, norm -def viewmatrix(lookat, up, pos) -> TensorType[...]: +def viewmatrix(lookat: torch.Tensor, up: torch.Tensor, pos: torch.Tensor) -> TensorType[...]: """Returns a camera transformation matrix. Args: diff --git a/nerfstudio/cameras/cameras.py b/nerfstudio/cameras/cameras.py index 01b71bfdb..e0ed08e4f 100644 --- a/nerfstudio/cameras/cameras.py +++ b/nerfstudio/cameras/cameras.py @@ -36,6 +36,8 @@ from nerfstudio.data.scene_box import SceneBox from nerfstudio.utils.misc import strtobool from nerfstudio.utils.tensor_dataclass import TensorDataclass +TORCH_DEVICE = Union[torch.device, str] # pylint: disable=invalid-name + class CameraType(Enum): """Supported camera types.""" @@ -109,7 +111,7 @@ class Cameras(TensorDataclass): ] ] = CameraType.PERSPECTIVE, times: Optional[TensorType["num_cameras"]] = None, - ): + ) -> None: """Initializes the Cameras object. Note on Input Tensor Dimensions: All of these tensors have items of dimensions TensorType[3, 4] @@ -148,7 +150,7 @@ class Cameras(TensorDataclass): self._use_nerfacc = strtobool(os.environ.get("INTERSECT_WITH_NERFACC", "TRUE")) - def _init_get_fc_xy(self, fc_xy, name): + def _init_get_fc_xy(self, fc_xy: Union[float, torch.Tensor], name: str) -> torch.Tensor: """ Parses the input focal length / principle point x or y and returns a tensor of the correct shape @@ -243,7 +245,7 @@ class Cameras(TensorDataclass): raise ValueError("Height must be an int, tensor, or None, received: " + str(type(h_w))) return h_w - def _init_get_times(self, times): + def _init_get_times(self, times: Union[None, torch.Tensor]) -> Union[None, torch.Tensor]: if times is None: times = None elif isinstance(times, torch.Tensor): @@ -255,7 +257,7 @@ class Cameras(TensorDataclass): return times @property - def device(self): + def device(self) -> TORCH_DEVICE: """Returns the device that the camera is on.""" return self.camera_to_worlds.device @@ -270,7 +272,7 @@ class Cameras(TensorDataclass): return self.width @property - def is_jagged(self): + def is_jagged(self) -> bool: """ Returns whether or not the cameras are "jagged" (i.e. the height and widths are different, meaning that you cannot concatenate the image coordinate maps together) diff --git a/nerfstudio/cameras/rays.py b/nerfstudio/cameras/rays.py index c65315c26..fd4e8129e 100644 --- a/nerfstudio/cameras/rays.py +++ b/nerfstudio/cameras/rays.py @@ -17,7 +17,7 @@ Some ray datastructures. """ import random from dataclasses import dataclass -from typing import Callable, Dict, Optional +from typing import Callable, Dict, Optional, Union import torch from torchtyping import TensorType @@ -25,6 +25,8 @@ from torchtyping import TensorType from nerfstudio.utils.math import Gaussians, conical_frustum_to_gaussian from nerfstudio.utils.tensor_dataclass import TensorDataclass +TORCH_DEVICE = Union[str, torch.device] # pylint: disable=invalid-name + @dataclass class Frustums(TensorDataclass): @@ -77,7 +79,7 @@ class Frustums(TensorDataclass): ) @classmethod - def get_mock_frustum(cls, device="cpu") -> "Frustums": + def get_mock_frustum(cls, device: Optional[TORCH_DEVICE] = "cpu") -> "Frustums": """Helper function to generate a placeholder frustum. Returns: @@ -169,7 +171,7 @@ class RayBundle(TensorDataclass): """ self.camera_indices = torch.ones_like(self.origins[..., 0:1]).long() * camera_index - def __len__(self): + def __len__(self) -> int: num_rays = torch.numel(self.origins) // self.origins.shape[-1] return num_rays diff --git a/nerfstudio/exporter/texture_utils.py b/nerfstudio/exporter/texture_utils.py index 9f2825ce3..4001d1c60 100644 --- a/nerfstudio/exporter/texture_utils.py +++ b/nerfstudio/exporter/texture_utils.py @@ -22,7 +22,7 @@ from __future__ import annotations import math from pathlib import Path -from typing import Optional, Tuple +from typing import Optional, Tuple, Union import mediapy as media import numpy as np @@ -39,6 +39,8 @@ from nerfstudio.utils.rich_utils import get_progress CONSOLE = Console(width=120) +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] @@ -56,7 +58,9 @@ def get_parallelogram_area( return (p[..., 0] - v0[..., 0]) * (v1[..., 1] - v0[..., 1]) - (p[..., 1] - v0[..., 1]) * (v1[..., 0] - v0[..., 0]) -def get_texture_image(num_pixels_w, num_pixels_h, device): +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]]: """Get a texture image.""" px_w = 1.0 / num_pixels_w px_h = 1.0 / num_pixels_h @@ -80,7 +84,11 @@ def unwrap_mesh_per_uv_triangle( faces: TensorType["num_faces", 3], vertex_normals: TensorType["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"], +]: """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 interpolation to get the origins and directions, per pixel, that are needed to render the NeRF with. @@ -331,7 +339,7 @@ def export_textured_mesh( unwrap_method: Literal["xatlas", "custom"] = "xatlas", raylen_method: Literal["edge", "none"] = "edge", num_pixels_per_side=1024, -): +) -> None: """Textures a mesh using the radiance field from the Pipeline. The mesh is written to an OBJ file in the output directory, along with the corresponding material and texture files. diff --git a/nerfstudio/exporter/tsdf_utils.py b/nerfstudio/exporter/tsdf_utils.py index 049f6354d..2249fde23 100644 --- a/nerfstudio/exporter/tsdf_utils.py +++ b/nerfstudio/exporter/tsdf_utils.py @@ -37,6 +37,8 @@ from nerfstudio.pipelines.base_pipeline import Pipeline CONSOLE = Console(width=120) +TORCH_DEVICE = Union[torch.device, str] # pylint: disable=invalid-name + @dataclass class TSDF: @@ -59,7 +61,7 @@ class TSDF: truncation_margin: float = 5.0 """Margin for truncation.""" - def to(self, device: str): + def to(self, device: TORCH_DEVICE): """Move the tensors to the specified device. Args: @@ -74,12 +76,12 @@ class TSDF: return self @property - def device(self): + def device(self) -> TORCH_DEVICE: """Returns the device that voxel_coords is on.""" return self.voxel_coords.device @property - def truncation(self): + def truncation(self) -> float: """Returns the truncation distance.""" # TODO: clean this up truncation = self.voxel_size[0] * self.truncation_margin @@ -136,7 +138,7 @@ class TSDF: return Mesh(vertices=vertices, faces=faces, normals=normals, colors=colors) @classmethod - def export_mesh(cls, mesh: Mesh, filename: str): + def export_mesh(cls, mesh: Mesh, filename: str) -> None: """Exports the mesh to a file. We use pymeshlab to export the mesh as a PLY file. @@ -172,7 +174,7 @@ class TSDF: 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, - ): + ) -> None: """Integrates a batch of depth images into the TSDF. Args: @@ -242,7 +244,6 @@ class TSDF: # Sequentially update the TSDF... for i in range(batch_size): - valid_points_i = valid_points[i] valid_points_i_shape = valid_points_i.view(*shape) # [xdim, ydim, zdim] @@ -281,7 +282,7 @@ def export_tsdf_mesh( use_bounding_box: bool = True, 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), -): +) -> None: """Export a TSDF mesh from a pipeline. Args: diff --git a/nerfstudio/plugins/registry.py b/nerfstudio/plugins/registry.py index 365f05a12..9e7a61c8d 100644 --- a/nerfstudio/plugins/registry.py +++ b/nerfstudio/plugins/registry.py @@ -21,6 +21,7 @@ import typing as t from rich.progress import Console +from nerfstudio.engine.trainer import TrainerConfig from nerfstudio.plugins.types import MethodSpecification if sys.version_info < (3, 10): @@ -30,7 +31,7 @@ else: CONSOLE = Console(width=120) -def discover_methods(): +def discover_methods() -> t.Tuple[t.Dict[str, TrainerConfig], t.Dict[str, str]]: """ Discovers all methods registered using the `nerfstudio.method_configs` entrypoint. """