From 58212543f5c45cb4555560fa0d4d6d1366c43448 Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Fri, 29 Jul 2022 14:17:38 -0700 Subject: [PATCH] Fix typing errors (#229) --- configs/graph_default.yaml | 2 + environment/jupyter.py | 3 - nerfactory/cameras/camera_paths.py | 11 +-- nerfactory/cameras/cameras.py | 3 +- nerfactory/cameras/rays.py | 7 +- nerfactory/cameras/utils.py | 4 +- nerfactory/cuda/backend.py | 1 + nerfactory/data/colmap_utils.py | 5 +- nerfactory/data/dataloader.py | 13 +++- nerfactory/data/datasets.py | 2 +- nerfactory/data/image_dataset.py | 10 +-- nerfactory/data/image_sampler.py | 23 ++++-- nerfactory/data/pixel_sampler.py | 5 +- nerfactory/data/structs.py | 28 +++---- nerfactory/engine/trainer.py | 75 +++++++++++-------- nerfactory/fields/base.py | 5 +- nerfactory/fields/density_fields/__init__.py | 13 ++++ .../fields/density_fields/density_grid.py | 8 +- nerfactory/fields/instant_ngp_field.py | 11 ++- nerfactory/fields/modules/base.py | 3 +- nerfactory/fields/modules/encoding.py | 8 ++ nerfactory/fields/modules/mlp.py | 9 ++- nerfactory/graphs/base.py | 26 ++++--- nerfactory/optimizers/optimizers.py | 14 ++-- nerfactory/renderers/renderers.py | 26 ++++--- nerfactory/utils/config.py | 2 +- nerfactory/utils/misc.py | 18 +++-- nerfactory/utils/writer.py | 4 +- pyproject.toml | 18 ++++- tests/viewer/test_viewer.py | 26 ++++--- 30 files changed, 236 insertions(+), 147 deletions(-) delete mode 100644 environment/jupyter.py create mode 100644 nerfactory/fields/density_fields/__init__.py diff --git a/configs/graph_default.yaml b/configs/graph_default.yaml index 8cb739628..96c026198 100644 --- a/configs/graph_default.yaml +++ b/configs/graph_default.yaml @@ -109,4 +109,6 @@ optimizers: viewer: enable: False zmq_url: tcp://127.0.0.1:6000 + min_render_image_height: 64 + max_render_image_height: 1024 num_rays_per_chunk: 4096 diff --git a/environment/jupyter.py b/environment/jupyter.py deleted file mode 100644 index 70dd75764..000000000 --- a/environment/jupyter.py +++ /dev/null @@ -1,3 +0,0 @@ -## https://github.com/jupyterlab/jupyterlab/issues/2049 -## Allow access to hidden files -c.ContentsManager.allow_hidden = True diff --git a/nerfactory/cameras/camera_paths.py b/nerfactory/cameras/camera_paths.py index 8c3a3ebae..77f6eb9ea 100644 --- a/nerfactory/cameras/camera_paths.py +++ b/nerfactory/cameras/camera_paths.py @@ -18,13 +18,14 @@ Code for camera paths. import copy from dataclasses import dataclass -from typing import List, Tuple +from typing import List, Optional, Tuple import torch import nerfactory.cameras.utils as camera_utils from nerfactory.cameras.cameras import ( Camera, + PinholeCamera, get_camera, get_intrinsics_from_intrinsics_matrix, ) @@ -38,7 +39,7 @@ class CameraPath: cameras: List[Camera] -def get_interpolated_camera_path(camera_a: Camera, camera_b: Camera, steps: int) -> CameraPath: +def get_interpolated_camera_path(camera_a: PinholeCamera, camera_b: PinholeCamera, steps: int) -> CameraPath: """Generate a camera path between two cameras. Args: @@ -72,10 +73,10 @@ def get_interpolated_camera_path(camera_a: Camera, camera_b: Camera, steps: int) def get_spiral_path( - camera: Camera, + camera: PinholeCamera, steps: int = 30, - radius: float = None, - radiuses: Tuple[float] = None, + radius: Optional[float] = None, + radiuses: Optional[Tuple[float]] = None, rots: int = 2, zrate: float = 0.5, ) -> CameraPath: diff --git a/nerfactory/cameras/cameras.py b/nerfactory/cameras/cameras.py index f1cf9f9c7..10e0b9a25 100644 --- a/nerfactory/cameras/cameras.py +++ b/nerfactory/cameras/cameras.py @@ -25,7 +25,6 @@ from torch.nn.functional import normalize from torchtyping import TensorType from nerfactory.cameras.rays import RayBundle -from nerfactory.utils.misc import is_not_none class Camera: @@ -157,7 +156,7 @@ class Camera: camera_to_world = self.camera_to_world.unsqueeze(0).repeat(height, width, 1, 1).to(device) # (num_rays, 3, 4) coords = self.get_image_coords().to(device) ray_bundle = self.generate_rays(intrinsics, camera_to_world, coords) - if is_not_none(self.camera_index): + if self.camera_index is not None: ray_bundle.set_camera_indices(camera_index=self.camera_index) return ray_bundle diff --git a/nerfactory/cameras/rays.py b/nerfactory/cameras/rays.py index a1a9b7a5b..b5d221fcf 100644 --- a/nerfactory/cameras/rays.py +++ b/nerfactory/cameras/rays.py @@ -23,7 +23,6 @@ import torch from torchtyping import TensorType from nerfactory.utils.math import Gaussians, conical_frustum_to_gaussian -from nerfactory.utils.misc import is_not_none from nerfactory.utils.tensor_dataclass import TensorDataclass @@ -143,7 +142,7 @@ class RaySamples(TensorDataclass): Returns: RaySamples: New set of masked samples. """ - if is_not_none(self.valid_mask): + if self.valid_mask is not None: return self[self.valid_mask[..., 0]] return self @@ -169,7 +168,7 @@ class RayBundle(TensorDataclass): nears: Optional[TensorType["num_rays", 1]] = None fars: Optional[TensorType["num_rays", 1]] = None valid_mask: Optional[TensorType["num_rays", 1, bool]] = None - num_rays_per_chunk: int = None + num_rays_per_chunk: Optional[int] = None def set_camera_indices(self, camera_index: int) -> None: """Sets all of the the camera indices to a specific camera index. @@ -229,7 +228,7 @@ class RayBundle(TensorDataclass): dists = bin_ends - bin_starts # [..., N_samples, 1] deltas = dists * torch.norm(self.directions[:, :], dim=-1)[..., None, None] - if is_not_none(self.camera_indices): + if self.camera_indices is not None: camera_indices = self.camera_indices[..., None] else: camera_indices = None diff --git a/nerfactory/cameras/utils.py b/nerfactory/cameras/utils.py index 1a65aecfe..5858fb6b8 100644 --- a/nerfactory/cameras/utils.py +++ b/nerfactory/cameras/utils.py @@ -143,7 +143,7 @@ def quaternion_from_matrix(matrix, isprecise=False): K /= 3.0 # quaternion is eigenvector of K that corresponds to largest eigenvalue w, V = np.linalg.eigh(K) - q = V[[3, 0, 1, 2], np.argmax(w)] + q = V[np.array([3, 0, 1, 2]), np.argmax(w)] if q[0] < 0.0: np.negative(q, q) return q @@ -167,6 +167,8 @@ def quaternion_slerp(quat0, quat1, fraction, spin=0, shortestpath=True): """ q0 = unit_vector(quat0[:4]) q1 = unit_vector(quat1[:4]) + if q0 is None or q1 is None: + raise ValueError("Input quaternions invalid.") if fraction == 0.0: return q0 elif fraction == 1.0: diff --git a/nerfactory/cuda/backend.py b/nerfactory/cuda/backend.py index a11f6b59b..283523e26 100644 --- a/nerfactory/cuda/backend.py +++ b/nerfactory/cuda/backend.py @@ -11,6 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +"""Setup cuda backend.""" import glob import os from subprocess import DEVNULL, call diff --git a/nerfactory/data/colmap_utils.py b/nerfactory/data/colmap_utils.py index d85872c1c..deb3b510e 100644 --- a/nerfactory/data/colmap_utils.py +++ b/nerfactory/data/colmap_utils.py @@ -427,8 +427,7 @@ def read_model(path, ext=""): elif detect_model_format(path, ".txt"): ext = ".txt" else: - print("Provide model format: '.bin' or '.txt'") - return + raise ValueError("Provide model format: '.bin' or '.txt'") if ext == ".txt": cameras = read_cameras_text(os.path.join(path, "cameras" + ext)) @@ -489,7 +488,7 @@ def rotmat2qvec(R): / 3.0 ) eigvals, eigvecs = np.linalg.eigh(K) - qvec = eigvecs[[3, 0, 1, 2], np.argmax(eigvals)] + qvec = eigvecs[np.array([3, 0, 1, 2]), np.argmax(eigvals)] if qvec[0] < 0: qvec *= -1 return qvec diff --git a/nerfactory/data/dataloader.py b/nerfactory/data/dataloader.py index d08a81817..f80285c90 100644 --- a/nerfactory/data/dataloader.py +++ b/nerfactory/data/dataloader.py @@ -59,10 +59,10 @@ def setup_dataset_train(config: DataConfig, device: str) -> Tuple[DatasetInputs, image_dataset_train = instantiate_from_dict_config(config.image_dataset_train, **dataset_inputs_train.as_dict()) # ImageSampler image_sampler_train = instantiate_from_dict_config( - config.dataloader_train.image_sampler, dataset=image_dataset_train, device=device + config.dataloader_train.image_sampler, dataset=image_dataset_train, device=device # type: ignore ) # PixelSampler - pixel_sampler_train = instantiate_from_dict_config(config.dataloader_train.pixel_sampler) + pixel_sampler_train = instantiate_from_dict_config(config.dataloader_train.pixel_sampler) # type: ignore # Dataloader dataloader_train = TrainDataloader(image_sampler_train, pixel_sampler_train) return dataset_inputs_train, dataloader_train @@ -255,6 +255,15 @@ class EvalDataloader: # pylint: disable=too-few-public-methods self.device = device self.kwargs = kwargs + @abstractmethod + def __iter__(self): + """Iterates over the dataset""" + return self + + @abstractmethod + def __next__(self) -> Tuple[RayBundle, Dict]: + """Returns the next batch of data""" + def get_camera(self, image_idx) -> Camera: """Get camera for the given image index""" intrinsics = self.intrinsics[image_idx].to(self.device) diff --git a/nerfactory/data/datasets.py b/nerfactory/data/datasets.py index 63cc126f4..96d884b43 100644 --- a/nerfactory/data/datasets.py +++ b/nerfactory/data/datasets.py @@ -462,7 +462,7 @@ class Friends(Dataset): camera_to_world[..., 3] = (camera_to_world[..., 3] - box_center) * box_scale_factor # center and rescale # --- semantics --- - semantics = Semantics() + semantics = None if self.include_semantics: thing_filenames = [ image_filename.replace("/images/", "/segmentations/thing/").replace(".jpg", ".png") diff --git a/nerfactory/data/image_dataset.py b/nerfactory/data/image_dataset.py index d58b5425c..da2792229 100644 --- a/nerfactory/data/image_dataset.py +++ b/nerfactory/data/image_dataset.py @@ -21,16 +21,16 @@ from typing import Dict, List, Optional, Union import numpy as np import numpy.typing as npt -import PIL import torch from PIL import Image +from torch.utils.data import Dataset from torchtyping import TensorType from nerfactory.data.structs import Semantics from nerfactory.utils.misc import is_not_none -class ImageDataset(torch.utils.data.Dataset): +class ImageDataset(Dataset): """Dataset that returns images.""" def __init__( @@ -79,7 +79,7 @@ class ImageDataset(torch.utils.data.Dataset): f"Image height {image_height} is not divisible by downscale_factor {self.downscale_factor}" ) pil_image = pil_image.resize( - (image_width // self.downscale_factor, image_height // self.downscale_factor), PIL.Image.BILINEAR + (image_width // self.downscale_factor, image_height // self.downscale_factor), Image.BILINEAR ) image = np.array(pil_image, dtype="uint8") # shape is (h, w, 3 or 4) assert len(image.shape) == 3 @@ -115,10 +115,10 @@ class ImageDataset(torch.utils.data.Dataset): data = {"image_idx": image_idx} assert is_not_none(image) data["image"] = image - if is_not_none(mask): + if mask is not None: assert mask.shape[:2] == image.shape[:2] data["mask"] = mask - if is_not_none(semantics): + if semantics is not None: assert semantics.shape[:2] == image.shape[:2] data["semantics"] = semantics return data diff --git a/nerfactory/data/image_sampler.py b/nerfactory/data/image_sampler.py index a271b42c1..be97b7669 100644 --- a/nerfactory/data/image_sampler.py +++ b/nerfactory/data/image_sampler.py @@ -21,13 +21,23 @@ from abc import abstractmethod from typing import Union import torch -from torch.utils.data import default_collate +from torch.utils.data import IterableDataset, default_collate +from nerfactory.data.image_dataset import ImageDataset from nerfactory.utils.misc import get_dict_to_torch -class ImageSampler(torch.utils.data.IterableDataset): - """Samples image_batch's.""" +class ImageSampler(IterableDataset): + """Samples image_batch's. + + Attributes: + dataset (torch.utils.data.Dataset): Dataset to sample from. + """ + + # dataset: ImageDataset + def __init__(self, dataset: ImageDataset): + super().__init__() + self.dataset = dataset @abstractmethod def __getitem__(self, idx): @@ -51,14 +61,14 @@ class CacheImageSampler(ImageSampler): def __init__( self, - dataset: torch.utils.data.Dataset, + dataset: ImageDataset, num_images_to_sample_from: int = -1, num_times_to_repeat_images: int = 0, device: Union[torch.device, str] = "cpu", ): - super().__init__() + super().__init__(dataset=dataset) self.dataset = dataset - self.cache_all_images = num_images_to_sample_from is -1 + self.cache_all_images = num_images_to_sample_from == -1 self.num_images_to_sample_from = len(self.dataset) if self.cache_all_images else num_images_to_sample_from self.num_times_to_repeat_images = num_times_to_repeat_images self.device = device @@ -76,6 +86,7 @@ class CacheImageSampler(ImageSampler): def _get_batch_list(self): """Returns a list of batches from the dataset attribute.""" indices = random.sample(range(len(self.dataset)), k=self.num_images_to_sample_from) + batch_list = [self.dataset.__getitem__(idx) for idx in indices] return batch_list diff --git a/nerfactory/data/pixel_sampler.py b/nerfactory/data/pixel_sampler.py index 1ee90d64b..9b7390b0b 100644 --- a/nerfactory/data/pixel_sampler.py +++ b/nerfactory/data/pixel_sampler.py @@ -43,6 +43,7 @@ def collate_image_dataset_batch(batch, num_rays_per_batch, keep_full_image: bool c, y, x = [i.flatten() for i in torch.split(indices, 1, dim=-1)] image = batch["image"][c, y, x] + mask, semantics = None, None if "mask" in batch: mask = batch["mask"][c, y, x] if "semantics" in batch: @@ -57,9 +58,9 @@ def collate_image_dataset_batch(batch, num_rays_per_batch, keep_full_image: bool "indices": indices, # with the abs camera indices "image": image, } - if "mask" in batch: + if mask is not None: collated_batch["mask"] = mask - if "semantics" in batch: + if semantics is not None: collated_batch["semantics"] = semantics if keep_full_image: diff --git a/nerfactory/data/structs.py b/nerfactory/data/structs.py index 733bd5343..5c73b31d7 100644 --- a/nerfactory/data/structs.py +++ b/nerfactory/data/structs.py @@ -17,7 +17,7 @@ Dataset input structures. """ from dataclasses import dataclass -from typing import Dict, List, Optional +from typing import Dict, List, Optional, Union import torch from torchtyping import TensorType @@ -37,12 +37,12 @@ class PointCloud: class Semantics: """Dataclass for semantic labels.""" - stuff_classes: Optional[List[str]] = None - stuff_colors: Optional[List[str]] = None - stuff_filenames: Optional[List[str]] = None - thing_classes: Optional[List[str]] = None - thing_colors: Optional[List[str]] = None - thing_filenames: Optional[List[str]] = None + stuff_filenames: List[str] + stuff_classes: List[str] + thing_filenames: List[str] + thing_classes: List[str] + stuff_colors: Optional[torch.Tensor] = None + thing_colors: Optional[torch.Tensor] = None @dataclass @@ -68,7 +68,7 @@ class SceneBounds: diff = self.aabb[1] - self.aabb[0] return self.aabb[0] + diff / 2.0 - def get_centered_and_scaled_scene_bounds(self, scale_factor=1.0): + def get_centered_and_scaled_scene_bounds(self, scale_factor: Union[float, torch.Tensor] = 1.0): """Returns a new box that has been shifted and rescaled to be centered about the origin.""" return SceneBounds(aabb=(self.aabb - self.get_center()) * scale_factor) @@ -110,13 +110,13 @@ class DatasetInputs: """ image_filenames: List[str] + intrinsics: torch.Tensor + camera_to_world: torch.Tensor downscale_factor: int = 1 - intrinsics: torch.tensor = None - camera_to_world: torch.tensor = None - mask_filenames: List[str] = None - depth_filenames: List[str] = None + mask_filenames: Optional[List[str]] = None + depth_filenames: Optional[List[str]] = None scene_bounds: SceneBounds = SceneBounds() - semantics: Semantics = Semantics() + semantics: Optional[Semantics] = None point_cloud: PointCloud = PointCloud() alpha_color: Optional[TensorType[3]] = None @@ -151,5 +151,5 @@ class BaseDataContainer: """ rays: RayBundle # Raybundle and the cameras will be merged into one thing in a later PR - cameras: torch.tensor + cameras: torch.Tensor ground_truth_pixels: Optional[TensorType["num_pixels", 3]] = None diff --git a/nerfactory/engine/trainer.py b/nerfactory/engine/trainer.py index ca78da24d..e11c58e43 100644 --- a/nerfactory/engine/trainer.py +++ b/nerfactory/engine/trainer.py @@ -18,23 +18,26 @@ Code to train model. import functools import logging import os -from typing import Dict +from typing import Dict, List, Optional import torch import torch.distributed as dist -from torch.cuda.amp import GradScaler +from torch.cuda.amp.grad_scaler import GradScaler from torch.nn.parallel import DistributedDataParallel as DDP from torchtyping import TensorType from nerfactory.cameras.rays import RayBundle from nerfactory.data.dataloader import ( EvalDataloader, + TrainDataloader, setup_dataset_eval, setup_dataset_train, ) -from nerfactory.graphs.base import setup_graph -from nerfactory.optimizers.optimizers import setup_optimizers +from nerfactory.data.structs import DatasetInputs +from nerfactory.graphs.base import Graph, setup_graph +from nerfactory.optimizers.optimizers import Optimizers, setup_optimizers from nerfactory.utils import profiler, writer +from nerfactory.utils.callbacks import Callback from nerfactory.utils.config import Config from nerfactory.utils.decorators import check_main_thread from nerfactory.utils.writer import EventName, TimeWriter @@ -62,12 +65,12 @@ class Trainer: self.mixed_precision = False logging.warning("Mixed precision is disabled for CPU training.") # dataset variables - self.dataset_inputs_train = None - self.dataloader_train = None - self.dataloader_eval = None + self.dataset_inputs_train: DatasetInputs + self.dataloader_train: TrainDataloader + self.dataloader_eval: EvalDataloader # model variables - self.graph = None - self.optimizers = None + self.graph = Graph + self.optimizers: Optimizers self.start_step = 0 # logging variables writer.setup_event_writers(config.logging, max_iter=config.trainer.max_num_iterations) @@ -76,7 +79,7 @@ class Trainer: self.visualizer_state = viewer_utils.VisualizerState(config.viewer) self.grad_scaler = GradScaler(enabled=self.mixed_precision) # training callbacks - self.callbacks = None + self.callbacks: List[Callback] def setup(self, test_mode=False): """Setup the Trainer by calling other setup functions. @@ -89,14 +92,13 @@ class Trainer: self.graph = setup_graph(self.config.graph, self.dataset_inputs_train, device=self.device) self.optimizers = setup_optimizers(self.config.optimizers, self.graph.get_param_groups()) - if self.config.trainer.resume_train.load_dir: - self._load_checkpoint() + self._load_checkpoint() if self.world_size > 1: self.graph = DDP(self.graph, device_ids=[self.local_rank]) dist.barrier(device_ids=[self.local_rank]) - self.callbacks = self.graph.get_training_callbacks() + self.callbacks = self.graph.get_training_callbacks() # type: ignore @classmethod def get_aggregated_loss(cls, loss_dict: Dict[str, torch.Tensor]) -> torch.Tensor: @@ -131,7 +133,7 @@ class Trainer: if step % self.config.trainer.steps_per_test == 0: self.eval_with_dataloader(self.dataloader_eval, step=step) self._write_out_storage(step) - self.visualizer_state.update_scene(step, self.graph) + self.visualizer_state.update_scene(step, self.graph) # type: ignore self._write_out_storage(num_iterations) @@ -152,15 +154,18 @@ class Trainer: def _load_checkpoint(self) -> None: """Helper function to load graph and optimizer from prespecified checkpoint""" load_config = self.config.trainer.resume_train - load_path = os.path.join(load_config.load_dir, f"step-{load_config.load_step:09d}.ckpt") - assert os.path.exists(load_path), f"Checkpoint {load_path} does not exist" - loaded_state = torch.load(load_path, map_location="cpu") - self.start_step = loaded_state["step"] + 1 - # load the checkpoints for graph and optimizer - self.graph.load_graph(loaded_state) - self.optimizers.load_optimizers(loaded_state) - self.grad_scaler.load_state_dict(loaded_state["scaler"]) - logging.info("done loading checkpoint from %s", load_path) + if load_config.load_dir is not None and load_config.load_step is not None: + load_path = os.path.join(load_config.load_dir, f"step-{load_config.load_step:09d}.ckpt") + assert os.path.exists(load_path), f"Checkpoint {load_path} does not exist" + loaded_state = torch.load(load_path, map_location="cpu") + self.start_step = loaded_state["step"] + 1 + # load the checkpoints for graph and optimizer + self.graph.load_graph(loaded_state) # type: ignore + self.optimizers.load_optimizers(loaded_state) + self.grad_scaler.load_state_dict(loaded_state["scaler"]) + logging.info("done loading checkpoint from %s", load_path) + else: + logging.info("No checkpoints to load, training from scratch") @check_main_thread def _save_checkpoint(self, output_dir: str, step: int) -> None: @@ -173,10 +178,14 @@ class Trainer: if not os.path.exists(output_dir): os.makedirs(output_dir) ckpt_path = os.path.join(output_dir, f"step-{step:09d}.ckpt") + if hasattr(self.graph, "module"): + model = self.graph.module.state_dict() + else: + model = self.graph.state_dict() torch.save( { "step": step, - "model": self.graph.module.state_dict() if hasattr(self.graph, "module") else self.graph.state_dict(), + "model": model, "optimizers": {k: v.state_dict() for (k, v) in self.optimizers.optimizers.items()}, "scalers": self.grad_scaler.state_dict(), }, @@ -197,9 +206,9 @@ class Trainer: """ self.optimizers.zero_grad_all() with torch.autocast(device_type=ray_indices.device.type, enabled=self.mixed_precision): - _, loss_dict, metrics_dict = self.graph.forward(ray_indices, batch=batch) + _, loss_dict, metrics_dict = self.graph.forward(ray_indices=ray_indices, batch=batch) # type: ignore loss = sum(loss_dict.values()) - self.grad_scaler.scale(loss).backward() + self.grad_scaler.scale(loss).backward() # type: ignore self.optimizers.optimizer_scaler_step_all(self.grad_scaler) self.grad_scaler.update() @@ -213,7 +222,7 @@ class Trainer: return loss_dict @profiler.time_function - def test_image(self, camera_ray_bundle: RayBundle, batch: dict, step: int = None) -> float: + def test_image(self, camera_ray_bundle: RayBundle, batch: dict, step: Optional[int] = None) -> float: """Test a specific image. Args: @@ -224,14 +233,16 @@ class Trainer: Returns: float: PSNR """ - self.graph.eval() + self.graph.eval() # type: ignore + if camera_ray_bundle.camera_indices is None: + raise ValueError("camera_ray_bundle.camera_indices is None during testing") image_idx = int(camera_ray_bundle.camera_indices[0, 0]) - outputs = self.graph.get_outputs_for_camera_ray_bundle(camera_ray_bundle) - psnr = self.graph.log_test_image_outputs(image_idx, step, batch, outputs) - self.graph.train() + outputs = self.graph.get_outputs_for_camera_ray_bundle(camera_ray_bundle) # type: ignore + psnr = self.graph.log_test_image_outputs(image_idx, step, batch, outputs) # type: ignore + self.graph.train() # type: ignore return psnr - def eval_with_dataloader(self, dataloader: EvalDataloader, step: int = None) -> None: + def eval_with_dataloader(self, dataloader: EvalDataloader, step: Optional[int] = None) -> None: """Run evaluation with a given dataloader. Args: diff --git a/nerfactory/fields/base.py b/nerfactory/fields/base.py index 2bf5dc44b..021d49f00 100644 --- a/nerfactory/fields/base.py +++ b/nerfactory/fields/base.py @@ -25,7 +25,6 @@ from torchtyping import TensorType from nerfactory.cameras.rays import Frustums, RaySamples from nerfactory.fields.modules.field_heads import FieldHeadNames -from nerfactory.utils.misc import is_not_none class Field(nn.Module): @@ -79,7 +78,7 @@ class Field(nn.Module): """ valid_mask = ray_samples.valid_mask - if is_not_none(valid_mask): + if valid_mask is not None: # Hacky handling of empty masks. Tests on a single ray but doesn't use results if not valid_mask.any(): ray_samples = RaySamples(frustums=Frustums.get_mock_frustum().to(valid_mask.device)) @@ -105,5 +104,5 @@ class Field(nn.Module): density, density_embedding = self.get_density(ray_samples) field_outputs = self.get_outputs(ray_samples, density_embedding=density_embedding) - field_outputs[FieldHeadNames.DENSITY] = density + field_outputs[FieldHeadNames.DENSITY] = density # type: ignore return field_outputs diff --git a/nerfactory/fields/density_fields/__init__.py b/nerfactory/fields/density_fields/__init__.py new file mode 100644 index 000000000..02655775e --- /dev/null +++ b/nerfactory/fields/density_fields/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2022 The Plenoptix Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/nerfactory/fields/density_fields/density_grid.py b/nerfactory/fields/density_fields/density_grid.py index 72cd53a23..bd9409850 100644 --- a/nerfactory/fields/density_fields/density_grid.py +++ b/nerfactory/fields/density_fields/density_grid.py @@ -51,8 +51,10 @@ class DensityGrid(nn.Module): on the scale but all shares the same center. For example: - - `center=0`, `base_scale=3`, `num_cascades=1` would create one grid (lvl=0) covers the aabb [[-1.5, -1.5, -1.5], [1.5, 1.5, 1.5]]. - - `center=0`, `base_scale=3`, `num_cascades=2` would create one grid (lvl=0) covers the aabb [[-1.5, -1.5, -1.5], [1.5, 1.5, 1.5]], + - `center=0`, `base_scale=3`, `num_cascades=1` would create one grid (lvl=0) covers + the aabb [[-1.5, -1.5, -1.5], [1.5, 1.5, 1.5]]. + - `center=0`, `base_scale=3`, `num_cascades=2` would create one grid (lvl=0) covers + the aabb [[-1.5, -1.5, -1.5], [1.5, 1.5, 1.5]], and one more grid (lvl=1) covers the aabb [[-3, -3, -3], [3, 3, 3]]. Normally you can simply set the `center` to the scene center, and `base_scale` to be the scene scale, and use @@ -225,6 +227,6 @@ class DensityGrid(nn.Module): densities = densities.view(*positions.shape[:-1], 1) return densities - def forward(self, x): + def forward(self, x): # pylint: disable=no-self-use """Not implemented.""" raise RuntimeError("Shouldn't be called!") diff --git a/nerfactory/fields/instant_ngp_field.py b/nerfactory/fields/instant_ngp_field.py index 834acd4bf..891e27979 100644 --- a/nerfactory/fields/instant_ngp_field.py +++ b/nerfactory/fields/instant_ngp_field.py @@ -17,10 +17,11 @@ Instant-NGP field implementations using tiny-cuda-nn, torch, .... """ -from typing import Tuple +from typing import Optional, Tuple import torch from torch.nn.parameter import Parameter +from torchtyping import TensorType from nerfactory.cameras.rays import RaySamples from nerfactory.data.structs import SceneBounds @@ -110,13 +111,17 @@ class TCNNInstantNGPField(Field): density = trunc_exp(density_before_activation.to(positions)) return density, base_mlp_out - def get_outputs(self, ray_samples: RaySamples, density_embedding=None): + def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None): # TODO: add valid_mask masking! # tcnn requires directions in the range [0,1] directions = get_normalized_directions(ray_samples.frustums.directions) directions_flat = directions.view(-1, 3) d = self.direction_encoding(directions_flat) - h = torch.cat([d, density_embedding.view(-1, self.geo_feat_dim)], dim=-1) + if density_embedding is None: + positions = SceneBounds.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) + h = torch.cat([d, positions.view(-1, 3)], dim=-1) + else: + h = torch.cat([d, density_embedding.view(-1, self.geo_feat_dim)], dim=-1) rgb = self.mlp_head(h).view(*ray_samples.frustums.directions.shape[:-1], -1).to(directions) return {FieldHeadNames.RGB: rgb} diff --git a/nerfactory/fields/modules/base.py b/nerfactory/fields/modules/base.py index 5cd14fc77..9f2e6295a 100644 --- a/nerfactory/fields/modules/base.py +++ b/nerfactory/fields/modules/base.py @@ -16,6 +16,7 @@ The field module baseclass. """ from abc import abstractmethod +from typing import Optional from torch import nn from torchtyping import TensorType @@ -29,7 +30,7 @@ class FieldModule(nn.Module): out_dim (int, optional): Ouput dimension to module. """ - def __init__(self, in_dim: int = None, out_dim: int = None) -> None: + def __init__(self, in_dim: Optional[int] = None, out_dim: Optional[int] = None) -> None: super().__init__() self.in_dim = in_dim self.out_dim = out_dim diff --git a/nerfactory/fields/modules/encoding.py b/nerfactory/fields/modules/encoding.py index b7086c0ad..914ea7b34 100644 --- a/nerfactory/fields/modules/encoding.py +++ b/nerfactory/fields/modules/encoding.py @@ -50,6 +50,8 @@ class Identity(Encoding): """Identity encoding (Does not modify input)""" def get_out_dim(self) -> int: + if self.in_dim is None: + 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"]: @@ -72,6 +74,8 @@ class ScalingAndOffset(Encoding): self.offset = offset def get_out_dim(self) -> int: + if self.in_dim is None: + 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"]: @@ -101,6 +105,8 @@ class NeRFEncoding(Encoding): self.include_input = include_input def get_out_dim(self) -> int: + if self.in_dim is None: + raise ValueError("Input dimension has not been set") out_dim = self.in_dim * self.num_frequencies * 2 if self.include_input: out_dim += self.in_dim @@ -156,6 +162,8 @@ class RFFEncoding(Encoding): if not scale > 0: raise ValueError("RFF encoding scale should be greater than zero") self.scale = scale + if self.in_dim is None: + raise ValueError("Input dimension has not been set") b_matrix = torch.normal(mean=0, std=self.scale, size=(self.in_dim, self.num_frequencies)) self.register_buffer(name="b_matrix", tensor=b_matrix) self.include_input = include_input diff --git a/nerfactory/fields/modules/mlp.py b/nerfactory/fields/modules/mlp.py index a55de10b1..3af30d012 100644 --- a/nerfactory/fields/modules/mlp.py +++ b/nerfactory/fields/modules/mlp.py @@ -42,7 +42,7 @@ class MLP(FieldModule): num_layers: int, layer_width: int, out_dim: Optional[int] = None, - skip_connections: Tuple[int] = (), + skip_connections: Optional[Tuple[int]] = None, activation: Optional[nn.Module] = nn.ReLU(), out_activation: Optional[nn.Module] = None, ) -> None: @@ -67,9 +67,10 @@ class MLP(FieldModule): else: for i in range(self.num_layers - 1): if i == 0: - assert i not in list(self.skip_connections), "Skip connection at layer 0 doesn't make sense." + if self.skip_connections is not None: + assert i not in list(self.skip_connections), "Skip connection at layer 0 doesn't make sense." layers.append(nn.Linear(self.in_dim, self.layer_width)) - elif i in self.skip_connections: + elif self.skip_connections is not None and i in self.skip_connections: layers.append(nn.Linear(self.layer_width + self.in_dim, self.layer_width)) else: layers.append(nn.Linear(self.layer_width, self.layer_width)) @@ -87,7 +88,7 @@ class MLP(FieldModule): """ x = in_tensor for i, layer in enumerate(self.layers): - if i in self.skip_connections: + if self.skip_connections is not None and i in self.skip_connections: x = torch.cat([in_tensor, x], -1) x = layer(x) if self.activation: diff --git a/nerfactory/graphs/base.py b/nerfactory/graphs/base.py index 013c09306..91b412071 100644 --- a/nerfactory/graphs/base.py +++ b/nerfactory/graphs/base.py @@ -39,18 +39,6 @@ from nerfactory.utils.misc import ( ) -@profiler.time_function -def setup_graph(config: GraphConfig, dataset_inputs: DatasetInputs, device: str) -> "Graph": - """Setup the graph. The dataset inputs should be set with the training data. - - Args: - dataset_inputs: The inputs which will be used to define the camera parameters. - """ - graph = instantiate_from_dict_config(DictConfig(config), **dataset_inputs.as_dict()) - graph.to(device) - return graph - - class AbstractGraph(nn.Module): """Highest level graph class. Somewhat useful to lift code up and out of the way.""" @@ -243,3 +231,17 @@ class Graph(AbstractGraph): def load_graph(self, loaded_state: Dict[str, Any]) -> None: """Load the checkpoint from the given path""" self.load_state_dict({key.replace("module.", ""): value for key, value in loaded_state["model"].items()}) + + +@profiler.time_function +def setup_graph(config: GraphConfig, dataset_inputs: DatasetInputs, device: str) -> Graph: + """Setup the graph. The dataset inputs should be set with the training data. + + Args: + dataset_inputs: The inputs which will be used to define the camera parameters. + """ + graph = instantiate_from_dict_config(DictConfig(config), **dataset_inputs.as_dict()) + graph.to(device) + if not isinstance(graph, Graph): + raise TypeError(f"Unable to instantiate Graph, got {type(graph)}") + return graph diff --git a/nerfactory/optimizers/optimizers.py b/nerfactory/optimizers/optimizers.py index e63fb3c88..5c466647b 100644 --- a/nerfactory/optimizers/optimizers.py +++ b/nerfactory/optimizers/optimizers.py @@ -16,19 +16,21 @@ Optimizers class. """ -from typing import Any, Dict, List +from typing import Any, Dict, List, Union import numpy as np from hydra.utils import instantiate from omegaconf import DictConfig -from torch.cuda.amp import GradScaler -from torch.nn import Parameter +from torch.cuda.amp.grad_scaler import GradScaler +from torch.nn.parameter import Parameter from torch.optim.lr_scheduler import LambdaLR from nerfactory.utils import writer -def setup_optimizers(config: DictConfig, param_groups: Dict[str, List[Parameter]]) -> "Optimizers": +def setup_optimizers( + config: Union[DictConfig, Dict[str, Any]], param_groups: Dict[str, List[Parameter]] +) -> "Optimizers": """Helper to set up the optimizers Args: @@ -82,7 +84,7 @@ class Optimizers: param_groups: A dictionary of parameter groups to optimize. """ - def __init__(self, config: DictConfig, param_groups: Dict[str, List[Parameter]]): + def __init__(self, config: Union[DictConfig, Dict[str, Any]], param_groups: Dict[str, List[Parameter]]): self.config = config self.optimizers = {} self.schedulers = {} @@ -108,7 +110,7 @@ class Optimizers: Args: param_group_name: name of scheduler to step forward """ - if self.config.param_group_name.scheduler: + if self.config.param_group_name.scheduler: # type: ignore self.schedulers[param_group_name].step() def zero_grad_all(self) -> None: diff --git a/nerfactory/renderers/renderers.py b/nerfactory/renderers/renderers.py index 672c6d6f5..22cffd050 100644 --- a/nerfactory/renderers/renderers.py +++ b/nerfactory/renderers/renderers.py @@ -51,10 +51,10 @@ class RGBRenderer(nn.Module): @classmethod def combine_rgb( cls, - rgb: TensorType[..., "num_samples", 3], - weights: TensorType[..., "num_samples", 1], + rgb: TensorType["bs":..., "num_samples", 3], + weights: TensorType["bs":..., "num_samples", 1], background_color: Optional[TensorType[3]] = None, - ) -> TensorType[..., 3]: + ) -> TensorType["bs":..., 3]: """Composite samples along ray and render color image Args: @@ -76,9 +76,9 @@ class RGBRenderer(nn.Module): def forward( self, - rgb: TensorType[..., "num_samples", 3], - weights: TensorType[..., "num_samples", 1], - ) -> TensorType[..., 3]: + rgb: TensorType["bs":..., "num_samples", 3], + weights: TensorType["bs":..., "num_samples", 1], + ) -> TensorType["bs":..., 3]: """Composite samples along ray and render color image Args: @@ -149,8 +149,8 @@ class AccumulationRenderer(nn.Module): @classmethod def forward( cls, - weights: TensorType[..., "num_samples", 1], - ) -> TensorType: + weights: TensorType["bs":..., "num_samples", 1], + ) -> TensorType["bs":..., 1]: """Composite samples along ray and calculate accumulation. Args: @@ -205,8 +205,8 @@ class UncertaintyRenderer(nn.Module): @classmethod def forward( - cls, betas: TensorType[..., "num_samples", 1], weights: TensorType[..., "num_samples", 1] - ) -> TensorType: + cls, betas: TensorType["bs":..., "num_samples", 1], weights: TensorType["bs":..., "num_samples", 1] + ) -> TensorType["bs":..., 1]: """Calculate uncertainty along the ray. Args: @@ -225,8 +225,10 @@ class SemanticRenderer(nn.Module): @classmethod def forward( - cls, semantics: TensorType[..., "num_samples", "num_classes"], weights: TensorType[..., "num_samples", 1] - ) -> TensorType[..., "num_classes"]: + cls, + semantics: TensorType["bs":..., "num_samples", "num_classes"], + weights: TensorType["bs":..., "num_samples", 1], + ) -> TensorType["bs":..., "num_classes"]: """_summary_""" sem = torch.sum(weights * semantics, dim=-2) return sem diff --git a/nerfactory/utils/config.py b/nerfactory/utils/config.py index d7d2168a1..f71551822 100644 --- a/nerfactory/utils/config.py +++ b/nerfactory/utils/config.py @@ -122,7 +122,7 @@ class Config: data: DataConfig = MISSING graph: GraphConfig = MISSING optimizers: Dict[str, Any] = MISSING - viewer: Dict[str, Any] = MISSING + viewer: ViewerConfig = MISSING # additional optional parameters here hydra: Optional[Dict[str, Any]] = None diff --git a/nerfactory/utils/misc.py b/nerfactory/utils/misc.py index 69f5a1333..4678ae34a 100644 --- a/nerfactory/utils/misc.py +++ b/nerfactory/utils/misc.py @@ -20,7 +20,7 @@ import hashlib import json from math import floor, log from pydoc import locate -from typing import Any, Callable, Dict +from typing import Any, Callable, Dict, Union import torch from omegaconf import DictConfig @@ -34,11 +34,11 @@ class DotDict(dict): def __getattr__(self, attr): return self[attr] - __setattr__ = dict.__setitem__ - __delattr__ = dict.__delitem__ + __setattr__ = dict.__setitem__ # type: ignore + __delattr__ = dict.__delitem__ # type: ignore -def get_dict_to_torch(stuff, device="cpu"): +def get_dict_to_torch(stuff, device: Union[torch.device, str] = "cpu"): """Set everything in the dict to the specified torch device.""" if isinstance(stuff, dict): for k, v in stuff.items(): @@ -74,13 +74,17 @@ def get_masked_dict(d, mask): return masked_dict -def instantiate_from_dict_config(dict_config: DictConfig, **kwargs): - """Our version of hydra's instantiate function.""" +def instantiate_from_dict_config(dict_config: Any, **kwargs): + """Our version of hydra's instantiate function. + + Args: + dict_config: DictConfig object to instantiate. It can be a dataclass or a dict but must have a `_target_` field. + """ dict_config_kwargs = {k: v for k, v in dict_config.items() if k != "_target_"} uninstantiated_class = locate(dict_config._target_) # pylint: disable=protected-access all_kwargs = dict_config_kwargs all_kwargs.update(kwargs) - instantiated_class = uninstantiated_class(**all_kwargs) + instantiated_class = uninstantiated_class(**all_kwargs) # type: ignore return instantiated_class diff --git a/nerfactory/utils/writer.py b/nerfactory/utils/writer.py index 531577a95..38bb4eb77 100644 --- a/nerfactory/utils/writer.py +++ b/nerfactory/utils/writer.py @@ -202,8 +202,8 @@ class TimeWriter: self.step = step self.write = write - self.start = None - self.duration = None + self.start: float = 0.0 + self.duration: float = 0.0 def __enter__(self): self.start = time() diff --git a/pyproject.toml b/pyproject.toml index e9c000bb4..6b63cf77b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,4 +43,20 @@ disable = [ addopts = "--typeguard-packages=nerfactory --torchtyping-patch-typeguard --disable-warnings --ignore=tests/cuda --ignore=tests/viewer" testpaths = [ "tests", -] \ No newline at end of file +] + +# pyright +[tool.pyright] +include = ["nerfactory"] +exclude = ["**/node_modules", + "**/__pycache__", +] +ignore = ["nerfactory/viewer", "nerfactory/cuda"] +defineConstant = { DEBUG = true } + +reportMissingImports = true +reportMissingTypeStubs = false +reportPrivateImportUsage = false + +pythonVersion = "3.8" +pythonPlatform = "Linux" diff --git a/tests/viewer/test_viewer.py b/tests/viewer/test_viewer.py index ed16dd033..0430ccc68 100644 --- a/tests/viewer/test_viewer.py +++ b/tests/viewer/test_viewer.py @@ -12,9 +12,9 @@ import umsgpack from tqdm import tqdm import nerfactory.viewer.server.geometry as g -from nerfactory.cameras.camera_paths import InterpolatedCameraPath +from nerfactory.cameras import camera_paths from nerfactory.cameras.cameras import get_camera -from nerfactory.data.utils import get_dataset_inputs +from nerfactory.data.datasets import Blender from nerfactory.utils.io import get_absolute_path from nerfactory.viewer.server.transformations import ( get_rotation_matrix, @@ -24,6 +24,7 @@ from nerfactory.viewer.server.viewer_utils import get_default_vis, set_camera def test_drawing(): + """Test drawing objects in the viewer.""" vis = get_default_vis() vis.delete() @@ -110,20 +111,21 @@ def test_drawing(): def test_streaming(): - pass + """Test streaming of data to the viewer.""" def test_rendering(): - pass + """Test rendering.""" def test_camera_trajectory(): - print("test_camera_trajectory") + """Test generating camera trajectory.""" vis = get_default_vis() # vis.delete() # sample poses from a dataset - dataset_inputs = get_dataset_inputs(data_directory="data/blender/lego", dataset_format="blender", split="train") + dataset = Blender(data_directory="data/blender/lego") + dataset_inputs = dataset.get_dataset_inputs(split="train") num_cameras = len(dataset_inputs.intrinsics) intrinsics = dataset_inputs.intrinsics camera_to_world = dataset_inputs.camera_to_world @@ -136,11 +138,10 @@ def test_camera_trajectory(): estimated_time = num_steps / fps print("estimated_time:", estimated_time) - camera_path = InterpolatedCameraPath(camera_a, camera_b) - cameras = camera_path.get_path(steps=num_steps) + camera_path = camera_paths.get_interpolated_camera_path(camera_a, camera_b, steps=num_steps) start = time.time() - for camera in cameras: + for camera in camera_path.cameras: set_camera(vis, camera) time.sleep(1 / fps) time_elapsed = time.time() - start @@ -148,8 +149,9 @@ def test_camera_trajectory(): def test_send_image_stream(): + """Test sending image stream.""" vis = get_default_vis() - vis_Background = vis["/Background"] + vis_background = vis["/Background"] video_filename = get_absolute_path("data/instant_ngp/bear/bear.MOV") images = [] @@ -162,16 +164,16 @@ def test_send_image_stream(): ret, image = cap.read() # break - num_images = len(images) fps = 30 for image in tqdm(images): rgba = cv2.cvtColor(image, cv2.COLOR_RGB2RGBA) encoded = umsgpack.Ext(0x12, rgba.tobytes()) - vis_Background.set_image(encoded) + vis_background.set_image(encoded) time.sleep(1 / fps) def test_perspective_camera(): + """Test perspective camera in viewer.""" vis = get_default_vis() vis.set_object(g.Box([0.5, 0.5, 0.5])) camera = g.PerspectiveCamera(fov=90)