diff --git a/README.md b/README.md index 8ac9ccf6c..966ceea5e 100644 --- a/README.md +++ b/README.md @@ -72,9 +72,13 @@ For more complex changes (e.g., running with your own data/setting up a new NeRF ## 1. Installation: Setup the environment +### Prerequisites + +CUDA must be installed on the system. This library has been tested with version 11.3. You can find more information about installing CUDA [here](https://docs.nvidia.com/cuda/cuda-quick-start-guide/index.html) + ### Create environment -We recommend using conda to manage dependencies. Make sure to install [Conda](https://docs.conda.io/en/latest/miniconda.html) before proceeding. +Nerfstudio requires `python >= 3.7`. We recommend using conda to manage dependencies. Make sure to install [Conda](https://docs.conda.io/en/latest/miniconda.html) before proceeding. ```bash conda create --name nerfstudio -y python=3.8 diff --git a/docs/conf.py b/docs/conf.py index 5fcddff41..2071057b9 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -24,10 +24,6 @@ project = "nerfstudio" copyright = "2022, nerfstudio Team" author = "nerfstudio Team" -# The full version, including alpha/beta/rc tags -release = "0.1.5.1" - - # -- General configuration --------------------------------------------------- # Add any Sphinx extension module names here, as strings. They can be diff --git a/docs/quickstart/first_nerf.md b/docs/quickstart/first_nerf.md index 971e2f633..5e9ffc50d 100644 --- a/docs/quickstart/first_nerf.md +++ b/docs/quickstart/first_nerf.md @@ -152,3 +152,23 @@ ns-eval --load-config={PATH_TO_CONFIG} --output-path=output.json ``` We also provide a train and evaluation script that allows you to do benchmarking on the classical Blender dataset (see our [benchmarking workflow](../developer_guides/debugging_tools/benchmarking.md)). + +## Multi-GPU Training + +Here we explain how to use multi-GPU training. This is the command to train the nerfacto model with 4 GPUs. We are using [PyTorch Distributed Data Parallel (DDP)](https://pytorch.org/tutorials/beginner/dist_overview.html), so we scale the learning rate with the number of GPUs used. Plotting will only be done for the first process. Note that you may want to play around with both learning rate and `_num_rays_per_batch` when using DDP. Below is a simple example for how you'd run the vanilla-nerf method with either 1 or 4 GPUs. We don't see much value at the moment for running the fast methods like nerfacto with more than one GPU. + +```python +# 1 GPU (1x LR) +export CUDA_VISIBLE_DEVICES=0 +ns-train vanilla-nerf \ + --machine.num-gpus 1 \ + --vis wandb \ + --optimizers.fields.optimizer.lr 5e-4 + +# 4 GPUs (4x LR) +export CUDA_VISIBLE_DEVICES=0,1,2,3 +ns-train vanilla-nerf \ + --machine.num-gpus 4 \ + --vis wandb \ + --optimizers.fields.optimizer.lr 20e-4 +``` diff --git a/docs/quickstart/installation.md b/docs/quickstart/installation.md index 9df96ab24..50199b95f 100644 --- a/docs/quickstart/installation.md +++ b/docs/quickstart/installation.md @@ -1,8 +1,12 @@ # Installation +## Prerequisites + +CUDA must be installed on the system. This library has been tested with version 11.3. You can find more information about installing CUDA [here](https://docs.nvidia.com/cuda/cuda-quick-start-guide/index.html) + ## Create environment -We recommend using conda to manage dependencies. Make sure to install [Conda](https://docs.conda.io/en/latest/miniconda.html) before preceding. +Nerfstudio requires `python >= 3.7`. We recommend using conda to manage dependencies. Make sure to install [Conda](https://docs.conda.io/en/latest/miniconda.html) before preceding. ```bash conda create --name nerfstudio -y python=3.8 diff --git a/nerfstudio/cameras/camera_optimizers.py b/nerfstudio/cameras/camera_optimizers.py index c773d5e71..891cd272a 100644 --- a/nerfstudio/cameras/camera_optimizers.py +++ b/nerfstudio/cameras/camera_optimizers.py @@ -20,13 +20,13 @@ from __future__ import annotations import functools from dataclasses import dataclass, field -from typing import Literal, Type, Union +from typing import Type, Union import torch import tyro from torch import nn from torchtyping import TensorType -from typing_extensions import assert_never +from typing_extensions import Literal, assert_never from nerfstudio.cameras.lie_groups import exp_map_SE3, exp_map_SO3xR3 from nerfstudio.configs import base_config as cfg diff --git a/nerfstudio/cameras/camera_utils.py b/nerfstudio/cameras/camera_utils.py index 6fd85b358..b0de30a9b 100644 --- a/nerfstudio/cameras/camera_utils.py +++ b/nerfstudio/cameras/camera_utils.py @@ -17,11 +17,12 @@ Camera transformation helper code. """ import math -from typing import List, Literal, Optional, Tuple +from typing import List, Optional, Tuple import numpy as np import torch from torchtyping import TensorType +from typing_extensions import Literal _EPS = np.finfo(float).eps * 4.0 @@ -406,20 +407,21 @@ 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 auto_orient_poses( - poses: TensorType["num_poses":..., 4, 4], method: Literal["pca", "up"] = "up" +def auto_orient_and_center_poses( + poses: TensorType["num_poses":..., 4, 4], method: Literal["pca", "up", "none"] = "up", center_poses: bool = True ) -> TensorType["num_poses":..., 3, 4]: """Orients and centers the poses. We provide two methods for orientation: pca and up. pca: Orient the poses so that the principal component of the points is aligned with the axes. This method works well when all of the cameras are in the same plane. - up: Orient the poses so that the average up vector is aligned with the z axis. This method works well when images are not at arbitrary angles. + Args: poses: The poses to orient. - method: The method to use for orientation. Either "pca" or "up". + method: The method to use for orientation. + center_poses: If True, the poses are centered around the origin. Returns: The oriented poses. @@ -428,16 +430,21 @@ def auto_orient_poses( translation = poses[..., :3, 3] mean_translation = torch.mean(translation, dim=0) - translation = translation - mean_translation + translation_diff = translation - mean_translation + + if center_poses: + translation = mean_translation + else: + translation = torch.zeros_like(mean_translation) if method == "pca": - _, eigvec = torch.linalg.eigh(translation.T @ translation) + _, eigvec = torch.linalg.eigh(translation_diff.T @ translation_diff) eigvec = torch.flip(eigvec, dims=(-1,)) if torch.linalg.det(eigvec) < 0: eigvec[:, 2] = -eigvec[:, 2] - transform = torch.cat([eigvec, eigvec @ -mean_translation[..., None]], dim=-1) + transform = torch.cat([eigvec, eigvec @ -translation[..., None]], dim=-1) oriented_poses = transform @ poses if oriented_poses.mean(axis=0)[2, 1] < 0: @@ -447,7 +454,10 @@ def auto_orient_poses( up = up / torch.linalg.norm(up) rotation = rotation_matrix(up, torch.Tensor([0, 0, 1])) - transform = torch.cat([rotation, rotation @ -mean_translation[..., None]], dim=-1) + transform = torch.cat([rotation, rotation @ -translation[..., None]], dim=-1) oriented_poses = transform @ poses + elif method == "none": + oriented_poses = poses + poses[:, :3, 3] -= translation return oriented_poses diff --git a/nerfstudio/cameras/cameras.py b/nerfstudio/cameras/cameras.py index 66afba2f1..8d66e3ae4 100644 --- a/nerfstudio/cameras/cameras.py +++ b/nerfstudio/cameras/cameras.py @@ -179,7 +179,7 @@ class Cameras: h_w = h_w.to(torch.int64).to(self.device).broadcast_to((self._num_cameras)) assert torch.all(h_w == h_w[0]), "Batched cameras of different types will be allowed in the future." elif h_w is None: - h_w = torch.Tensor(c_x_y.to(torch.int64).to(self.device) * 2).broadcast_to((self._num_cameras)) + h_w = torch.Tensor((c_x_y * 2).to(torch.int64).to(self.device)).broadcast_to((self._num_cameras)) else: raise ValueError("Height must be an int, tensor, or None, received: " + str(type(h_w))) return h_w diff --git a/nerfstudio/configs/base_config.py b/nerfstudio/configs/base_config.py index e2c518d37..3f1aff36a 100644 --- a/nerfstudio/configs/base_config.py +++ b/nerfstudio/configs/base_config.py @@ -21,10 +21,11 @@ from __future__ import annotations from dataclasses import dataclass from datetime import datetime from pathlib import Path -from typing import Any, Dict, List, Literal, Optional, Tuple, Type +from typing import Any, Dict, List, Optional, Tuple, Type import yaml from rich.console import Console +from typing_extensions import Literal from nerfstudio.configs.config_utils import to_immutable_dict diff --git a/nerfstudio/data/datamanagers.py b/nerfstudio/data/datamanagers.py index c0ad8c7a6..c3cf367c3 100644 --- a/nerfstudio/data/datamanagers.py +++ b/nerfstudio/data/datamanagers.py @@ -294,19 +294,12 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method def setup_train(self): """Sets up the data loaders for training""" assert self.train_dataset is not None - if self.world_size > 1: - sampler = DistributedSampler( - self.train_dataset, num_replicas=self.world_size, rank=self.local_rank, shuffle=True, seed=42 - ) - else: - sampler = None self.train_image_dataloader = CacheDataloader( self.train_dataset, num_images_to_sample_from=self.config.train_num_images_to_sample_from, device=self.device, num_workers=self.world_size * 4, pin_memory=True, - sampler=sampler, ) self.iter_train_image_dataloader = iter(self.train_image_dataloader) self.train_pixel_sampler = PixelSampler(self.config.train_num_rays_per_batch) @@ -321,19 +314,12 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method def setup_eval(self): """Sets up the data loader for evaluation""" assert self.eval_dataset is not None - if self.world_size > 1: - sampler = DistributedSampler( - self.eval_dataset, num_replicas=self.world_size, rank=self.local_rank, shuffle=True, seed=42 - ) - else: - sampler = None self.eval_image_dataloader = CacheDataloader( self.eval_dataset, num_images_to_sample_from=self.config.eval_num_images_to_sample_from, device=self.device, num_workers=self.world_size * 4, pin_memory=True, - sampler=sampler, ) self.iter_eval_image_dataloader = iter(self.eval_image_dataloader) self.eval_pixel_sampler = PixelSampler(self.config.eval_num_rays_per_batch) diff --git a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py index 1e29f6b11..66cfbd967 100644 --- a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py +++ b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py @@ -18,12 +18,13 @@ from __future__ import annotations import math from dataclasses import dataclass, field from pathlib import Path, PureWindowsPath -from typing import Literal, Optional, Type +from typing import Optional, Type import numpy as np import torch from PIL import Image from rich.console import Console +from typing_extensions import Literal from nerfstudio.cameras import camera_utils from nerfstudio.cameras.cameras import CAMERA_MODEL_TO_TYPE, Cameras, CameraType @@ -53,8 +54,12 @@ class NerfstudioDataParserConfig(DataParserConfig): """How much to downscale images. If not set, images are chosen such that the max dimension is <1600px.""" scene_scale: float = 1.0 """How much to scale the region of interest by.""" - orientation_method: Literal["pca", "up"] = "up" + orientation_method: Literal["pca", "up", "none"] = "up" """The method to use for orientation.""" + center_poses: bool = True + """Whether to center the poses.""" + auto_scale_poses: bool = True + """Whether to automatically scale the poses to fit in +/- 1 bounding box.""" train_split_percentage: float = 0.9 """The percent of images to use for training. The remaining images are for eval.""" @@ -112,10 +117,15 @@ class Nerfstudio(DataParser): raise ValueError(f"Unknown dataparser split {split}") poses = torch.from_numpy(np.array(poses).astype(np.float32)) - poses = camera_utils.auto_orient_poses(poses, method=self.config.orientation_method) + poses = camera_utils.auto_orient_and_center_poses( + poses, method=self.config.orientation_method, center_poses=self.config.center_poses + ) # Scale poses - scale_factor = 1.0 / torch.max(torch.abs(poses[:, :3, 3])) + scale_factor = 1.0 + if self.config.auto_scale_poses: + scale_factor /= torch.max(torch.abs(poses[:, :3, 3])) + poses[:, :3, 3] *= scale_factor * self.config.scale_factor # Choose image_filenames and poses based on split, but after auto orient and scaling the poses. diff --git a/nerfstudio/data/dataparsers/record3d_dataparser.py b/nerfstudio/data/dataparsers/record3d_dataparser.py index ef20f21ba..513f9b4c6 100644 --- a/nerfstudio/data/dataparsers/record3d_dataparser.py +++ b/nerfstudio/data/dataparsers/record3d_dataparser.py @@ -17,11 +17,12 @@ from __future__ import annotations from dataclasses import dataclass, field from pathlib import Path -from typing import Literal, Type +from typing import Type import numpy as np import torch from scipy.spatial.transform import Rotation +from typing_extensions import Literal from nerfstudio.cameras import camera_utils from nerfstudio.cameras.cameras import Cameras, CameraType @@ -105,9 +106,9 @@ class Record3D(DataParser): # convert to Tensors poses = torch.from_numpy(poses[:, :3, :4]) - poses = camera_utils.auto_orient_poses(pose_utils.to4x4(poses), method=self.config.orientation_method)[ - :, :3, :4 - ] + poses = camera_utils.auto_orient_and_center_poses( + pose_utils.to4x4(poses), method=self.config.orientation_method + )[:, :3, :4] # Centering poses poses[:, :3, 3] = poses[:, :3, 3] - torch.mean(poses[:, :3, 3], dim=0) diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index 686cfcca4..5b9b60c73 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -21,7 +21,7 @@ import dataclasses import functools import os import time -from typing import Any, Dict, List, Tuple +from typing import Dict, List, Tuple import torch from rich.console import Console @@ -48,20 +48,6 @@ from nerfstudio.viewer.server import viewer_utils CONSOLE = Console(width=120) -def train_loop(local_rank: int, world_size: int, config: cfg.Config) -> Any: - """Main training function that sets up and runs the trainer per process - - Args: - local_rank: current rank of process - world_size: total number of gpus available - config: config file specifying training regimen - """ - trainer = Trainer(config, local_rank, world_size) - trainer.setup() - trainer.train() - return 0 - - class Trainer: """Trainer class @@ -122,7 +108,7 @@ class Trainer: Args: test_mode: Whether to setup for testing. Defaults to False. """ - self.pipeline: VanillaPipeline = self.config.pipeline.setup( + self.pipeline = self.config.pipeline.setup( device=self.device, test_mode=test_mode, world_size=self.world_size, local_rank=self.local_rank ) self.optimizers = setup_optimizers(self.config, self.pipeline.get_param_groups()) diff --git a/nerfstudio/field_components/encodings.py b/nerfstudio/field_components/encodings.py index b65bcb08b..7e149e531 100644 --- a/nerfstudio/field_components/encodings.py +++ b/nerfstudio/field_components/encodings.py @@ -17,13 +17,14 @@ Encoding functions """ from abc import abstractmethod -from typing import Literal, Optional +from typing import Optional import numpy as np import torch import torch.nn.functional as F from torch import nn from torchtyping import TensorType +from typing_extensions import Literal from nerfstudio.field_components.base_field_component import FieldComponent from nerfstudio.utils.math import components_from_spherical_harmonics, expected_sin diff --git a/nerfstudio/model_components/losses.py b/nerfstudio/model_components/losses.py index 14564a5b3..d06b27828 100644 --- a/nerfstudio/model_components/losses.py +++ b/nerfstudio/model_components/losses.py @@ -30,47 +30,36 @@ LOSSES = {"L1": L1Loss, "MSE": MSELoss} EPS = 1.0e-7 -def searchsorted(a: TensorType[..., "N"], v: TensorType[..., "M"]) -> TensorType[..., "M"]: - """ - https://github.com/kakaobrain/NeRF-Factory/blob/f61bb8744a5cb4820a4d968fb3bfbed777550f4a/src/model/mipnerf360/helper.py#L109 - https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/stepfun.py#L30 +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"]: + """Faster version of - - """ - i = torch.arange(a.shape[-1], device=a.device) - v_ge_a = v[..., None, :] >= a[..., :, None] - idx_lo = torch.where(v_ge_a, i[..., :, None], i[..., :1, None]).max(dim=-2).values - idx_hi = torch.where(~v_ge_a, i[..., :, None], i[..., -1:, None]).min(dim=-2).values - return idx_lo, idx_hi - - -def inner_outer( - t0: TensorType[..., "num_samples+1"], t1: TensorType[..., "num_samples+1"], y1: TensorType[..., "num_samples"] -): - """ https://github.com/kakaobrain/NeRF-Factory/blob/f61bb8744a5cb4820a4d968fb3bfbed777550f4a/src/model/mipnerf360/helper.py#L117 https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/stepfun.py#L64 Args: - t0: interval edges - t1: interval edges + t0_starts: start of the interval edges + t0_ends: end of the interval edges + t1_starts: start of the interval edges + t1_ends: end of the interval edges y1: weights """ - cy1 = torch.cat([torch.zeros_like(y1[..., :1]), torch.cumsum(y1, dim=-1)], dim=-1) - idx_lo, idx_hi = searchsorted(t1, t0) - cy1_lo = torch.take_along_dim(cy1, idx_lo, dim=-1) - cy1_hi = torch.take_along_dim(cy1, idx_hi, dim=-1) + idx_lo = torch.searchsorted(t1_starts.contiguous(), t0_starts.contiguous(), side="right") - 1 + idx_lo = torch.clamp(idx_lo, min=0, max=y1.shape[-1] - 1) + idx_hi = torch.searchsorted(t1_ends.contiguous(), t0_ends.contiguous(), side="right") + idx_hi = torch.clamp(idx_hi, min=0, max=y1.shape[-1] - 1) + cy1_lo = torch.take_along_dim(cy1[..., :-1], idx_lo, dim=-1) + cy1_hi = torch.take_along_dim(cy1[..., 1:], idx_hi, dim=-1) + y0_outer = cy1_hi - cy1_lo - y0_outer = cy1_hi[..., 1:] - cy1_lo[..., :-1] - y0_inner = torch.where( - idx_hi[..., :-1] <= idx_lo[..., 1:], - cy1_lo[..., 1:] - cy1_hi[..., :-1], - torch.zeros_like(cy1_lo[..., 1:]), - ) - - return y0_inner, y0_outer + return y0_outer def lossfun_outer( @@ -89,7 +78,7 @@ def lossfun_outer( t_env: interval edges of the upper bound enveloping historgram w_env: weights that should upper bound the inner (t,w) histogram """ - _, w_outer = inner_outer(t, t_env, w_env) + w_outer = outer(t[..., :-1], t[..., 1:], t_env[..., :-1], t_env[..., 1:], w_env) return torch.clip(w - w_outer, min=0) ** 2 / (w + EPS) diff --git a/nerfstudio/model_components/renderers.py b/nerfstudio/model_components/renderers.py index 64b8bced3..025194ee4 100644 --- a/nerfstudio/model_components/renderers.py +++ b/nerfstudio/model_components/renderers.py @@ -27,12 +27,13 @@ Example: """ import math -from typing import Literal, Optional, Union +from typing import Optional, Union import nerfacc import torch from torch import nn from torchtyping import TensorType +from typing_extensions import Literal from nerfstudio.cameras.rays import RaySamples from nerfstudio.utils.math import components_from_spherical_harmonics diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 6f6293f60..6929c2fda 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -19,7 +19,7 @@ NeRF implementation that combines many recent advancements. from __future__ import annotations from dataclasses import dataclass, field -from typing import Dict, List, Literal, Tuple, Type +from typing import Dict, List, Tuple, Type import numpy as np import torch @@ -27,6 +27,7 @@ from torch.nn import Parameter from torchmetrics import PeakSignalNoiseRatio from torchmetrics.functional import structural_similarity_index_measure from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity +from typing_extensions import Literal from nerfstudio.cameras.rays import RayBundle from nerfstudio.engine.callbacks import ( diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 2858c881d..7c64e268f 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -227,9 +227,7 @@ class VanillaPipeline(Pipeline): self.world_size = world_size if world_size > 1: - self._model = typing.cast( - Model, typing.cast(Model, DDP(self._model, device_ids=[local_rank], find_unused_parameters=True)) - ) + self._model = typing.cast(Model, DDP(self._model, device_ids=[local_rank], find_unused_parameters=True)) dist.barrier(device_ids=[local_rank]) @property diff --git a/nerfstudio/utils/comms.py b/nerfstudio/utils/comms.py index e009f6e92..9be472531 100644 --- a/nerfstudio/utils/comms.py +++ b/nerfstudio/utils/comms.py @@ -13,40 +13,38 @@ # limitations under the License. """functionality to handle multiprocessing syncing and communicating""" -import torch import torch.distributed as dist -_LOCAL_PROCESS_GROUP = None +LOCAL_PROCESS_GROUP = None + + +def is_dist_avail_and_initialized() -> bool: + """Returns True if distributed is available and initialized.""" + return dist.is_available() and dist.is_initialized() def get_world_size() -> int: """Get total number of available gpus""" - if not dist.is_available(): - return 1 - if not dist.is_initialized(): + if not is_dist_avail_and_initialized(): return 1 return dist.get_world_size() def get_rank() -> int: """Get global rank of current thread""" - if not dist.is_available(): - return 0 - if not dist.is_initialized(): + if not is_dist_avail_and_initialized(): return 0 return dist.get_rank() def get_local_rank() -> int: """The rank of the current process within the local (per-machine) process group.""" - if not dist.is_available(): - return 0 - if not dist.is_initialized(): + if not is_dist_avail_and_initialized(): return 0 assert ( - _LOCAL_PROCESS_GROUP is not None + LOCAL_PROCESS_GROUP is not None ), "Local process group is not created! Please use launch() to spawn processes!" - return dist.get_rank(group=_LOCAL_PROCESS_GROUP) + return dist.get_rank(group=LOCAL_PROCESS_GROUP) def get_local_size() -> int: @@ -54,11 +52,9 @@ def get_local_size() -> int: The size of the per-machine process group, i.e. the number of processes per machine. """ - if not dist.is_available(): + if not is_dist_avail_and_initialized(): return 1 - if not dist.is_initialized(): - return 1 - return dist.get_world_size(group=_LOCAL_PROCESS_GROUP) + return dist.get_world_size(group=LOCAL_PROCESS_GROUP) def is_main_process() -> bool: @@ -66,21 +62,16 @@ def is_main_process() -> bool: return get_rank() == 0 -def synchronize(world_size): +def synchronize(): """ Helper function to synchronize (barrier) among all processes when using distributed training """ - if not dist.is_available(): - return - if not dist.is_initialized(): - return - world_size = dist.get_world_size() - if world_size == 1: + if dist.get_world_size() == 1: return if dist.get_backend() == dist.Backend.NCCL: # This argument is needed to avoid warnings. # It's valid only for NCCL backend. - dist.barrier(device_ids=[torch.cuda.current_device()]) + dist.barrier(device_ids=[get_local_rank()]) else: dist.barrier() diff --git a/nerfstudio/utils/writer.py b/nerfstudio/utils/writer.py index bd9741a7d..e1ed07057 100644 --- a/nerfstudio/utils/writer.py +++ b/nerfstudio/utils/writer.py @@ -166,7 +166,6 @@ def write_out_storage(): EVENT_STORAGE.clear() -@check_main_thread def setup_local_writer(config: cfg.LoggingConfig, max_iter: int, banner_messages: Optional[List[str]] = None) -> None: """Initialization of all event writers specified in config diff --git a/nerfstudio/viewer/app/src/index.scss b/nerfstudio/viewer/app/src/index.scss index d6dd69216..c39522b58 100644 --- a/nerfstudio/viewer/app/src/index.scss +++ b/nerfstudio/viewer/app/src/index.scss @@ -175,8 +175,9 @@ $sidebar-width: 100% !default; .CameraPanel-slider-container { margin-top: 40px; - margin-left: 20px; - margin-right: 20px; + margin-left: auto; + margin-right: auto; + width: 70%; } .CameraPanel-slider-button-container { diff --git a/nerfstudio/viewer/app/src/modules/Scene/Scene.jsx b/nerfstudio/viewer/app/src/modules/Scene/Scene.jsx index 8ad2bc5d7..dd0d36fc6 100644 --- a/nerfstudio/viewer/app/src/modules/Scene/Scene.jsx +++ b/nerfstudio/viewer/app/src/modules/Scene/Scene.jsx @@ -253,7 +253,12 @@ export function get_scene_tree() { mouseVector.x = 2 * (e.clientX / size.x) - 1; mouseVector.y = 1 - 2 * ((e.clientY - BANNER_HEIGHT) / size.y); - if (mouseVector.x > 1 || mouseVector.x < -1 || mouseVector.y > 1 || mouseVector.y < -1) { + if ( + mouseVector.x > 1 || + mouseVector.x < -1 || + mouseVector.y > 1 || + mouseVector.y < -1 + ) { if (selectedCam !== null) { selectedCam.material.color = new THREE.Color(1, 1, 1); selectedCam = null; @@ -268,7 +273,9 @@ export function get_scene_tree() { selectedCam.material.color = new THREE.Color(1, 1, 1); selectedCam = null; } - const filtered_intersections = intersections.filter(isect => checkVisibility(isect.object)); + const filtered_intersections = intersections.filter((isect) => + checkVisibility(isect.object), + ); if (filtered_intersections.length > 0) { selectedCam = filtered_intersections[0].object; selectedCam.material.color = new THREE.Color(0xfab300); diff --git a/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPanel.jsx b/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPanel.jsx index 948e166f4..9f1cf207d 100644 --- a/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPanel.jsx +++ b/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPanel.jsx @@ -15,12 +15,15 @@ import { MeshLine, MeshLineMaterial } from 'meshline'; import { useContext, useEffect } from 'react'; import { useDispatch, useSelector } from 'react-redux'; +import AddAPhotoIcon from '@mui/icons-material/AddAPhoto'; import ArrowBackIosNewIcon from '@mui/icons-material/ArrowBackIosNew'; import ArrowForwardIosIcon from '@mui/icons-material/ArrowForwardIos'; import { CSS2DObject } from 'three/examples/jsm/renderers/CSS2DRenderer'; import ContentPasteGoIcon from '@mui/icons-material/ContentPasteGo'; import DeleteIcon from '@mui/icons-material/Delete'; import FirstPageIcon from '@mui/icons-material/FirstPage'; +import FileDownloadOutlinedIcon from '@mui/icons-material/FileDownloadOutlined'; +import FileUploadOutlinedIcon from '@mui/icons-material/FileUploadOutlined'; import IconButton from '@mui/material/IconButton'; import LastPageIcon from '@mui/icons-material/LastPage'; import PauseIcon from '@mui/icons-material/Pause'; @@ -423,14 +426,27 @@ export default function CameraPanel(props) { }); } + const keyframes = []; + for (let i = 0; i < cameras.length; i += 1) { + const camera = cameras[i]; + keyframes.push({ + matrix: JSON.stringify(camera.matrix.toArray()), + fov: camera.fov, + aspect: camera_render.aspect, + }); + } + // const myData const camera_path_object = { - keyframes: [], + keyframes, render_height, render_width, camera_path, fps, seconds, + smoothness_value, + is_cycle, + is_linear, }; return camera_path_object; }; @@ -461,6 +477,58 @@ export default function CameraPanel(props) { URL.revokeObjectURL(href); }; + const load_camera_path = (camera_path_object) => { + // TODO UI for getting json + + const new_camera_list = []; + + setRenderHeight(camera_path_object.render_height); + setUIRenderHeight(camera_path_object.render_height); + setRenderWidth(camera_path_object.render_width); + setUIRenderWidth(camera_path_object.render_width); + + setFps(camera_path_object.fps); + setUIfps(camera_path_object.fps); + + setSeconds(camera_path_object.seconds); + setUISeconds(camera_path_object.seconds); + + set_smoothness_value(camera_path_object.smoothness_value); + setIsCycle(camera_path_object.is_cycle); + setIsLinear(camera_path_object.is_linear); + + for (let i = 0; i < camera_path_object.keyframes.length; i += 1) { + const keyframe = camera_path_object.keyframes[i]; + const camera = new THREE.PerspectiveCamera( + keyframe.fov, + keyframe.aspect, + 0.1, + 1000, + ); + + const mat = new THREE.Matrix4(); + mat.fromArray(JSON.parse(keyframe.matrix)); + // camera.matrix = mat; + set_camera_position(camera, mat); + new_camera_list.push(camera); + } + + setCameras(new_camera_list); + reset_slider_render_on_add(new_camera_list); + }; + + const uploadCameraPath = (e) => { + const fileUpload = e.target.files[0]; + + const fr = new FileReader(); + fr.onload = (res) => { + const camera_path_object = JSON.parse(res.target.result); + load_camera_path(camera_path_object); + }; + + fr.readAsText(fileUpload); + }; + const copy_cmd_to_clipboard = () => { console.log('copy_cmd_to_clipboard'); @@ -497,40 +565,25 @@ export default function CameraPanel(props) { }); }; - const setUp = () => { - const rot = camera_main.rotation; - const unitY = new THREE.Vector3(0, 1, 0); - const upVec = unitY.applyEuler(rot); - - const grid = sceneTree.find_object_no_create(['Grid']); - grid.setRotationFromEuler(rot); - - const pos = new THREE.Vector3(); - camera_main.getWorldPosition(pos); - camera_main.up.set(upVec.x, upVec.y, upVec.z); - sceneTree.metadata.camera_controls.updateCameraUp(); - sceneTree.metadata.camera_controls.setLookAt(pos.x, pos.y, pos.z, 0, 0, 0); - const points = [new THREE.Vector3(0, 0, 0), upVec.multiplyScalar(2)]; - const geometry = new THREE.BufferGeometry().setFromPoints(points); - const material = new THREE.LineBasicMaterial({ - color: 0xaa46fc, - linewidth: 1, - }); - const line = new THREE.LineSegments(geometry, material); - sceneTree.set_object_from_path(['Viewer Up Vector'], line); - }; - return (
-
- -
-
@@ -538,6 +591,7 @@ export default function CameraPanel(props) { size="small" className="CameraPanel-top-button" variant="outlined" + startIcon={} onClick={export_camera_path} > Export Path @@ -551,7 +605,147 @@ export default function CameraPanel(props) {
+
+ { + if (e.target.validity.valid) { + setUIRenderHeight(e.target.value); + } + }} + onBlur={(e) => { + if (e.target.validity.valid) { + if (e.target.value !== '') { + setRenderHeight(parseInt(e.target.value, 10)); + } else { + setUIRenderHeight(render_height); + } + } + }} + value={ui_render_height} + error={ui_render_height <= 0} + helperText={ui_render_height <= 0 ? 'Required' : ''} + variant="standard" + /> + { + if (e.target.validity.valid) { + setUIRenderWidth(e.target.value); + } + }} + onBlur={(e) => { + if (e.target.validity.valid) { + if (e.target.value !== '') { + setRenderWidth(parseInt(e.target.value, 10)); + } else { + setUIRenderWidth(render_width); + } + } + }} + value={ui_render_width} + error={ui_render_width <= 0} + helperText={ui_render_width <= 0 ? 'Required' : ''} + variant="standard" + /> + { + if (e.target.validity.valid) { + setUIFieldOfView(e.target.value); + } + }} + onBlur={(e) => { + if (e.target.validity.valid) { + if (e.target.value !== '') { + setFOV(parseInt(e.target.value, 10)); + } else { + setUIFieldOfView(field_of_view); + } + } + }} + value={ui_field_of_view} + error={ui_field_of_view <= 0} + helperText={ui_field_of_view <= 0 ? 'Required' : ''} + variant="standard" + /> +
+
+ { + if (e.target.validity.valid) { + setUISeconds(e.target.value); + } + }} + onBlur={(e) => { + if (e.target.validity.valid) { + if (e.target.value !== '') { + setSeconds(parseInt(e.target.value, 10)); + } else { + setUISeconds(seconds); + } + } + }} + value={ui_seconds} + error={ui_seconds <= 0} + helperText={ui_seconds <= 0 ? 'Required' : ''} + variant="standard" + /> + { + if (e.target.validity.valid) { + setUIfps(e.target.value); + } + }} + onBlur={(e) => { + if (e.target.validity.valid) { + if (e.target.value !== '') { + setFps(parseInt(e.target.value, 10)); + } else { + setUIfps(fps); + } + } + }} + value={ui_fps} + error={ui_fps <= 0} + helperText={ui_fps <= 0 ? 'Required' : ''} + variant="standard" + /> +
+
+ +
{!is_cycle ? ( @@ -709,136 +903,6 @@ export default function CameraPanel(props) {
-
- { - if (e.target.validity.valid) { - setUIRenderHeight(e.target.value); - } - }} - onBlur={(e) => { - if (e.target.validity.valid) { - if (e.target.value !== '') { - setRenderHeight(parseInt(e.target.value, 10)); - } else { - setUIRenderHeight(render_height); - } - } - }} - value={ui_render_height} - error={ui_render_height <= 0} - helperText={ui_render_height <= 0 ? 'Required' : ''} - variant="standard" - /> - { - if (e.target.validity.valid) { - setUIRenderWidth(e.target.value); - } - }} - onBlur={(e) => { - if (e.target.validity.valid) { - if (e.target.value !== '') { - setRenderWidth(parseInt(e.target.value, 10)); - } else { - setUIRenderWidth(render_width); - } - } - }} - value={ui_render_width} - error={ui_render_width <= 0} - helperText={ui_render_width <= 0 ? 'Required' : ''} - variant="standard" - /> - { - if (e.target.validity.valid) { - setUIFieldOfView(e.target.value); - } - }} - onBlur={(e) => { - if (e.target.validity.valid) { - if (e.target.value !== '') { - setFOV(parseInt(e.target.value, 10)); - } else { - setUIFieldOfView(field_of_view); - } - } - }} - value={ui_field_of_view} - error={ui_field_of_view <= 0} - helperText={ui_field_of_view <= 0 ? 'Required' : ''} - variant="standard" - /> -
-
- { - if (e.target.validity.valid) { - setUISeconds(e.target.value); - } - }} - onBlur={(e) => { - if (e.target.validity.valid) { - if (e.target.value !== '') { - setSeconds(parseInt(e.target.value, 10)); - } else { - setUISeconds(seconds); - } - } - }} - value={ui_seconds} - error={ui_seconds <= 0} - helperText={ui_seconds <= 0 ? 'Required' : ''} - variant="standard" - /> - { - if (e.target.validity.valid) { - setUIfps(e.target.value); - } - }} - onBlur={(e) => { - if (e.target.validity.valid) { - if (e.target.value !== '') { - setFps(parseInt(e.target.value, 10)); - } else { - setUIfps(fps); - } - } - }} - value={ui_fps} - error={ui_fps <= 0} - helperText={ui_fps <= 0 ? 'Required' : ''} - variant="standard" - /> -
{ + const mat = new THREE.Matrix4(); + mat.fromArray(matrix.elements); + mat.decompose(camera.position, camera.quaternion, camera.scale); + const unit = new THREE.Vector3(0, 0, -1); + const viewDirection = unit.applyMatrix4(mat); + sceneTree.metadata.camera_controls.setLookAt( + camera.position.x, + camera.position.y, + camera.position.z, + viewDirection.x, + viewDirection.y, + viewDirection.z, + ); +}; + +interface ListItemProps { + name: String; + sceneTree: SceneNode; + scene_node: SceneNode; + level: Number; + groupVisible: Boolean; + canSnap: Boolean; +} + +function MenuItems(props: ListItemProps) { + const name = props.name; + const sceneTree = props.sceneTree; + const scene_node = props.scene_node; + const level = props.level; + const groupVisible = props.groupVisible; + const canSnap = props.canSnap; + + // TODO: sort the keys by string + const terminal = Object.keys(scene_node.children).includes(''); + + const getCamera = (node) => { + return node.object.children[0]; + }; + + const num_children = Object.keys(scene_node.children).length; + if (num_children === 0) { + return null; + } + + const [open, setOpen] = React.useState(true); + const toggleOpen = () => { + setOpen(!open); + }; + + const [visible, setVisible] = React.useState(groupVisible); + if (terminal) { + scene_node.object.visible = visible; + } + + const toggleVisible = (e) => { + e.stopPropagation(); + setVisible(!visible); + }; + + React.useEffect(() => { + setVisible(groupVisible); + }, [groupVisible]); + + return ( + <> + + + + + + + {canSnap && ( + + snap_to_camera( + sceneTree, + sceneTree.metadata.camera, + getCamera(scene_node).matrix, + ) + } + sx={{ mr: 1 }} + > + + + )} + + + {visible ? : } + + + {terminal + ? null + : (() => { + if (open) { + return ; + } + return ; + })()} + + + {terminal ? null : ( + + + {Object.keys(scene_node.children) + .filter((key) => { + if (key === 'Camera') { + return false; // skip + } + return true; + }) + .map((key) => ( + + ))} + + + )} + + ); +} + +interface ClickableListProps { + sceneTree: object; +} + +function ClickableList(props: ClickableListProps) { + const sceneTree = props.sceneTree; + + return ( + + + + ); +} + +export default function ScenePanel(props) { + const sceneTree = props.sceneTree; + const camera_main = sceneTree.find_object_no_create([ + 'Cameras', + 'Main Camera', + ]); + const setUp = () => { + const rot = camera_main.rotation; + const unitY = new THREE.Vector3(0, 1, 0); + const upVec = unitY.applyEuler(rot); + + const grid = sceneTree.find_object_no_create(['Grid']); + grid.setRotationFromEuler(rot); + + const pos = new THREE.Vector3(); + camera_main.getWorldPosition(pos); + camera_main.up.set(upVec.x, upVec.y, upVec.z); + sceneTree.metadata.camera_controls.updateCameraUp(); + sceneTree.metadata.camera_controls.setLookAt(pos.x, pos.y, pos.z, 0, 0, 0); + const points = [new THREE.Vector3(0, 0, 0), upVec.multiplyScalar(2)]; + const geometry = new THREE.BufferGeometry().setFromPoints(points); + const material = new THREE.LineBasicMaterial({ + color: 0xaa46fc, + linewidth: 1, + }); + const line = new THREE.LineSegments(geometry, material); + sceneTree.set_object_from_path(['Viewer Up Vector'], line); + }; + return ( +
+
+ +
+ +
+ ); +} diff --git a/nerfstudio/viewer/app/src/modules/SidePanel/ScenePanel/index.jsx b/nerfstudio/viewer/app/src/modules/SidePanel/ScenePanel/index.jsx new file mode 100644 index 000000000..87dc14f57 --- /dev/null +++ b/nerfstudio/viewer/app/src/modules/SidePanel/ScenePanel/index.jsx @@ -0,0 +1,3 @@ +import ScenePanel from './ScenePanel'; + +export default ScenePanel; diff --git a/nerfstudio/viewer/app/src/modules/SidePanel/SidePanel.jsx b/nerfstudio/viewer/app/src/modules/SidePanel/SidePanel.jsx index 51368df70..5f1da13cc 100644 --- a/nerfstudio/viewer/app/src/modules/SidePanel/SidePanel.jsx +++ b/nerfstudio/viewer/app/src/modules/SidePanel/SidePanel.jsx @@ -3,33 +3,21 @@ import * as React from 'react'; import * as THREE from 'three'; -import { FaTractor } from 'react-icons/fa'; - import Box from '@mui/material/Box'; -import { Collapse, IconButton } from '@mui/material'; import Divider from '@mui/material/Divider'; -import List from '@mui/material/List'; -import ListItemButton from '@mui/material/ListItemButton'; -import ListItemText from '@mui/material/ListItemText'; -import ListItemIcon from '@mui/material/ListItemIcon'; import { Leva } from 'leva'; import Tab from '@mui/material/Tab'; import Tabs from '@mui/material/Tabs'; import Typography from '@mui/material/Typography'; import { CameraAltRounded, - ExpandLess, - ExpandMore, TuneRounded, - Videocam, - Visibility, - VisibilityOff, WidgetsRounded, } from '@mui/icons-material/'; import StatusPanel from './StatusPanel'; -import SceneNode from '../../SceneNode'; import LevaTheme from '../ConfigPanel/leva_theme.json'; import CameraPanel from './CameraPanel'; +import ScenePanel from './ScenePanel'; import { RenderControls } from '../ConfigPanel/ConfigPanel'; import { LogPanel } from '../LogPanel/LogPanel'; @@ -81,165 +69,6 @@ function a11yProps(index: number) { }; } -interface ListItemProps { - name: String; - sceneTree: SceneNode; - scene_node: SceneNode; - level: Number; - groupVisible: Boolean; - canSnap: Boolean; -} - -function MenuItems(props: ListItemProps) { - const name = props.name; - const sceneTree = props.sceneTree; - const scene_node = props.scene_node; - const level = props.level; - const groupVisible = props.groupVisible; - const canSnap = props.canSnap; - - // TODO: sort the keys by string - const terminal = Object.keys(scene_node.children).includes(''); - - const getCamera = (node) => { - return node.object.children[0]; - }; - - const num_children = Object.keys(scene_node.children).length; - if (num_children === 0) { - return null; - } - - const [open, setOpen] = React.useState(true); - const toggleOpen = () => { - setOpen(!open); - }; - - const [visible, setVisible] = React.useState(groupVisible); - if (terminal) { - scene_node.object.visible = visible; - } - - const toggleVisible = (e) => { - e.stopPropagation(); - setVisible(!visible); - }; - - React.useEffect(() => { - setVisible(groupVisible); - }, [groupVisible]); - - return ( - <> - - - - - - - {canSnap && ( - - snap_to_camera( - sceneTree, - sceneTree.metadata.camera, - getCamera(scene_node).matrix, - ) - } - sx={{ mr: 1 }} - > - - - )} - - - {visible ? : } - - - {terminal - ? null - : (() => { - if (open) { - return ; - } - return ; - })()} - - - {terminal ? null : ( - - - {Object.keys(scene_node.children) - .filter((key) => { - if (key === 'Camera') { - return false; // skip - } - return true; - }) - .map((key) => ( - - ))} - - - )} - - ); -} - -interface ClickableListProps { - sceneTree: object; -} - -function ClickableList(props: ClickableListProps) { - const sceneTree = props.sceneTree; - - return ( - - - - ); -} - interface BasicTabsProps { sceneTree: object; } @@ -290,7 +119,7 @@ export function BasicTabs(props: BasicTabsProps) {
- +
diff --git a/nerfstudio/viewer/app/src/themes/theme.ts b/nerfstudio/viewer/app/src/themes/theme.ts index a9ce1fbf8..4bfed2ce3 100644 --- a/nerfstudio/viewer/app/src/themes/theme.ts +++ b/nerfstudio/viewer/app/src/themes/theme.ts @@ -123,5 +123,13 @@ export const appTheme = createTheme({ }, }, }, + MuiList: { + styleOverrides: { + padding: { + paddingTop: 4, + paddingBottom: 4, + }, + }, + }, }, }); diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index 0a831c5fe..70848d08d 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/nerfstudio/viewer/server/viewer_utils.py @@ -55,6 +55,13 @@ warnings.filterwarnings("ignore", category=CryptographyDeprecationWarning) CONSOLE = Console(width=120) +def get_viewer_version() -> str: + """Get the version of the viewer.""" + json_filename = os.path.join(os.path.dirname(__file__), "../app/package.json") + version = load_from_json(Path(json_filename))["version"] + return version + + @check_main_thread def setup_viewer(config: cfg.ViewerConfig, log_filename: Path): """Sets up the viewer if enabled @@ -230,8 +237,7 @@ class ViewerState: ) # TODO(ethan): log the output of the viewer bridge server in a file where the training logs go CONSOLE.line() - json_filename = os.path.join(os.path.dirname(__file__), "../app/package.json") - version = load_from_json(Path(json_filename))["version"] + version = get_viewer_version() websocket_url = f"ws://localhost:{self.config.websocket_port}" self.viewer_url = f"https://viewer.nerf.studio/versions/{version}/?websocket_url={websocket_url}" CONSOLE.rule(characters="=") diff --git a/pyproject.toml b/pyproject.toml index a4630212c..4e9dee650 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,11 +4,11 @@ build-backend = "setuptools.build_meta" [project] name = "nerfstudio" -version = "0.1.5.1" +version = "0.1.6" description = "All-in-one repository for state-of-the-art NeRFs" readme = "README.md" license = { text="Apache 2.0"} -requires-python = ">=3.8" +requires-python = ">=3.7.3" classifiers = [ "Development Status :: 3 - Alpha", "Programming Language :: Python", @@ -30,7 +30,7 @@ dependencies = [ "mediapy==1.1.0", "msgpack==1.0.4", "msgpack_numpy==0.4.8", - "nerfacc==0.2.0", + "nerfacc==0.2.1", "opencv-python==4.6.0.66", "plotly==5.7.0", "protobuf==3.20.0", diff --git a/scripts/completions/install.py b/scripts/completions/install.py index 3d413aa16..37cfd9fbc 100755 --- a/scripts/completions/install.py +++ b/scripts/completions/install.py @@ -10,12 +10,12 @@ import shutil import stat import subprocess import sys -from typing import List, Literal, Union +from typing import List, Union import tyro from rich.console import Console from rich.prompt import Confirm -from typing_extensions import assert_never +from typing_extensions import Literal, assert_never ConfigureMode = Literal["install", "uninstall"] ShellType = Literal["zsh", "bash"] @@ -175,10 +175,7 @@ def _update_rc( assert_never(mode) -def main( - mode: ConfigureMode = "install", - /, -) -> None: +def main(mode: ConfigureMode = "install") -> None: """Main script. Args: diff --git a/scripts/downloads/download_data.py b/scripts/downloads/download_data.py old mode 100644 new mode 100755 index d75e4a1f1..aa96b25c1 --- a/scripts/downloads/download_data.py +++ b/scripts/downloads/download_data.py @@ -4,11 +4,12 @@ import shutil import sys import zipfile from pathlib import Path -from typing import Literal, Optional +from typing import Optional import gdown import tyro from rich.console import Console +from typing_extensions import Literal CONSOLE = Console(width=120) @@ -99,7 +100,7 @@ def download_nerfstudio(save_dir: Path, capture_name: str): record3d_file_ids = { - "bear": grab_file_id("https://drive.google.com/file/d/1WRZohWMRj0nNlYFIEBwkddDoGPvLTzkR/view?usp=sharing"), + "bear": grab_file_id("https://drive.google.com/file/d/1WRZohWMRj0nNlYFIEBwkddDoGPvLTzkR/view?usp=sharing") } @@ -108,8 +109,25 @@ def download_record3d(save_dir: Path, capture_name: str): download_capture_name(save_dir, "record3d", capture_name, capture_name_to_file_id=record3d_file_ids) +def download_dnerf(save_dir: Path): + """Download the D-NeRF dataset (https://github.com/albertpumarola/D-NeRF).""" + # TODO: give this code the same structure as download_nerfstudio + + final_path = save_dir / Path("dnerf") + if os.path.exists(final_path): + shutil.rmtree(str(final_path)) + download_path = save_dir / "dnerf_data.zip" + os.system(f"curl -L https://www.dropbox.com/s/raw/0bf6fl0ye2vz3vr/data.zip > {download_path}") + with zipfile.ZipFile(download_path, "r") as zip_ref: + zip_ref.extractall(str(save_dir)) + unzip_path = save_dir / Path("data") + final_path = save_dir / Path("dnerf") + unzip_path.rename(final_path) + download_path.unlink(missing_ok=True) + + def main( - dataset: Literal["blender", "friends", "nerfstudio", "record3d"], + dataset: Literal["blender", "friends", "nerfstudio", "record3d", "dnerf"], capture_name: Optional[DatasetName] = None, # type: ignore save_dir: Path = Path("data/"), ): @@ -148,6 +166,10 @@ def main( download_nerfstudio(save_dir, capture_name) if dataset == "record3d": download_record3d(save_dir, "bear") + if dataset == "dnerf": + if capture_name is not None: + CONSOLE.print("Capture name is ignored when downloading from the dnerf dataset.") + download_dnerf(save_dir) def entrypoint(): diff --git a/scripts/process_data.py b/scripts/process_data.py index 8dbf5d6f3..d62918c35 100755 --- a/scripts/process_data.py +++ b/scripts/process_data.py @@ -9,7 +9,7 @@ from contextlib import nullcontext from dataclasses import dataclass from enum import Enum from pathlib import Path -from typing import List, Literal, Optional, Tuple, Union +from typing import List, Optional, Tuple, Union import appdirs import numpy as np @@ -19,7 +19,7 @@ import yaml from PIL import Image from rich.console import Console from rich.progress import track -from typing_extensions import Annotated +from typing_extensions import Annotated, Literal from nerfstudio.utils import colmap_utils, install_checks @@ -166,7 +166,7 @@ def convert_video_to_images( spacing = num_frames // num_frames_target if spacing > 1: - ffmpeg_cmd += f" -vf 'thumbnail={spacing},setpts=N/TB' -r 1" + ffmpeg_cmd += f" -vf thumbnail={spacing},setpts=N/TB -r 1" else: CONSOLE.print("[bold red]Can't satify requested number of frames. Extracting all frames.") diff --git a/scripts/render.py b/scripts/render.py index 5ada54af4..b37d7dfdb 100644 --- a/scripts/render.py +++ b/scripts/render.py @@ -8,7 +8,7 @@ import dataclasses import json import sys from pathlib import Path -from typing import Literal, Optional +from typing import Optional import mediapy as media import torch @@ -21,7 +21,7 @@ from rich.progress import ( TextColumn, TimeRemainingColumn, ) -from typing_extensions import assert_never +from typing_extensions import Literal, assert_never from nerfstudio.cameras.camera_paths import get_path_from_json, get_spiral_path from nerfstudio.cameras.cameras import Cameras diff --git a/scripts/train.py b/scripts/train.py index 86dd77d68..37717a11c 100755 --- a/scripts/train.py +++ b/scripts/train.py @@ -48,7 +48,7 @@ from rich.console import Console from nerfstudio.configs import base_config as cfg from nerfstudio.configs.config_utils import convert_markup_to_ansi from nerfstudio.configs.method_configs import AnnotatedBaseConfigUnion -from nerfstudio.engine.trainer import train_loop +from nerfstudio.engine.trainer import Trainer from nerfstudio.utils import comms, profiler CONSOLE = Console(width=120) @@ -59,7 +59,7 @@ torch.backends.cudnn.benchmark = True # type: ignore def _find_free_port() -> str: - """Find free port on open socket""" + """Finds a free port.""" sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.bind(("", 0)) port = sock.getsockname()[1] @@ -74,6 +74,20 @@ def _set_random_seed(seed) -> None: torch.manual_seed(seed) +def train_loop(local_rank: int, world_size: int, config: cfg.Config, global_rank: int = 0): + """Main training function that sets up and runs the trainer per process + + Args: + local_rank: current rank of process + world_size: total number of gpus available + config: config file specifying training regimen + """ + _set_random_seed(config.machine.seed + global_rank) + trainer = Trainer(config, local_rank, world_size) + trainer.setup() + trainer.train() + + def _distributed_worker( local_rank: int, main_func: Callable, @@ -85,55 +99,47 @@ def _distributed_worker( timeout: timedelta = DEFAULT_TIMEOUT, ) -> Any: """Spawned distributed worker that handles the initialization of process group and handles the - training process on multiple processes + training process on multiple processes. Args: - local_rank (int): current rank of process - main_func (Callable): function that will be called by the distributed workers - world_size (int): total number of gpus available - num_gpus_per_machine (int): number of GPUs per machine - machine_rank (int): rank of this machine - dist_url (str): url to connect to for distributed jobs, including protocol - e.g. "tcp://127.0.0.1:8686". - Can be set to "auto" to automatically select a free port on localhost - config (Config): config file specifying training regimen - timeout (timedelta, optional): timeout of the distributed workers Defaults to DEFAULT_TIMEOUT. + local_rank: Current rank of process. + main_func: Function that will be called by the distributed workers. + world_size: Total number of gpus available. + num_gpus_per_machine: Number of GPUs per machine. + machine_rank: Rank of this machine. + dist_url: URL to connect to for distributed jobs, including protocol + E.g., "tcp://127.0.0.1:8686". + It can be set to "auto" to automatically select a free port on localhost. + config: Config specifying training regimen. + timeout: Timeout of the distributed workers. Raises: e: Exception in initializing the process group Returns: - Any: TODO(): determine the return type + Any: TODO: determine the return type """ assert torch.cuda.is_available(), "cuda is not available. Please check your installation." global_rank = machine_rank * num_gpus_per_machine + local_rank - try: - dist.init_process_group( - backend="NCCL", - init_method=dist_url, - world_size=world_size, - rank=global_rank, - timeout=timeout, - ) - except Exception as e: - CONSOLE.log(f"Process group URL: {dist_url}") - raise e - assert comms._LOCAL_PROCESS_GROUP is None # pylint: disable=protected-access + dist.init_process_group( + backend="nccl", + init_method=dist_url, + world_size=world_size, + rank=global_rank, + timeout=timeout, + ) + assert comms.LOCAL_PROCESS_GROUP is None num_machines = world_size // num_gpus_per_machine for i in range(num_machines): ranks_on_i = list(range(i * num_gpus_per_machine, (i + 1) * num_gpus_per_machine)) pg = dist.new_group(ranks_on_i) if i == machine_rank: - comms._LOCAL_PROCESS_GROUP = pg # pylint: disable=protected-access + comms.LOCAL_PROCESS_GROUP = pg assert num_gpus_per_machine <= torch.cuda.device_count() - torch.cuda.set_device(local_rank) - _set_random_seed(config.machine.seed + global_rank) - comms.synchronize(world_size) - - output = main_func(local_rank, world_size, config) - comms.synchronize(world_size) + output = main_func(local_rank, world_size, config, global_rank) + comms.synchronize() dist.destroy_process_group() return output @@ -160,23 +166,20 @@ def launch( """ assert config is not None world_size = num_machines * num_gpus_per_machine - if world_size == 0: - # Using only CPU and one process. - _set_random_seed(config.machine.seed) - main_func(local_rank=0, world_size=0, config=config) - elif world_size == 1: - # Using one gpu and one process. - _set_random_seed(config.machine.seed) + if world_size <= 1: + # world_size=0 uses one CPU in one process. + # world_size=1 uses one GPU in one process. try: - main_func(local_rank=0, world_size=1, config=config) + main_func(local_rank=0, world_size=world_size, config=config) except KeyboardInterrupt: - print(traceback.format_exc()) + # print the stack trace + CONSOLE.print(traceback.format_exc()) finally: profiler.flush_profiler(config.logging) elif world_size > 1: # Using multiple gpus with multiple processes. if dist_url == "auto": - assert num_machines == 1, "dist_url=auto not supported in multi-machine jobs." + assert num_machines == 1, "dist_url=auto is not supported for multi-machine jobs." port = _find_free_port() dist_url = f"tcp://127.0.0.1:{port}" if num_machines > 1 and dist_url.startswith("file://"): @@ -185,6 +188,7 @@ def launch( process_context = mp.spawn( _distributed_worker, nprocs=num_gpus_per_machine, + join=False, args=( main_func, world_size, @@ -195,16 +199,18 @@ def launch( timeout, ), ) + # process_context won't be None because join=False, so it's okay to assert this + # for Pylance reasons assert process_context is not None try: process_context.join() except KeyboardInterrupt: for i, process in enumerate(process_context.processes): if process.is_alive(): - CONSOLE.log("Terminating process %s", str(i)) + CONSOLE.log(f"Terminating process {i}...") process.terminate() process.join() - CONSOLE.log("Process %s finished", str(i)) + CONSOLE.log(f"Process {i} finished.") finally: profiler.flush_profiler(config.logging) @@ -226,8 +232,8 @@ def main(config: cfg.Config) -> None: config.save_config() launch( - train_loop, - config.machine.num_gpus, + main_func=train_loop, + num_gpus_per_machine=config.machine.num_gpus, num_machines=config.machine.num_machines, machine_rank=config.machine.machine_rank, dist_url=config.machine.dist_url, diff --git a/tests/test_train.py b/tests/test_train.py index 0816088c8..6b69590da 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -12,7 +12,7 @@ import pytest from nerfstudio.configs.base_config import Config from nerfstudio.configs.method_configs import method_configs from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig -from nerfstudio.engine.trainer import train_loop +from scripts.train import train_loop BLACKLIST = ["base", "semantic-nerfw", "instant-ngp", "nerfacto"]