diff --git a/.github/workflows/core_code_checks.yml b/.github/workflows/core_code_checks.yml index 179b4cd9d..761d4841b 100644 --- a/.github/workflows/core_code_checks.yml +++ b/.github/workflows/core_code_checks.yml @@ -36,6 +36,9 @@ jobs: run: ruff docs/ nerfstudio/ tests/ - name: Run Black run: black docs/ nerfstudio/ tests/ --check + - name: Run Pyright + run: | + pyright - name: Test with pytest run: | pytest diff --git a/docs/reference/contributing.md b/docs/reference/contributing.md index 5b2a82ad0..dd0fdf5b4 100644 --- a/docs/reference/contributing.md +++ b/docs/reference/contributing.md @@ -18,6 +18,7 @@ Below are the various tooling features our team uses to maintain this codebase. | --------------- | ---------------------------------------------------------- | | Formatting | [Black](https://black.readthedocs.io/en/stable/) | | Linter | [Ruff](https://beta.ruff.rs/docs/) | +| Type checking | [Pyright](https://github.com/microsoft/pyright) | | Testing | [pytest](https://docs.pytest.org/en/7.1.x/) | | Docs | [Sphinx](https://www.sphinx-doc.org/en/master/) | | Docstring style | [Google](https://google.github.io/styleguide/pyguide.html) | @@ -53,7 +54,8 @@ conda install -c conda-forge pandoc This will perform the following checks and actions: - - Black/ Linting style check: Ensures code is consistently and properly formatted. + - Formatting and linting: Ensures code is consistently and properly formatted. + - Type checking: Ensures static type safety. - Pytests: Runs pytests locally to make sure added code does not break existing logic. - Documentation build: Builds docs locally. Ensures changes do not result in warnings/errors. - Licensing: Automatically adds licensing headers to the correct files. diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 9cfafc263..2e92fb2ba 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -38,7 +38,6 @@ from typing import ( import torch from torch import nn from torch.nn import Parameter -from torch.utils.data import Dataset from torch.utils.data.distributed import DistributedSampler from typing_extensions import TypeVar @@ -154,8 +153,8 @@ class DataManager(nn.Module): """ - train_dataset: Optional[Dataset] = None - eval_dataset: Optional[Dataset] = None + train_dataset: Optional[InputDataset] = None + eval_dataset: Optional[InputDataset] = None train_sampler: Optional[DistributedSampler] = None eval_sampler: Optional[DistributedSampler] = None includes_time: bool = False @@ -278,9 +277,9 @@ class DataManager(nn.Module): """Returns the number of rays per batch for evaluation.""" raise NotImplementedError - def get_datapath(self) -> Optional[Path]: + @abstractmethod + def get_datapath(self) -> Path: """Returns the path to the data. This is used to determine where to save camera paths.""" - return None def get_training_callbacks( self, training_callback_attributes: TrainingCallbackAttributes @@ -325,7 +324,7 @@ class VanillaDataManagerConfig(DataManagerConfig): camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig() """Specifies the camera pose optimizer used during training. Helpful if poses are noisy, such as for data from Record3D.""" - collate_fn: Callable[[Any], Any] = staticmethod(nerfstudio_collate) + collate_fn: Callable[[Any], Any] = nerfstudio_collate """Specifies the collate function to use for the train and eval dataloaders.""" camera_res_scale_factor: float = 1.0 """The scale factor for scaling spatial data such as images, mask, semantics diff --git a/nerfstudio/data/dataparsers/dycheck_dataparser.py b/nerfstudio/data/dataparsers/dycheck_dataparser.py index 1d7319023..50a581f28 100644 --- a/nerfstudio/data/dataparsers/dycheck_dataparser.py +++ b/nerfstudio/data/dataparsers/dycheck_dataparser.py @@ -74,7 +74,7 @@ def upscale(img, scale: int) -> np.ndarray: return resized -def rescale(img, scale_factor: float, interpolation=cv2.INTER_AREA) -> np.ndarray: # type: ignore +def rescale(img, scale_factor: float, interpolation: int = cv2.INTER_AREA) -> np.ndarray: """Function from DyCheck's repo. Rescale an image. Args: @@ -316,13 +316,13 @@ class Dycheck(DataParser): d = self.config.downscale_factor if not image_filenames[0].exists(): CONSOLE.print(f"downscale factor {d}x not exist, converting") - ori_h, ori_w = cv2.imread(str(self.data / f"rgb/1x/{frame_names[0]}.png")).shape[:2] # type: ignore + ori_h, ori_w = cv2.imread(str(self.data / f"rgb/1x/{frame_names[0]}.png")).shape[:2] (self.data / f"rgb/{d}x").mkdir(exist_ok=True) h, w = ori_h // d, ori_w // d for frame in frame_names: - cv2.imwrite( # type: ignore + cv2.imwrite( str(self.data / f"rgb/{d}x/{frame}.png"), - cv2.resize(cv2.imread(str(self.data / f"rgb/1x/{frame}.png")), [h, w]), # type: ignore + cv2.resize(cv2.imread(str(self.data / f"rgb/1x/{frame}.png")), (h, w)), ) CONSOLE.print("finished") @@ -331,8 +331,8 @@ class Dycheck(DataParser): (self.data / f"processed_depth/{d}x").mkdir(exist_ok=True, parents=True) for idx, frame in enumerate(frame_names): depth = np.load(self.data / f"depth/1x/{frame}.npy") - mask = rescale((depth != 0).astype(np.uint8) * 255, 1 / d, cv2.INTER_AREA) # type: ignore - depth = rescale(depth, 1 / d, cv2.INTER_AREA) # type: ignore + mask = rescale((depth != 0).astype(np.uint8) * 255, 1 / d, cv2.INTER_AREA) + depth = rescale(depth, 1 / d, cv2.INTER_AREA) depth[mask != 255] = 0 depth = _rescale_depth(depth, cams[idx]) np.save(str(self.data / f"processed_depth/{d}x/{frame}.npy"), depth) diff --git a/nerfstudio/data/datasets/base_dataset.py b/nerfstudio/data/datasets/base_dataset.py index 8ab110688..ea7cfd1ef 100644 --- a/nerfstudio/data/datasets/base_dataset.py +++ b/nerfstudio/data/datasets/base_dataset.py @@ -29,6 +29,7 @@ from PIL import Image from torch import Tensor from torch.utils.data import Dataset +from nerfstudio.cameras.cameras import Cameras from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs from nerfstudio.data.utils.data_utils import get_image_mask_tensor_from_path @@ -41,6 +42,8 @@ class InputDataset(Dataset): scale_factor: The scaling factor for the dataparser outputs """ + cameras: Cameras + def __init__(self, dataparser_outputs: DataparserOutputs, scale_factor: float = 1.0): super().__init__() self._dataparser_outputs = dataparser_outputs diff --git a/nerfstudio/data/utils/data_utils.py b/nerfstudio/data/utils/data_utils.py index 0f0e911d1..92a85ef8b 100644 --- a/nerfstudio/data/utils/data_utils.py +++ b/nerfstudio/data/utils/data_utils.py @@ -61,7 +61,7 @@ def get_depth_image_from_path( height: int, width: int, scale_factor: float, - interpolation: int = cv2.INTER_NEAREST, # type: ignore + interpolation: int = cv2.INTER_NEAREST, ) -> torch.Tensor: """Loads, rescales and resizes depth images. Filepath points to a 16-bit or 32-bit depth image, or a numpy array `*.npy`. @@ -78,9 +78,9 @@ def get_depth_image_from_path( """ if filepath.suffix == ".npy": image = np.load(filepath) * scale_factor - image = cv2.resize(image, (width, height), interpolation=interpolation) # type: ignore + image = cv2.resize(image, (width, height), interpolation=interpolation) else: - image = cv2.imread(str(filepath.absolute()), cv2.IMREAD_ANYDEPTH) # type: ignore + image = cv2.imread(str(filepath.absolute()), cv2.IMREAD_ANYDEPTH) image = image.astype(np.float64) * scale_factor - image = cv2.resize(image, (width, height), interpolation=interpolation) # type: ignore + image = cv2.resize(image, (width, height), interpolation=interpolation) return torch.from_numpy(image[:, :, np.newaxis]) diff --git a/nerfstudio/data/utils/dataloaders.py b/nerfstudio/data/utils/dataloaders.py index b3b4aea4a..a6216d16d 100644 --- a/nerfstudio/data/utils/dataloaders.py +++ b/nerfstudio/data/utils/dataloaders.py @@ -21,7 +21,7 @@ import concurrent.futures import multiprocessing import random from abc import abstractmethod -from typing import Dict, Optional, Sized, Tuple, Union +from typing import Any, Callable, Dict, Optional, Sized, Tuple, Union import torch from rich.progress import track @@ -54,7 +54,7 @@ class CacheDataloader(DataLoader): num_images_to_sample_from: int = -1, num_times_to_repeat_images: int = -1, device: Union[torch.device, str] = "cpu", - collate_fn=nerfstudio_collate, + collate_fn: Callable[[Any], Any] = nerfstudio_collate, **kwargs, ): self.dataset = dataset diff --git a/nerfstudio/exporter/texture_utils.py b/nerfstudio/exporter/texture_utils.py index 415f658c1..8a304a7b8 100644 --- a/nerfstudio/exporter/texture_utils.py +++ b/nerfstudio/exporter/texture_utils.py @@ -21,7 +21,7 @@ from __future__ import annotations import math from pathlib import Path -from typing import Literal, Tuple, Union +from typing import Literal, Optional, Tuple, Union import mediapy as media import numpy as np @@ -324,7 +324,7 @@ def export_textured_mesh( mesh: Mesh, pipeline: Pipeline, output_dir: Path, - px_per_uv_triangle: int, + px_per_uv_triangle: Optional[int], unwrap_method: Literal["xatlas", "custom"] = "xatlas", raylen_method: Literal["edge", "none"] = "edge", num_pixels_per_side=1024, @@ -338,7 +338,7 @@ def export_textured_mesh( mesh: The mesh to texture. pipeline: The pipeline to use for texturing. output_dir: The directory to write the textured mesh to. - px_per_uv_triangle: The number of pixels per side of UV triangle. + px_per_uv_triangle: The number of pixels per side of UV triangle. Required for "custom" unwrapping. unwrap_method: The method to use for unwrapping the mesh. offset_method: The method to use for computing the ray length to render. num_pixels_per_side: The number of pixels per side of the texture image. @@ -363,6 +363,7 @@ def export_textured_mesh( CONSOLE.print("[bold green]:white_check_mark: Unwrapped mesh with xatlas method") elif unwrap_method == "custom": CONSOLE.print("Unwrapping mesh with custom method...") + assert px_per_uv_triangle is not None texture_coordinates, origins, directions = unwrap_mesh_per_uv_triangle( vertices, faces, vertex_normals, px_per_uv_triangle ) diff --git a/nerfstudio/exporter/tsdf_utils.py b/nerfstudio/exporter/tsdf_utils.py index 70f05caed..d144686c9 100644 --- a/nerfstudio/exporter/tsdf_utils.py +++ b/nerfstudio/exporter/tsdf_utils.py @@ -304,6 +304,7 @@ def export_tsdf_mesh( device = pipeline.device + assert pipeline.datamanager.train_dataset is not None dataparser_outputs = pipeline.datamanager.train_dataset._dataparser_outputs # initialize the TSDF volume diff --git a/nerfstudio/field_components/temporal_distortions.py b/nerfstudio/field_components/temporal_distortions.py index 468b2112b..0a63a122f 100644 --- a/nerfstudio/field_components/temporal_distortions.py +++ b/nerfstudio/field_components/temporal_distortions.py @@ -16,7 +16,7 @@ import abc from enum import Enum -from typing import Any, Dict, Optional, Tuple +from typing import Any, Dict, Tuple import torch from jaxtyping import Float @@ -30,7 +30,7 @@ class TemporalDistortion(nn.Module): """Apply spatial distortions as a function of time""" @abc.abstractmethod - def forward(self, positions: Float[Tensor, "*bs 3"], times: Optional[Float[Tensor, "1"]]) -> Float[Tensor, "*bs 3"]: + def forward(self, positions: Float[Tensor, "*bs 3"], times: Float[Tensor, "*bs 1"]) -> Float[Tensor, "*bs 3"]: """ Args: positions: Samples to translate as a function of time @@ -86,9 +86,7 @@ class DNeRFDistortion(TemporalDistortion): skip_connections=skip_connections, ) - def forward(self, positions, times=None): - if times is None: - return None + def forward(self, positions: Float[Tensor, "*bs 3"], times: Float[Tensor, "*bs 1"]) -> Float[Tensor, "*bs 3"]: p = self.position_encoding(positions) t = self.temporal_encoding(times) return self.mlp_deform(torch.cat([p, t], dim=-1)) diff --git a/nerfstudio/model_components/scene_colliders.py b/nerfstudio/model_components/scene_colliders.py index dfff2dc6a..5dfbd1b3b 100644 --- a/nerfstudio/model_components/scene_colliders.py +++ b/nerfstudio/model_components/scene_colliders.py @@ -33,7 +33,7 @@ class SceneCollider(nn.Module): self.kwargs = kwargs super().__init__() - def set_nears_and_fars(self, ray_bundle) -> RayBundle: + def set_nears_and_fars(self, ray_bundle: RayBundle) -> RayBundle: """To be implemented.""" raise NotImplementedError diff --git a/nerfstudio/models/vanilla_nerf.py b/nerfstudio/models/vanilla_nerf.py index e600ac2ad..4a23bf0e5 100644 --- a/nerfstudio/models/vanilla_nerf.py +++ b/nerfstudio/models/vanilla_nerf.py @@ -143,7 +143,11 @@ class NeRFModel(Model): # uniform sampling ray_samples_uniform = self.sampler_uniform(ray_bundle) if self.temporal_distortion is not None: - offsets = self.temporal_distortion(ray_samples_uniform.frustums.get_positions(), ray_samples_uniform.times) + offsets = None + if ray_samples_uniform.times is not None: + offsets = self.temporal_distortion( + ray_samples_uniform.frustums.get_positions(), ray_samples_uniform.times + ) ray_samples_uniform.frustums.set_offsets(offsets) # coarse field: @@ -159,7 +163,9 @@ class NeRFModel(Model): # pdf sampling ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_coarse) if self.temporal_distortion is not None: - offsets = self.temporal_distortion(ray_samples_pdf.frustums.get_positions(), ray_samples_pdf.times) + offsets = None + if ray_samples_pdf.times is not None: + offsets = self.temporal_distortion(ray_samples_pdf.frustums.get_positions(), ray_samples_pdf.times) ray_samples_pdf.frustums.set_offsets(offsets) # fine field: diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index a17dfbb0c..ebc1c2498 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -39,7 +39,6 @@ from torch.nn.parallel import DistributedDataParallel as DDP from nerfstudio.configs import base_config as cfg from nerfstudio.data.datamanagers.base_datamanager import ( DataManager, - DataManagerConfig, VanillaDataManager, VanillaDataManagerConfig, ) @@ -205,7 +204,7 @@ class VanillaPipelineConfig(cfg.InstantiateConfig): _target: Type = field(default_factory=lambda: VanillaPipeline) """target class to instantiate""" - datamanager: DataManagerConfig = VanillaDataManagerConfig() + datamanager: VanillaDataManagerConfig = VanillaDataManagerConfig() """specifies the datamanager config""" model: ModelConfig = ModelConfig() """specifies the model config""" diff --git a/nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py b/nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py index 3203b16b0..b80923d3e 100644 --- a/nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py +++ b/nerfstudio/process_data/colmap_converter_to_nerfstudio_dataset.py @@ -203,6 +203,10 @@ class ColmapConverterToNerfstudioDataset(BaseConverterToNerfstudioDataset): elif sfm_tool == "hloc": if mask_path is not None: raise RuntimeError("Cannot use a mask with hloc. Please remove the cropping options " "and try again.") + + assert feature_type is not None + assert matcher_type is not None + assert matcher_type != "NN" # Only used for colmap. hloc_utils.run_hloc( image_dir=self.image_dir, colmap_dir=self.absolute_colmap_path, diff --git a/nerfstudio/process_data/colmap_utils.py b/nerfstudio/process_data/colmap_utils.py index b94f530cc..4f17026f1 100644 --- a/nerfstudio/process_data/colmap_utils.py +++ b/nerfstudio/process_data/colmap_utils.py @@ -582,7 +582,7 @@ def create_sfm_depth( out_name = str(im_data.name) depth_path = output_dir / out_name - cv2.imwrite(str(depth_path), depth_img) + cv2.imwrite(str(depth_path), depth_img) # type: ignore image_id_to_depth_path[im_id] = depth_path @@ -594,13 +594,13 @@ def create_sfm_depth( overlay = 255.0 * colormaps.apply_depth_colormap(torch.from_numpy(depth_flat)).numpy() overlay = overlay.reshape([H, W, 3]) input_image_path = input_images_dir / im_data.name - input_image = cv2.imread(str(input_image_path)) + input_image = cv2.imread(str(input_image_path)) # type: ignore debug = 0.3 * input_image + 0.7 + overlay out_name = out_name + ".debug.jpg" output_path = output_dir / "debug_depth" / out_name output_path.parent.mkdir(parents=True, exist_ok=True) - cv2.imwrite(str(output_path), debug.astype(np.uint8)) + cv2.imwrite(str(output_path), debug.astype(np.uint8)) # type: ignore return image_id_to_depth_path diff --git a/nerfstudio/process_data/equirect_utils.py b/nerfstudio/process_data/equirect_utils.py index 573181b5c..4217c97f2 100644 --- a/nerfstudio/process_data/equirect_utils.py +++ b/nerfstudio/process_data/equirect_utils.py @@ -190,6 +190,7 @@ def generate_planar_projections_from_equirectangular( v_rad = torch.pi * v_deg / 180.0 u_rad = torch.pi * u_deg / 180.0 pers_image = equi2pers(im, rots={"roll": 0, "pitch": v_rad, "yaw": u_rad}) * 255.0 + assert isinstance(pers_image, torch.Tensor) pers_image = (pers_image.permute(1, 2, 0)).type(torch.uint8).to("cpu").numpy() cv2.imwrite(f"{output_dir}/{i[:-4]}_{count}.jpg", pers_image) count += 1 diff --git a/nerfstudio/process_data/hloc_utils.py b/nerfstudio/process_data/hloc_utils.py index 2dbf58788..85049b5bf 100644 --- a/nerfstudio/process_data/hloc_utils.py +++ b/nerfstudio/process_data/hloc_utils.py @@ -42,7 +42,7 @@ else: _HAS_HLOC = True try: - from pixsfm.refine_hloc import PixSfM + from pixsfm.refine_hloc import PixSfM # type: ignore except ImportError: _HAS_PIXSFM = False else: diff --git a/nerfstudio/process_data/process_data_utils.py b/nerfstudio/process_data/process_data_utils.py index 15c414dcc..d211b9b14 100644 --- a/nerfstudio/process_data/process_data_utils.py +++ b/nerfstudio/process_data/process_data_utils.py @@ -395,7 +395,7 @@ def find_tool_feature_matcher_combination( "sosnet", "disk", ], - Literal["superglue", "superglue-fast", "NN-superpoint", "NN-ratio", "NN-mutual", "adalam"], + Literal["NN", "superglue", "superglue-fast", "NN-superpoint", "NN-ratio", "NN-mutual", "adalam"], ], ]: """Find a valid combination of sfm tool, feature type, and matcher type. @@ -452,7 +452,7 @@ def generate_circle_mask(height: int, width: int, percent_radius) -> Optional[np mask = np.zeros((height, width), dtype=np.uint8) center = (width // 2, height // 2) radius = int(percent_radius * np.sqrt(width**2 + height**2) / 2.0) - cv2.circle(mask, center, radius, 1, -1) # type: ignore + cv2.circle(mask, center, radius, 1, -1) return mask @@ -523,7 +523,7 @@ def save_mask( The path to the mask file or None if no mask is needed. """ image_path = next(image_dir.glob("frame_*")) - image = cv2.imread(str(image_path)) # type: ignore + image = cv2.imread(str(image_path)) height, width = image.shape[:2] mask = generate_mask(height, width, crop_factor, percent_radius) if mask is None: @@ -531,17 +531,17 @@ def save_mask( mask *= 255 mask_path = image_dir.parent / "masks" mask_path.mkdir(exist_ok=True) - cv2.imwrite(str(mask_path / "mask.png"), mask) # type: ignore + cv2.imwrite(str(mask_path / "mask.png"), mask) downscale_factors = [2**i for i in range(num_downscales + 1)[1:]] for downscale in downscale_factors: mask_path_i = image_dir.parent / f"masks_{downscale}" mask_path_i.mkdir(exist_ok=True) mask_path_i = mask_path_i / "mask.png" - mask_i = cv2.resize( # type: ignore + mask_i = cv2.resize( mask, (width // downscale, height // downscale), - interpolation=cv2.INTER_NEAREST, # type: ignore + interpolation=cv2.INTER_NEAREST, ) - cv2.imwrite(str(mask_path_i), mask_i) # type: ignore + cv2.imwrite(str(mask_path_i), mask_i) CONSOLE.log(":tada: Generated and saved masks.") return mask_path / "mask.png" diff --git a/nerfstudio/process_data/realitycapture_utils.py b/nerfstudio/process_data/realitycapture_utils.py index 9230ae96a..c18afbd44 100644 --- a/nerfstudio/process_data/realitycapture_utils.py +++ b/nerfstudio/process_data/realitycapture_utils.py @@ -17,7 +17,7 @@ import csv import json from pathlib import Path -from typing import List +from typing import Dict, List import numpy as np from PIL import Image @@ -27,7 +27,7 @@ from nerfstudio.utils.rich_utils import CONSOLE def realitycapture_to_json( - image_filename_map: List[Path], + image_filename_map: Dict[str, Path], csv_filename: Path, output_dir: Path, verbose: bool = False, diff --git a/nerfstudio/scripts/datasets/process_nuscenes_masks.py b/nerfstudio/scripts/datasets/process_nuscenes_masks.py index 77efc3432..afbd573d6 100644 --- a/nerfstudio/scripts/datasets/process_nuscenes_masks.py +++ b/nerfstudio/scripts/datasets/process_nuscenes_masks.py @@ -100,15 +100,17 @@ class ProcessNuScenesMasks: mask[-100:] = 0 for box in boxes: - # Dont mask out static objects (static in all frames) + # Dont' mask out static objects (static in all frames) instance_token = nusc.get("sample_annotation", box.token)["instance_token"] if not instances_is_dynamic[instance_token]: continue - # project box to image plane and rasterize each face + # Project box to image plane and rasterize each face corners_3d = box.corners() corners = view_points(corners_3d, intrinsics, normalize=True)[:2, :] corners = np.round(corners).astype(int).T + + # Type ignores needed because fillPoly expects cv2.Mat cv2.fillPoly(mask, [corners[[0, 1, 2, 3]]], 0) # front # type: ignore cv2.fillPoly(mask, [corners[[4, 5, 6, 7]]], 0) # back # type: ignore cv2.fillPoly(mask, [corners[[0, 1, 5, 4]]], 0) # top # type: ignore @@ -117,27 +119,30 @@ class ProcessNuScenesMasks: cv2.fillPoly(mask, [corners[[1, 2, 6, 5]]], 0) # right # type: ignore maskname = os.path.split(camera_data["filename"])[1].replace("jpg", "png") - cv2.imwrite(str(self.output_dir / "masks" / camera / maskname), mask * 255) # type: ignore + cv2.imwrite( + str(self.output_dir / "masks" / camera / maskname), + mask * 255, # type: ignore + ) if self.verbose: - img = cv2.imread(str(self.data_dir / camera_data["filename"])) # type: ignore + img = cv2.imread(str(self.data_dir / camera_data["filename"])) mask = ~mask.astype(bool) - img[mask, :] -= np.minimum(img[mask, :], 100) + img[mask, :] = img[mask, :] - np.minimum(img[mask, :], 100) viz.append(img) if self.verbose: if len(viz) == 6: viz = np.vstack((np.hstack(viz[:3]), np.hstack(viz[3:]))) - viz = cv2.resize(viz, (int(1600 * 3 / 3), int(900 * 2 / 3))) # type: ignore + viz = cv2.resize(viz, (int(1600 * 3 / 3), int(900 * 2 / 3))) elif len(viz) == 3: viz = np.hstack(viz[:3]) - viz = cv2.resize(viz, (int(1600 * 3 / 3), int(900 / 3))) # type: ignore + viz = cv2.resize(viz, (int(1600 * 3 / 3), int(900 / 3))) elif len(viz) == 1: viz = viz[0] else: raise ValueError("Only support 1 or 3 or 6 cameras for viz") - cv2.imshow("", viz) # type: ignore - cv2.waitKey(1) # type: ignore + cv2.imshow("", viz) + cv2.waitKey(1) def entrypoint(): diff --git a/nerfstudio/scripts/github/run_actions.py b/nerfstudio/scripts/github/run_actions.py index 0f8618921..1081cee57 100644 --- a/nerfstudio/scripts/github/run_actions.py +++ b/nerfstudio/scripts/github/run_actions.py @@ -23,7 +23,7 @@ from rich.style import Style from nerfstudio.utils.rich_utils import CONSOLE -LOCAL_TESTS = ["Run license checks", "Run Ruff", "Run Black", "Test with pytest"] +LOCAL_TESTS = ["Run license checks", "Run Ruff", "Run Black", "Run Pyright", "Test with pytest"] def run_command(command: str, continue_on_fail: bool = False) -> bool: diff --git a/nerfstudio/scripts/render.py b/nerfstudio/scripts/render.py index 5ba8371f3..072ff0888 100644 --- a/nerfstudio/scripts/render.py +++ b/nerfstudio/scripts/render.py @@ -51,6 +51,7 @@ from nerfstudio.cameras.camera_paths import ( get_spiral_path, ) from nerfstudio.cameras.cameras import Cameras, CameraType +from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager from nerfstudio.data.scene_box import SceneBox from nerfstudio.model_components import renderers from nerfstudio.pipelines.base_pipeline import Pipeline @@ -416,6 +417,7 @@ class SpiralRender(BaseRender): install_checks.check_ffmpeg_installed() + assert isinstance(pipeline.datamanager, VanillaDataManager) steps = int(self.frame_rate * self.seconds) camera_start = pipeline.datamanager.eval_dataloader.get_camera(image_idx=0).flatten() camera_path = get_spiral_path(camera_start, steps=steps, radius=self.radius) diff --git a/pyproject.toml b/pyproject.toml index 3ae136e56..e4647ed93 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -79,6 +79,12 @@ dev = [ "typeguard==2.13.3", "ruff==0.0.267", "sshconf==0.2.5", + "pycolmap==0.3.0", + "hloc@git+https://github.com/cvg/Hierarchical-Localization.git", + "diffusers==0.16.1", + "opencv-stubs==0.0.7", + "transformers==4.29.2", + "pyright==1.1.308", ] # Documentation related packages @@ -139,7 +145,7 @@ exclude = ["**/node_modules", ] defineConstant = { DEBUG = true } -reportMissingImports = true +reportMissingImports = "warning" reportMissingTypeStubs = false reportPrivateImportUsage = false diff --git a/tests/field_components/test_temporal_distortions.py b/tests/field_components/test_temporal_distortions.py index 548a5fdcf..8d5abd191 100644 --- a/tests/field_components/test_temporal_distortions.py +++ b/tests/field_components/test_temporal_distortions.py @@ -15,9 +15,10 @@ def test_dnerf_distortion(): num_rays = 1024 num_samples = 256 positions = torch.rand((num_rays, num_samples, 3), dtype=torch.float32, device=device) - assert distortion.forward(positions, None) is None times = torch.rand_like(positions[..., :1]) - assert distortion.forward(positions, times).shape == positions.shape + out = distortion.forward(positions, times) + assert out is not None + assert out.shape == positions.shape if __name__ == "__main__": diff --git a/tests/field_components/test_temporal_grid.py b/tests/field_components/test_temporal_grid.py index e539b3030..137f4fc2d 100644 --- a/tests/field_components/test_temporal_grid.py +++ b/tests/field_components/test_temporal_grid.py @@ -34,6 +34,7 @@ def test_temporal_grid(): weight = torch.randn_like(out) (out * weight).sum().backward() assert torch.all(out == 0.5) + assert model.embeddings.grad is not None assert model.embeddings.grad.sum() - weight.sum() < 0.01 assert torch.all(model.embeddings.grad[2:, :] == 0) assert torch.all(model.embeddings.grad[:, 1:] == 0) diff --git a/tests/pipelines/test_vanilla_pipeline.py b/tests/pipelines/test_vanilla_pipeline.py index 2f6f6e4a9..2218294b7 100644 --- a/tests/pipelines/test_vanilla_pipeline.py +++ b/tests/pipelines/test_vanilla_pipeline.py @@ -75,7 +75,7 @@ def test_load_state_dict(): state_dict["_model.param"].mul_(2) pipeline.load_pipeline(state_dict, 0) assert was_called - assert pipeline.model.param[0].item() == 2 + assert getattr(pipeline.model, "param")[0].item() == 2 # preparation for another test state_dict["_model.param"].mul_(2) @@ -91,4 +91,4 @@ def test_load_state_dict(): # load DDP checkpoint pipeline.load_pipeline(ddp_state_dict, 0) assert was_called - assert pipeline.model.param[0].item() == 4 + assert getattr(pipeline.model, "param")[0].item() == 4 diff --git a/tests/test_train.py b/tests/test_train.py index e1b977fd3..a559a0217 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -12,6 +12,7 @@ from nerfstudio.configs.method_configs import method_configs from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig from nerfstudio.data.dataparsers.minimal_dataparser import MinimalDataParserConfig from nerfstudio.engine.trainer import TrainerConfig +from nerfstudio.models.vanilla_nerf import VanillaModelConfig from nerfstudio.scripts.train import train_loop BLACKLIST = [ @@ -46,8 +47,10 @@ def set_reduced_config(config: TrainerConfig): # reduce model factors if hasattr(config.pipeline.model, "num_coarse_samples"): + assert isinstance(config.pipeline.model, VanillaModelConfig) config.pipeline.model.num_coarse_samples = 4 if hasattr(config.pipeline.model, "num_importance_samples"): + assert isinstance(config.pipeline.model, VanillaModelConfig) config.pipeline.model.num_importance_samples = 4 # remove viewer config.viewer.quit_on_train_completion = True diff --git a/tests/utils/test_tensor_dataclass.py b/tests/utils/test_tensor_dataclass.py index 05c2eb43a..8971c83f8 100644 --- a/tests/utils/test_tensor_dataclass.py +++ b/tests/utils/test_tensor_dataclass.py @@ -1,8 +1,8 @@ """ Test tensor dataclass """ -from dataclasses import dataclass -from typing import Dict +from dataclasses import dataclass, field +from typing import Generic, Optional, TypeVar import pytest import torch @@ -17,14 +17,17 @@ class DummyNestedClass(TensorDataclass): x: torch.Tensor +MaybeTensorDataclass = TypeVar("MaybeTensorDataclass") + + @dataclass -class DummyTensorDataclass(TensorDataclass): +class DummyTensorDataclass(TensorDataclass, Generic[MaybeTensorDataclass]): """Dummy dataclass""" a: torch.Tensor b: torch.Tensor - c: DummyNestedClass = None - d: Dict = None + c: MaybeTensorDataclass + d: dict = field(default_factory=dict) def test_init(): @@ -34,7 +37,7 @@ def test_init(): class Dummy(TensorDataclass): """Dummy dataclass""" - dummy_vals: torch.Tensor = None + dummy_vals: Optional[torch.Tensor] = None Dummy(dummy_vals=torch.ones(1)) with pytest.raises(ValueError): @@ -46,19 +49,19 @@ def test_broadcasting(): a = torch.ones((4, 6, 3)) b = torch.ones((6, 2)) - tensor_dataclass = DummyTensorDataclass(a=a, b=b) + tensor_dataclass = DummyTensorDataclass(a=a, b=b, c=None) assert tensor_dataclass.b.shape == (4, 6, 2) a = torch.ones((4, 6, 3)) b = torch.ones(2) - tensor_dataclass = DummyTensorDataclass(a=a, b=b) + tensor_dataclass = DummyTensorDataclass(a=a, b=b, c=None) assert tensor_dataclass.b.shape == (4, 6, 2) # Invalid broadcasting a = torch.ones((4, 6, 3)) b = torch.ones((3, 2)) with pytest.raises(RuntimeError): - tensor_dataclass = DummyTensorDataclass(a=a, b=b) + tensor_dataclass = DummyTensorDataclass(a=a, b=b, c=None) def test_tensor_ops(): @@ -67,7 +70,7 @@ def test_tensor_ops(): a = torch.ones((4, 6, 3)) b = torch.ones((6, 2)) d = {"t1": torch.ones((4, 6, 3)), "t2": {"t3": torch.ones((6, 4))}} - tensor_dataclass = DummyTensorDataclass(a=a, b=b, d=d) + tensor_dataclass = DummyTensorDataclass(a=a, b=b, c=None, d=d) assert tensor_dataclass.shape == (4, 6) assert tensor_dataclass.a.shape == (4, 6, 3) @@ -114,6 +117,7 @@ def test_tensor_ops(): tensor_dataclass = DummyTensorDataclass( a=torch.ones((2, 3, 4, 5)), b=torch.ones((4, 5)), + c=None, d={"t1": torch.ones((2, 3, 4, 5)), "t2": {"t3": torch.ones((4, 5))}}, ) assert tensor_dataclass[0, ...].shape == (3, 4) @@ -168,7 +172,7 @@ def test_nested_class(): def test_iter(): """Test iterating over tensor dataclass""" - tensor_dataclass = DummyTensorDataclass(a=torch.ones((3, 4, 5)), b=torch.ones((3, 4, 5))) + tensor_dataclass = DummyTensorDataclass(a=torch.ones((3, 4, 5)), b=torch.ones((3, 4, 5)), c=None) for batch in tensor_dataclass: assert batch.shape == (4,) assert batch.a.shape == (4, 5)