From 44dafa3a4ec31d41aa49da644cb840df932d3ffe Mon Sep 17 00:00:00 2001 From: SiN <58646324+bro-sin@users.noreply.github.com> Date: Sat, 13 Jan 2024 14:24:27 +0800 Subject: [PATCH 1/4] Add --no-cache-dir option to pip install in Dockerfile (#2739) --- Dockerfile | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/Dockerfile b/Dockerfile index 8c9fe4f31c..4c00a8faa9 100644 --- a/Dockerfile +++ b/Dockerfile @@ -119,33 +119,33 @@ ENV PATH="${PATH}:/home/user/.local/bin" SHELL ["/bin/bash", "-c"] # Upgrade pip and install packages. -RUN python3.10 -m pip install --upgrade pip setuptools pathtools promise pybind11 +RUN python3.10 -m pip install --no-cache-dir --upgrade pip setuptools pathtools promise pybind11 # Install pytorch and submodules -RUN CUDA_VER=${CUDA_VERSION%.*} && CUDA_VER=${CUDA_VER//./} && python3.10 -m pip install \ +RUN CUDA_VER=${CUDA_VERSION%.*} && CUDA_VER=${CUDA_VER//./} && python3.10 -m pip install --no-cache-dir \ torch==2.0.1+cu${CUDA_VER} \ torchvision==0.15.2+cu${CUDA_VER} \ --extra-index-url https://download.pytorch.org/whl/cu${CUDA_VER} # Install tynyCUDNN (we need to set the target architectures as environment variable first). ENV TCNN_CUDA_ARCHITECTURES=${CUDA_ARCHITECTURES} -RUN python3.10 -m pip install git+https://github.com/NVlabs/tiny-cuda-nn.git@v1.6#subdirectory=bindings/torch +RUN python3.10 -m pip install --no-cache-dir git+https://github.com/NVlabs/tiny-cuda-nn.git@v1.6#subdirectory=bindings/torch # Install pycolmap, required by hloc. RUN git clone --branch v0.4.0 --recursive https://github.com/colmap/pycolmap.git && \ cd pycolmap && \ - python3.10 -m pip install . && \ + python3.10 -m pip install --no-cache-dir . && \ cd .. # Install hloc 1.4 as alternative feature detector and matcher option for nerfstudio. RUN git clone --branch master --recursive https://github.com/cvg/Hierarchical-Localization.git && \ cd Hierarchical-Localization && \ git checkout v1.4 && \ - python3.10 -m pip install -e . && \ + python3.10 -m pip install --no-cache-dir -e . && \ cd .. # Install pyceres from source RUN git clone --branch v1.0 --recursive https://github.com/cvg/pyceres.git && \ cd pyceres && \ - python3.10 -m pip install -e . && \ + python3.10 -m pip install --no-cache-dir -e . && \ cd .. # Install pixel perfect sfm. @@ -153,10 +153,10 @@ RUN git clone --recursive https://github.com/cvg/pixel-perfect-sfm.git && \ cd pixel-perfect-sfm && \ git reset --hard 40f7c1339328b2a0c7cf71f76623fb848e0c0357 && \ git clean -df && \ - python3.10 -m pip install -e . && \ + python3.10 -m pip install --no-cache-dir -e . && \ cd .. -RUN python3.10 -m pip install omegaconf +RUN python3.10 -m pip install --no-cache-dir omegaconf # Copy nerfstudio folder and give ownership to user. ADD . /home/user/nerfstudio USER root @@ -165,7 +165,7 @@ USER ${USER_ID} # Install nerfstudio dependencies. RUN cd nerfstudio && \ - python3.10 -m pip install -e . && \ + python3.10 -m pip install --no-cache-dir -e . && \ cd .. # Change working directory From 69ed6d40935ced2a295702423dac8c0020e3a7d1 Mon Sep 17 00:00:00 2001 From: Brent Yi Date: Sat, 13 Jan 2024 12:31:39 -0800 Subject: [PATCH 2/4] Improve viewer banner when server is bound to 0.0.0.0 (#2761) * Improve viewer banner when server is bound to 0.0.0.0 * Add comment * Update viewer.py --- nerfstudio/engine/trainer.py | 9 +++---- nerfstudio/scripts/viewer/run_viewer.py | 4 ++-- nerfstudio/utils/writer.py | 2 +- nerfstudio/viewer/viewer.py | 32 +++++++++++++++---------- 4 files changed, 28 insertions(+), 19 deletions(-) diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index b51f5b5565..35bb428441 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -21,11 +21,12 @@ import functools import os import time +from collections import defaultdict from dataclasses import dataclass, field from pathlib import Path from threading import Lock -from typing import Dict, List, Literal, Optional, Tuple, Type, cast, DefaultDict -from collections import defaultdict +from typing import DefaultDict, Dict, List, Literal, Optional, Tuple, Type, cast + import torch from nerfstudio.configs.experiment_config import ExperimentConfig from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes, TrainingCallbackLocation @@ -36,8 +37,8 @@ from nerfstudio.utils.misc import step_check from nerfstudio.utils.rich_utils import CONSOLE from nerfstudio.utils.writer import EventName, TimeWriter -from nerfstudio.viewer_legacy.server.viewer_state import ViewerLegacyState from nerfstudio.viewer.viewer import Viewer as ViewerState +from nerfstudio.viewer_legacy.server.viewer_state import ViewerLegacyState from rich import box, style from rich.panel import Panel from rich.table import Table @@ -182,7 +183,7 @@ def setup(self, test_mode: Literal["test", "val", "inference"] = "val") -> None: train_lock=self.train_lock, share=self.config.viewer.make_share_url, ) - banner_messages = [f"Viewer at: {self.viewer_state.viewer_url}"] + banner_messages = self.viewer_state.viewer_info self._check_viewer_warnings() self._load_checkpoint() diff --git a/nerfstudio/scripts/viewer/run_viewer.py b/nerfstudio/scripts/viewer/run_viewer.py index 8d679fb226..f071478d50 100644 --- a/nerfstudio/scripts/viewer/run_viewer.py +++ b/nerfstudio/scripts/viewer/run_viewer.py @@ -99,7 +99,7 @@ def _start_viewer(config: TrainerConfig, pipeline: Pipeline, step: int): datapath=pipeline.datamanager.get_datapath(), pipeline=pipeline, ) - banner_messages = [f"Legaccy viewer at: {viewer_state.viewer_url}"] + banner_messages = [f"Legacy viewer at: {viewer_state.viewer_url}"] if config.vis == "viewer": viewer_state = ViewerState( config.viewer, @@ -108,7 +108,7 @@ def _start_viewer(config: TrainerConfig, pipeline: Pipeline, step: int): pipeline=pipeline, share=config.viewer.make_share_url, ) - banner_messages = [f"Viewer at: {viewer_state.viewer_url}"] + banner_messages = viewer_state.viewer_info # We don't need logging, but writer.GLOBAL_BUFFER needs to be populated config.logging.local_writer.enable = False diff --git a/nerfstudio/utils/writer.py b/nerfstudio/utils/writer.py index 59e6c5a713..aacddbe4a9 100644 --- a/nerfstudio/utils/writer.py +++ b/nerfstudio/utils/writer.py @@ -518,7 +518,7 @@ def _print_stats(self, latest_map, padding=" "): for i, mssg in enumerate(self.past_mssgs): pad_len = len(max(self.past_mssgs, key=len)) - style = "\x1b[6;30;42m" if self.banner_len and i >= len(self.past_mssgs) - self.banner_len + 1 else "" + style = "\x1b[30;42m" if self.banner_len and i >= len(self.past_mssgs) - self.banner_len + 1 else "" print(f"{style}{mssg:{padding}<{pad_len}} \x1b[0m") else: print(curr_mssg) diff --git a/nerfstudio/viewer/viewer.py b/nerfstudio/viewer/viewer.py index d17b93e069..00d1925529 100644 --- a/nerfstudio/viewer/viewer.py +++ b/nerfstudio/viewer/viewer.py @@ -23,7 +23,6 @@ import numpy as np import torch import torchvision -import viser import viser.theme import viser.transforms as vtf from nerfstudio.cameras.camera_optimizers import CameraOptimizer @@ -34,15 +33,17 @@ from nerfstudio.pipelines.base_pipeline import Pipeline from nerfstudio.utils.decorators import check_main_thread, decorate_all from nerfstudio.utils.writer import GLOBAL_BUFFER, EventName -from nerfstudio.viewer_legacy.server import viewer_utils from nerfstudio.viewer.control_panel import ControlPanel from nerfstudio.viewer.export_panel import populate_export_tab from nerfstudio.viewer.render_panel import populate_render_tab from nerfstudio.viewer.render_state_machine import RenderAction, RenderStateMachine from nerfstudio.viewer.utils import CameraState, parse_object from nerfstudio.viewer.viewer_elements import ViewerControl, ViewerElement +from nerfstudio.viewer_legacy.server import viewer_utils from typing_extensions import assert_never +import viser + if TYPE_CHECKING: from nerfstudio.engine.trainer import Trainer @@ -63,13 +64,12 @@ class Viewer: share: print a shareable URL Attributes: - viewer_url: url to open viewer + viewer_info: information string for the viewer viser_server: the viser server """ - viewer_url: str + viewer_info: List[str] viser_server: viser.ViserServer - camera_state: Optional[CameraState] = None def __init__( self, @@ -107,15 +107,23 @@ def __init__( self.viser_server = viser.ViserServer(host=config.websocket_host, port=websocket_port) # Set the name of the URL either to the share link if available, or the localhost + share_url = None if share: - url = self.viser_server.request_share_url() - if url is not None: - print("Couldn't make share URL") - self.viewer_url = url - else: - self.viewer_url = f"http://{config.websocket_host}:{websocket_port}" + share_url = self.viser_server.request_share_url() + if share_url is None: + print("Couldn't make share URL!") + + if share_url is not None: + self.viewer_info = [f"Viewer at: http://localhost:{websocket_port} or {share_url}"] + elif config.websocket_host == "0.0.0.0": + # 0.0.0.0 is not a real IP address and was confusing people, so + # we'll just print localhost instead. There are some security + # (and IPv6 compatibility) implications here though, so we should + # note that the server is bound to 0.0.0.0! + self.viewer_info = [f"Viewer running locally at: http://localhost:{websocket_port} (listening on 0.0.0.0)"] else: - self.viewer_url = f"http://{config.websocket_host}:{websocket_port}" + self.viewer_info = [f"Viewer running locally at: http://{config.websocket_host}:{websocket_port}"] + buttons = ( viser.theme.TitlebarButton( text="Getting Started", From f0aadf9de8b9da4c576a207ad399450bbac4315e Mon Sep 17 00:00:00 2001 From: Matthew Hoffman Date: Mon, 15 Jan 2024 13:43:14 -0800 Subject: [PATCH 3/4] Use ruff as linter and formatter (#2765) * Switch from black to ruff format * Sort dependencies with ruff check * Run ruff format --- .github/workflows/core_code_checks.yml | 8 ++-- .pre-commit-config.yaml | 13 +++--- .vscode/settings.json | 8 ++-- docs/reference/contributing.md | 17 ++++---- nerfstudio/cameras/camera_optimizers.py | 6 ++- nerfstudio/cameras/camera_utils.py | 2 +- nerfstudio/cameras/cameras.py | 4 +- nerfstudio/configs/dataparser_configs.py | 2 +- nerfstudio/configs/method_configs.py | 21 +++------ .../data/datamanagers/base_datamanager.py | 12 +----- .../random_cameras_datamanager.py | 2 +- .../dataparsers/arkitscenes_dataparser.py | 6 +-- .../data/dataparsers/blender_dataparser.py | 6 +-- .../data/dataparsers/colmap_dataparser.py | 20 ++++----- .../data/dataparsers/dnerf_dataparser.py | 6 +-- .../data/dataparsers/dycheck_dataparser.py | 6 +-- .../dataparsers/instant_ngp_dataparser.py | 14 ++---- .../data/dataparsers/minimal_dataparser.py | 7 +-- .../data/dataparsers/nerfosr_dataparser.py | 6 +-- .../data/dataparsers/nerfstudio_dataparser.py | 12 ++---- .../data/dataparsers/nuscenes_dataparser.py | 6 +-- .../dataparsers/phototourism_dataparser.py | 11 +---- .../data/dataparsers/scannet_dataparser.py | 6 +-- .../data/dataparsers/scannetpp_dataparser.py | 5 +-- .../data/dataparsers/sdfstudio_dataparser.py | 6 +-- .../data/dataparsers/sitcoms3d_dataparser.py | 7 +-- nerfstudio/data/datasets/depth_dataset.py | 16 +++---- nerfstudio/data/pixel_samplers.py | 11 +---- nerfstudio/data/scene_box.py | 4 +- nerfstudio/data/utils/pixel_sampling_utils.py | 2 +- nerfstudio/engine/trainer.py | 9 ++-- nerfstudio/exporter/exporter_utils.py | 2 +- nerfstudio/exporter/marching_cubes.py | 4 +- nerfstudio/field_components/__init__.py | 3 +- nerfstudio/field_components/encodings.py | 6 +-- nerfstudio/field_components/mlp.py | 5 +-- nerfstudio/fields/semantic_nerf_field.py | 7 +-- nerfstudio/fields/vanilla_nerf_field.py | 7 +-- nerfstudio/generative/deepfloyd.py | 9 +--- .../generative/positional_text_embeddings.py | 7 +-- nerfstudio/generative/stable_diffusion.py | 5 +-- nerfstudio/generative/utils.py | 1 + nerfstudio/models/base_model.py | 2 +- nerfstudio/models/depth_nerfacto.py | 2 +- nerfstudio/models/gaussian_splatting.py | 4 +- nerfstudio/models/generfacto.py | 43 ++++++++----------- nerfstudio/models/instant_ngp.py | 12 +----- nerfstudio/models/mipnerf.py | 6 +-- nerfstudio/models/neus.py | 6 +-- nerfstudio/models/neus_facto.py | 11 +---- nerfstudio/models/semantic_nerfw.py | 6 +-- nerfstudio/models/tensorf.py | 23 +++------- nerfstudio/models/vanilla_nerf.py | 8 +--- nerfstudio/pipelines/dynamic_batch.py | 2 +- nerfstudio/plugins/registry_dataparser.py | 2 +- nerfstudio/process_data/odm_utils.py | 6 +-- nerfstudio/process_data/process_data_utils.py | 1 + .../scripts/blender/nerfstudio_blender.py | 2 +- nerfstudio/scripts/completions/install.py | 3 +- nerfstudio/scripts/downloads/download_data.py | 2 +- nerfstudio/scripts/exporter.py | 6 +-- nerfstudio/scripts/github/run_actions.py | 4 +- nerfstudio/scripts/process_data.py | 2 +- nerfstudio/scripts/render.py | 28 +++--------- nerfstudio/scripts/viewer/run_viewer.py | 2 +- nerfstudio/utils/eval_utils.py | 2 +- nerfstudio/utils/misc.py | 6 +-- nerfstudio/utils/profiler.py | 19 +------- nerfstudio/utils/rich_utils.py | 10 +---- nerfstudio/utils/writer.py | 5 ++- nerfstudio/viewer/control_panel.py | 3 +- nerfstudio/viewer/export_panel.py | 3 +- nerfstudio/viewer/render_panel.py | 10 +++-- nerfstudio/viewer/render_state_machine.py | 4 +- nerfstudio/viewer/utils.py | 3 +- nerfstudio/viewer/viewer.py | 6 +-- nerfstudio/viewer/viewer_elements.py | 2 +- nerfstudio/viewer_legacy/viser/__init__.py | 3 +- nerfstudio/viewer_legacy/viser/gui.py | 20 +-------- pyproject.toml | 13 +++--- .../dataparsers/test_nerfstudio_dataparser.py | 5 +-- tests/pipelines/test_vanilla_pipeline.py | 9 +--- tests/process_data/test_process_images.py | 10 +++-- 83 files changed, 206 insertions(+), 427 deletions(-) diff --git a/.github/workflows/core_code_checks.yml b/.github/workflows/core_code_checks.yml index 761d4841b6..0f1fcf0caf 100644 --- a/.github/workflows/core_code_checks.yml +++ b/.github/workflows/core_code_checks.yml @@ -32,10 +32,10 @@ jobs: - name: Check notebook cell metadata run: | python ./nerfstudio/scripts/docs/add_nb_tags.py --check - - name: Run Ruff - run: ruff docs/ nerfstudio/ tests/ - - name: Run Black - run: black docs/ nerfstudio/ tests/ --check + - name: Run Ruff Linter + run: ruff check docs/ nerfstudio/ tests/ + - name: Run Ruff Formatter + run: ruff format docs/ nerfstudio/ tests/ --check - name: Run Pyright run: | pyright diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 4ccc2a61f8..bdb8ff4f49 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -12,16 +12,15 @@ repos: files: '.*' pass_filenames: false - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v3.2.0 + rev: v4.5.0 hooks: - id: trailing-whitespace - id: end-of-file-fixer - repo: https://github.com/charliermarsh/ruff-pre-commit - # Ruff version. - rev: 'v0.0.267' + rev: v0.1.13 hooks: - id: ruff -- repo: https://github.com/psf/black - rev: '23.3.0' - hooks: - - id: black + types_or: [ python, pyi, jupyter ] + args: [ --fix ] + - id: ruff-format + types_or: [ python, pyi, jupyter ] diff --git a/.vscode/settings.json b/.vscode/settings.json index 8b0f1d7501..c365ae07b4 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -24,16 +24,16 @@ "typescript.suggestionActions.enabled": false, "javascript.suggestionActions.enabled": false, "[python]": { - "editor.defaultFormatter": "ms-python.black-formatter", + "editor.defaultFormatter": "charliermarsh.ruff", "editor.codeActionsOnSave": { - "source.organizeImports": true, - "source.fixAll": true + "source.organizeImports": "explicit", + "source.fixAll": "explicit" } }, "editor.formatOnSave": true, + "editor.rulers": [120], "python.envFile": "${workspaceFolder}/.env", "python.formatting.provider": "none", - "black-formatter.args": ["--line-length=120"], "python.linting.pylintEnabled": false, "python.linting.flake8Enabled": false, "python.linting.enabled": true, diff --git a/docs/reference/contributing.md b/docs/reference/contributing.md index bf0af33db3..818a18dbac 100644 --- a/docs/reference/contributing.md +++ b/docs/reference/contributing.md @@ -14,15 +14,14 @@ In addition to code contributions, we also encourage contributors to add their o Below are the various tooling features our team uses to maintain this codebase. -| Tooling | Support | -| --------------- | ---------------------------------------------------------- | -| Formatting | [Black](https://black.readthedocs.io/en/stable/) | -| Linter | [Ruff](https://beta.ruff.rs/docs/) | -| Type checking | [Pyright](https://github.com/microsoft/pyright) | -| Testing | [pytest](https://docs.pytest.org/en/7.1.x/) | -| Docs | [Sphinx](https://www.sphinx-doc.org/en/master/) | -| Docstring style | [Google](https://google.github.io/styleguide/pyguide.html) | -| JS Linting | [eslint](https://eslint.org/) | +| Tooling | Support | +| -------------------- | ---------------------------------------------------------- | +| Formatting & Linting | [Ruff](https://beta.ruff.rs/docs/) | +| Type checking | [Pyright](https://github.com/microsoft/pyright) | +| Testing | [pytest](https://docs.pytest.org/en/7.1.x/) | +| Docs | [Sphinx](https://www.sphinx-doc.org/en/master/) | +| Docstring style | [Google](https://google.github.io/styleguide/pyguide.html) | +| JS Linting | [eslint](https://eslint.org/) | ## Requirements diff --git a/nerfstudio/cameras/camera_optimizers.py b/nerfstudio/cameras/camera_optimizers.py index 4aa69a69b6..b192ad8b93 100644 --- a/nerfstudio/cameras/camera_optimizers.py +++ b/nerfstudio/cameras/camera_optimizers.py @@ -27,13 +27,13 @@ from torch import Tensor, nn from typing_extensions import assert_never +from nerfstudio.cameras.cameras import Cameras from nerfstudio.cameras.lie_groups import exp_map_SE3, exp_map_SO3xR3 from nerfstudio.cameras.rays import RayBundle from nerfstudio.configs.base_config import InstantiateConfig -from nerfstudio.utils import poses as pose_utils from nerfstudio.engine.optimizers import OptimizerConfig from nerfstudio.engine.schedulers import SchedulerConfig -from nerfstudio.cameras.cameras import Cameras +from nerfstudio.utils import poses as pose_utils @dataclass @@ -60,6 +60,7 @@ class CameraOptimizerConfig(InstantiateConfig): def __post_init__(self): if self.optimizer is not None: import warnings + from nerfstudio.utils.rich_utils import CONSOLE CONSOLE.print( @@ -70,6 +71,7 @@ def __post_init__(self): if self.scheduler is not None: import warnings + from nerfstudio.utils.rich_utils import CONSOLE CONSOLE.print( diff --git a/nerfstudio/cameras/camera_utils.py b/nerfstudio/cameras/camera_utils.py index 4456566237..7c98ef080c 100644 --- a/nerfstudio/cameras/camera_utils.py +++ b/nerfstudio/cameras/camera_utils.py @@ -346,7 +346,7 @@ def _compute_residual_and_jacobian( xd: torch.Tensor, yd: torch.Tensor, distortion_params: torch.Tensor, -) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor,]: +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Auxiliary function of radial_and_tangential_undistort() that computes residuals and jacobians. Adapted from MultiNeRF: https://github.com/google-research/multinerf/blob/b02228160d3179300c7d499dca28cb9ca3677f32/internal/camera_utils.py#L427-L474 diff --git a/nerfstudio/cameras/cameras.py b/nerfstudio/cameras/cameras.py index 3086298e2c..f81323b45b 100644 --- a/nerfstudio/cameras/cameras.py +++ b/nerfstudio/cameras/cameras.py @@ -670,7 +670,7 @@ def _generate_rays_from_coords( assert c2w.shape == num_rays_shape + (3, 4) def _compute_rays_for_omnidirectional_stereo( - eye: Literal["left", "right"] + eye: Literal["left", "right"], ) -> Tuple[Float[Tensor, "num_rays_shape 3"], Float[Tensor, "3 num_rays_shape 3"]]: """Compute the rays for an omnidirectional stereo camera @@ -727,7 +727,7 @@ def _compute_rays_for_omnidirectional_stereo( return ods_origins_circle, directions_stack def _compute_rays_for_vr180( - eye: Literal["left", "right"] + eye: Literal["left", "right"], ) -> Tuple[Float[Tensor, "num_rays_shape 3"], Float[Tensor, "3 num_rays_shape 3"]]: """Compute the rays for a VR180 camera diff --git a/nerfstudio/configs/dataparser_configs.py b/nerfstudio/configs/dataparser_configs.py index e3fa66ec5c..d9e11b3045 100644 --- a/nerfstudio/configs/dataparser_configs.py +++ b/nerfstudio/configs/dataparser_configs.py @@ -33,9 +33,9 @@ from nerfstudio.data.dataparsers.nuscenes_dataparser import NuScenesDataParserConfig from nerfstudio.data.dataparsers.phototourism_dataparser import PhototourismDataParserConfig from nerfstudio.data.dataparsers.scannet_dataparser import ScanNetDataParserConfig +from nerfstudio.data.dataparsers.scannetpp_dataparser import ScanNetppDataParserConfig from nerfstudio.data.dataparsers.sdfstudio_dataparser import SDFStudioDataParserConfig from nerfstudio.data.dataparsers.sitcoms3d_dataparser import Sitcoms3DDataParserConfig -from nerfstudio.data.dataparsers.scannetpp_dataparser import ScanNetppDataParserConfig from nerfstudio.plugins.registry_dataparser import discover_dataparsers dataparsers = { diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 1f1a09edbb..55ba44eb41 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -26,23 +26,16 @@ from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig from nerfstudio.configs.base_config import ViewerConfig from nerfstudio.configs.external_methods import get_external_methods -from nerfstudio.data.datamanagers.base_datamanager import ( - VanillaDataManager, - VanillaDataManagerConfig, -) +from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager, VanillaDataManagerConfig +from nerfstudio.data.datamanagers.full_images_datamanager import FullImageDatamanagerConfig from nerfstudio.data.datamanagers.parallel_datamanager import ParallelDataManagerConfig -from nerfstudio.data.datamanagers.random_cameras_datamanager import ( - RandomCamerasDataManagerConfig, -) +from nerfstudio.data.datamanagers.random_cameras_datamanager import RandomCamerasDataManagerConfig from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig +from nerfstudio.data.dataparsers.colmap_dataparser import ColmapDataParserConfig from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig -from nerfstudio.data.dataparsers.instant_ngp_dataparser import ( - InstantNGPDataParserConfig, -) +from nerfstudio.data.dataparsers.instant_ngp_dataparser import InstantNGPDataParserConfig from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig -from nerfstudio.data.dataparsers.phototourism_dataparser import ( - PhototourismDataParserConfig, -) +from nerfstudio.data.dataparsers.phototourism_dataparser import PhototourismDataParserConfig from nerfstudio.data.dataparsers.sdfstudio_dataparser import SDFStudioDataParserConfig from nerfstudio.data.dataparsers.sitcoms3d_dataparser import Sitcoms3DDataParserConfig from nerfstudio.data.datasets.depth_dataset import DepthDataset @@ -62,7 +55,6 @@ from nerfstudio.models.gaussian_splatting import GaussianSplattingModelConfig from nerfstudio.models.generfacto import GenerfactoModelConfig from nerfstudio.models.instant_ngp import InstantNGPModelConfig -from nerfstudio.data.dataparsers.colmap_dataparser import ColmapDataParserConfig from nerfstudio.models.mipnerf import MipNerfModel from nerfstudio.models.nerfacto import NerfactoModelConfig from nerfstudio.models.neus import NeuSModelConfig @@ -71,7 +63,6 @@ from nerfstudio.models.tensorf import TensoRFModelConfig from nerfstudio.models.vanilla_nerf import NeRFModel, VanillaModelConfig from nerfstudio.pipelines.base_pipeline import VanillaPipelineConfig -from nerfstudio.data.datamanagers.full_images_datamanager import FullImageDatamanagerConfig from nerfstudio.pipelines.dynamic_batch import DynamicBatchPipelineConfig from nerfstudio.plugins.registry import discover_methods diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 322a1e8617..2131d5260b 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -54,16 +54,8 @@ from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig from nerfstudio.data.datasets.base_dataset import InputDataset -from nerfstudio.data.pixel_samplers import ( - PatchPixelSamplerConfig, - PixelSampler, - PixelSamplerConfig, -) -from nerfstudio.data.utils.dataloaders import ( - CacheDataloader, - FixedIndicesEvalDataloader, - RandIndicesEvalDataloader, -) +from nerfstudio.data.pixel_samplers import PatchPixelSamplerConfig, PixelSampler, PixelSamplerConfig +from nerfstudio.data.utils.dataloaders import CacheDataloader, FixedIndicesEvalDataloader, RandIndicesEvalDataloader from nerfstudio.data.utils.nerfstudio_collate import nerfstudio_collate from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes from nerfstudio.model_components.ray_generators import RayGenerator diff --git a/nerfstudio/data/datamanagers/random_cameras_datamanager.py b/nerfstudio/data/datamanagers/random_cameras_datamanager.py index 7704f2f161..6c638ec0e0 100644 --- a/nerfstudio/data/datamanagers/random_cameras_datamanager.py +++ b/nerfstudio/data/datamanagers/random_cameras_datamanager.py @@ -24,8 +24,8 @@ import torch from rich.progress import Console -from torch.nn import Parameter from torch import Tensor +from torch.nn import Parameter from typing_extensions import Literal from nerfstudio.cameras.cameras import Cameras diff --git a/nerfstudio/data/dataparsers/arkitscenes_dataparser.py b/nerfstudio/data/dataparsers/arkitscenes_dataparser.py index ada6f941a2..8dbd7d8bb9 100644 --- a/nerfstudio/data/dataparsers/arkitscenes_dataparser.py +++ b/nerfstudio/data/dataparsers/arkitscenes_dataparser.py @@ -24,11 +24,7 @@ from nerfstudio.cameras import camera_utils from nerfstudio.cameras.cameras import Cameras, CameraType -from nerfstudio.data.dataparsers.base_dataparser import ( - DataParser, - DataParserConfig, - DataparserOutputs, -) +from nerfstudio.data.dataparsers.base_dataparser import DataParser, DataParserConfig, DataparserOutputs from nerfstudio.data.scene_box import SceneBox diff --git a/nerfstudio/data/dataparsers/blender_dataparser.py b/nerfstudio/data/dataparsers/blender_dataparser.py index 9d6524d4e0..34fc3c52b8 100644 --- a/nerfstudio/data/dataparsers/blender_dataparser.py +++ b/nerfstudio/data/dataparsers/blender_dataparser.py @@ -24,11 +24,7 @@ import torch from nerfstudio.cameras.cameras import Cameras, CameraType -from nerfstudio.data.dataparsers.base_dataparser import ( - DataParser, - DataParserConfig, - DataparserOutputs, -) +from nerfstudio.data.dataparsers.base_dataparser import DataParser, DataParserConfig, DataparserOutputs from nerfstudio.data.scene_box import SceneBox from nerfstudio.utils.colors import get_color from nerfstudio.utils.io import load_from_json diff --git a/nerfstudio/data/dataparsers/colmap_dataparser.py b/nerfstudio/data/dataparsers/colmap_dataparser.py index c478a318f3..a0ee580e35 100644 --- a/nerfstudio/data/dataparsers/colmap_dataparser.py +++ b/nerfstudio/data/dataparsers/colmap_dataparser.py @@ -17,8 +17,8 @@ import sys from dataclasses import dataclass, field -from pathlib import Path from functools import partial +from pathlib import Path from typing import List, Literal, Optional, Type import numpy as np @@ -31,15 +31,15 @@ from nerfstudio.data.dataparsers.base_dataparser import DataParser, DataParserConfig, DataparserOutputs from nerfstudio.data.scene_box import SceneBox from nerfstudio.data.utils import colmap_parsing_utils as colmap_utils -from nerfstudio.process_data.colmap_utils import parse_colmap_camera_params -from nerfstudio.utils.scripts import run_command -from nerfstudio.utils.rich_utils import CONSOLE, status from nerfstudio.data.utils.dataparsers_utils import ( + get_train_eval_split_all, get_train_eval_split_filename, get_train_eval_split_fraction, get_train_eval_split_interval, - get_train_eval_split_all, ) +from nerfstudio.process_data.colmap_utils import parse_colmap_camera_params +from nerfstudio.utils.rich_utils import CONSOLE, status +from nerfstudio.utils.scripts import run_command MAX_AUTO_RESOLUTION = 1600 @@ -66,7 +66,7 @@ class ColmapDataParserConfig(DataParserConfig): """Whether to automatically scale the poses to fit in +/- 1 bounding box.""" eval_mode: Literal["fraction", "filename", "interval", "all"] = "interval" """ - The method to use for splitting the dataset into train and eval. + The method to use for splitting the dataset into train and eval. Fraction splits based on a percentage for train and the remaining for eval. Filename splits based on filenames containing train/eval. Interval uses every nth frame for eval (used by most academic papers, e.g. MipNerf360, GSplat). @@ -284,15 +284,11 @@ def _generate_dataparser_outputs(self, split: str = "train", **kwargs): if "depth_path" in frame: depth_filenames.append(Path(frame["depth_path"])) - assert len(mask_filenames) == 0 or ( - len(mask_filenames) == len(image_filenames) - ), """ + assert len(mask_filenames) == 0 or (len(mask_filenames) == len(image_filenames)), """ Different number of image and mask filenames. You should check that mask_path is specified for every frame (or zero frames) in transforms.json. """ - assert len(depth_filenames) == 0 or ( - len(depth_filenames) == len(image_filenames) - ), """ + assert len(depth_filenames) == 0 or (len(depth_filenames) == len(image_filenames)), """ Different number of image and depth filenames. You should check that depth_file_path is specified for every frame (or zero frames) in transforms.json. """ diff --git a/nerfstudio/data/dataparsers/dnerf_dataparser.py b/nerfstudio/data/dataparsers/dnerf_dataparser.py index 0c01eedba1..478b1fc38c 100644 --- a/nerfstudio/data/dataparsers/dnerf_dataparser.py +++ b/nerfstudio/data/dataparsers/dnerf_dataparser.py @@ -24,11 +24,7 @@ import torch from nerfstudio.cameras.cameras import Cameras, CameraType -from nerfstudio.data.dataparsers.base_dataparser import ( - DataParser, - DataParserConfig, - DataparserOutputs, -) +from nerfstudio.data.dataparsers.base_dataparser import DataParser, DataParserConfig, DataparserOutputs from nerfstudio.data.scene_box import SceneBox from nerfstudio.utils.colors import get_color from nerfstudio.utils.io import load_from_json diff --git a/nerfstudio/data/dataparsers/dycheck_dataparser.py b/nerfstudio/data/dataparsers/dycheck_dataparser.py index 9af4f58001..90f5e9e978 100644 --- a/nerfstudio/data/dataparsers/dycheck_dataparser.py +++ b/nerfstudio/data/dataparsers/dycheck_dataparser.py @@ -25,11 +25,7 @@ import torch from nerfstudio.cameras.cameras import Cameras, CameraType -from nerfstudio.data.dataparsers.base_dataparser import ( - DataParser, - DataParserConfig, - DataparserOutputs, -) +from nerfstudio.data.dataparsers.base_dataparser import DataParser, DataParserConfig, DataparserOutputs from nerfstudio.data.scene_box import SceneBox from nerfstudio.utils.colors import get_color from nerfstudio.utils.io import load_from_json diff --git a/nerfstudio/data/dataparsers/instant_ngp_dataparser.py b/nerfstudio/data/dataparsers/instant_ngp_dataparser.py index 6fc6047225..86780dcc75 100644 --- a/nerfstudio/data/dataparsers/instant_ngp_dataparser.py +++ b/nerfstudio/data/dataparsers/instant_ngp_dataparser.py @@ -18,7 +18,7 @@ from dataclasses import dataclass, field from pathlib import Path -from typing import Dict, Tuple, Type, Literal +from typing import Dict, Literal, Tuple, Type import imageio import numpy as np @@ -26,17 +26,13 @@ from nerfstudio.cameras import camera_utils from nerfstudio.cameras.cameras import Cameras, CameraType -from nerfstudio.data.dataparsers.base_dataparser import ( - DataParser, - DataParserConfig, - DataparserOutputs, -) +from nerfstudio.data.dataparsers.base_dataparser import DataParser, DataParserConfig, DataparserOutputs from nerfstudio.data.scene_box import SceneBox from nerfstudio.data.utils.dataparsers_utils import ( + get_train_eval_split_all, get_train_eval_split_filename, get_train_eval_split_fraction, get_train_eval_split_interval, - get_train_eval_split_all, ) from nerfstudio.utils.io import load_from_json from nerfstudio.utils.rich_utils import CONSOLE @@ -109,9 +105,7 @@ def _generate_dataparser_outputs(self, split="train"): mask_filenames.append(mask_fname) if num_skipped_image_filenames >= 0: CONSOLE.print(f"Skipping {num_skipped_image_filenames} files in dataset split {split}.") - assert ( - len(image_filenames) != 0 - ), """ + assert len(image_filenames) != 0, """ No image files found. You should check the file_paths in the transforms.json file to make sure they are correct. """ diff --git a/nerfstudio/data/dataparsers/minimal_dataparser.py b/nerfstudio/data/dataparsers/minimal_dataparser.py index 1aa40c12a0..afcef15b90 100644 --- a/nerfstudio/data/dataparsers/minimal_dataparser.py +++ b/nerfstudio/data/dataparsers/minimal_dataparser.py @@ -24,12 +24,7 @@ import torch from nerfstudio.cameras.cameras import Cameras -from nerfstudio.data.dataparsers.base_dataparser import ( - DataParser, - DataParserConfig, - DataparserOutputs, - Semantics, -) +from nerfstudio.data.dataparsers.base_dataparser import DataParser, DataParserConfig, DataparserOutputs, Semantics from nerfstudio.data.scene_box import SceneBox diff --git a/nerfstudio/data/dataparsers/nerfosr_dataparser.py b/nerfstudio/data/dataparsers/nerfosr_dataparser.py index 690ed88e67..db884cb736 100644 --- a/nerfstudio/data/dataparsers/nerfosr_dataparser.py +++ b/nerfstudio/data/dataparsers/nerfosr_dataparser.py @@ -30,11 +30,7 @@ from nerfstudio.cameras import camera_utils from nerfstudio.cameras.cameras import Cameras, CameraType -from nerfstudio.data.dataparsers.base_dataparser import ( - DataParser, - DataParserConfig, - DataparserOutputs, -) +from nerfstudio.data.dataparsers.base_dataparser import DataParser, DataParserConfig, DataparserOutputs from nerfstudio.data.scene_box import SceneBox diff --git a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py index 7f6f21b885..aaa88c7691 100644 --- a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py +++ b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py @@ -29,10 +29,10 @@ from nerfstudio.data.dataparsers.base_dataparser import DataParser, DataParserConfig, DataparserOutputs from nerfstudio.data.scene_box import SceneBox from nerfstudio.data.utils.dataparsers_utils import ( + get_train_eval_split_all, get_train_eval_split_filename, get_train_eval_split_fraction, get_train_eval_split_interval, - get_train_eval_split_all, ) from nerfstudio.utils.io import load_from_json from nerfstudio.utils.rich_utils import CONSOLE @@ -62,7 +62,7 @@ class NerfstudioDataParserConfig(DataParserConfig): """Whether to automatically scale the poses to fit in +/- 1 bounding box.""" eval_mode: Literal["fraction", "filename", "interval", "all"] = "fraction" """ - The method to use for splitting the dataset into train and eval. + The method to use for splitting the dataset into train and eval. Fraction splits based on a percentage for train and the remaining for eval. Filename splits based on filenames containing train/eval. Interval uses every nth frame for eval. @@ -179,15 +179,11 @@ def _generate_dataparser_outputs(self, split="train"): depth_fname = self._get_fname(depth_filepath, data_dir, downsample_folder_prefix="depths_") depth_filenames.append(depth_fname) - assert len(mask_filenames) == 0 or ( - len(mask_filenames) == len(image_filenames) - ), """ + assert len(mask_filenames) == 0 or (len(mask_filenames) == len(image_filenames)), """ Different number of image and mask filenames. You should check that mask_path is specified for every frame (or zero frames) in transforms.json. """ - assert len(depth_filenames) == 0 or ( - len(depth_filenames) == len(image_filenames) - ), """ + assert len(depth_filenames) == 0 or (len(depth_filenames) == len(image_filenames)), """ Different number of image and depth filenames. You should check that depth_file_path is specified for every frame (or zero frames) in transforms.json. """ diff --git a/nerfstudio/data/dataparsers/nuscenes_dataparser.py b/nerfstudio/data/dataparsers/nuscenes_dataparser.py index 0e899a68c9..19d215b140 100644 --- a/nerfstudio/data/dataparsers/nuscenes_dataparser.py +++ b/nerfstudio/data/dataparsers/nuscenes_dataparser.py @@ -25,11 +25,7 @@ from nuscenes.nuscenes import NuScenes as NuScenesDatabase from nerfstudio.cameras.cameras import Cameras, CameraType -from nerfstudio.data.dataparsers.base_dataparser import ( - DataParser, - DataParserConfig, - DataparserOutputs, -) +from nerfstudio.data.dataparsers.base_dataparser import DataParser, DataParserConfig, DataparserOutputs from nerfstudio.data.scene_box import SceneBox diff --git a/nerfstudio/data/dataparsers/phototourism_dataparser.py b/nerfstudio/data/dataparsers/phototourism_dataparser.py index 9249ae42a4..538bad2f0a 100644 --- a/nerfstudio/data/dataparsers/phototourism_dataparser.py +++ b/nerfstudio/data/dataparsers/phototourism_dataparser.py @@ -25,19 +25,12 @@ from nerfstudio.cameras import camera_utils from nerfstudio.cameras.cameras import Cameras, CameraType -from nerfstudio.data.dataparsers.base_dataparser import ( - DataParser, - DataParserConfig, - DataparserOutputs, -) +from nerfstudio.data.dataparsers.base_dataparser import DataParser, DataParserConfig, DataparserOutputs from nerfstudio.data.scene_box import SceneBox # TODO(1480) use pycolmap instead of colmap_parsing_utils # import pycolmap -from nerfstudio.data.utils.colmap_parsing_utils import ( - read_cameras_binary, - read_images_binary, -) +from nerfstudio.data.utils.colmap_parsing_utils import read_cameras_binary, read_images_binary from nerfstudio.utils.rich_utils import CONSOLE diff --git a/nerfstudio/data/dataparsers/scannet_dataparser.py b/nerfstudio/data/dataparsers/scannet_dataparser.py index 90a53f8a75..eb0c9ae19b 100644 --- a/nerfstudio/data/dataparsers/scannet_dataparser.py +++ b/nerfstudio/data/dataparsers/scannet_dataparser.py @@ -24,11 +24,7 @@ from nerfstudio.cameras import camera_utils from nerfstudio.cameras.cameras import Cameras, CameraType -from nerfstudio.data.dataparsers.base_dataparser import ( - DataParser, - DataParserConfig, - DataparserOutputs, -) +from nerfstudio.data.dataparsers.base_dataparser import DataParser, DataParserConfig, DataparserOutputs from nerfstudio.data.scene_box import SceneBox diff --git a/nerfstudio/data/dataparsers/scannetpp_dataparser.py b/nerfstudio/data/dataparsers/scannetpp_dataparser.py index bbd260730f..18afa4f423 100644 --- a/nerfstudio/data/dataparsers/scannetpp_dataparser.py +++ b/nerfstudio/data/dataparsers/scannetpp_dataparser.py @@ -21,6 +21,7 @@ import numpy as np import torch + from nerfstudio.cameras import camera_utils from nerfstudio.cameras.cameras import CAMERA_MODEL_TO_TYPE, Cameras, CameraType from nerfstudio.data.dataparsers.base_dataparser import DataParser, DataParserConfig, DataparserOutputs @@ -108,9 +109,7 @@ def _generate_dataparser_outputs(self, split="train"): else: i_train.append(idx) - assert len(mask_filenames) == 0 or ( - len(mask_filenames) == len(image_filenames) - ), """ + assert len(mask_filenames) == 0 or (len(mask_filenames) == len(image_filenames)), """ Different number of image and mask filenames. You should check that mask_path is specified for every frame (or zero frames) in transforms.json. """ diff --git a/nerfstudio/data/dataparsers/sdfstudio_dataparser.py b/nerfstudio/data/dataparsers/sdfstudio_dataparser.py index 430a8aca8b..540aa65d54 100644 --- a/nerfstudio/data/dataparsers/sdfstudio_dataparser.py +++ b/nerfstudio/data/dataparsers/sdfstudio_dataparser.py @@ -23,11 +23,7 @@ from nerfstudio.cameras import camera_utils from nerfstudio.cameras.cameras import Cameras, CameraType -from nerfstudio.data.dataparsers.base_dataparser import ( - DataParser, - DataParserConfig, - DataparserOutputs, -) +from nerfstudio.data.dataparsers.base_dataparser import DataParser, DataParserConfig, DataparserOutputs from nerfstudio.data.scene_box import SceneBox from nerfstudio.utils.io import load_from_json diff --git a/nerfstudio/data/dataparsers/sitcoms3d_dataparser.py b/nerfstudio/data/dataparsers/sitcoms3d_dataparser.py index 1d4bc73c36..221a76cf0c 100644 --- a/nerfstudio/data/dataparsers/sitcoms3d_dataparser.py +++ b/nerfstudio/data/dataparsers/sitcoms3d_dataparser.py @@ -27,12 +27,7 @@ import torch from nerfstudio.cameras.cameras import Cameras, CameraType -from nerfstudio.data.dataparsers.base_dataparser import ( - DataParser, - DataParserConfig, - DataparserOutputs, - Semantics, -) +from nerfstudio.data.dataparsers.base_dataparser import DataParser, DataParserConfig, DataparserOutputs, Semantics from nerfstudio.data.scene_box import SceneBox from nerfstudio.utils.io import load_from_json diff --git a/nerfstudio/data/datasets/depth_dataset.py b/nerfstudio/data/datasets/depth_dataset.py index e17ac63157..f2f1ab86ad 100644 --- a/nerfstudio/data/datasets/depth_dataset.py +++ b/nerfstudio/data/datasets/depth_dataset.py @@ -16,24 +16,22 @@ Depth dataset. """ -from typing import Dict +import json +from pathlib import Path +from typing import Dict, Union import numpy as np +import torch +from PIL import Image +from rich.progress import track from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs -from nerfstudio.model_components import losses from nerfstudio.data.datasets.base_dataset import InputDataset from nerfstudio.data.utils.data_utils import get_depth_image_from_path +from nerfstudio.model_components import losses from nerfstudio.utils.misc import torch_compile from nerfstudio.utils.rich_utils import CONSOLE -from typing import Union -from PIL import Image -import torch -from rich.progress import track -from pathlib import Path -import json - class DepthDataset(InputDataset): """Dataset that returns images and depths. If no depths are found, then we generate them with Zoe Depth. diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index c7ce9e073e..0fc58bde01 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -18,20 +18,13 @@ import random from dataclasses import dataclass, field -from typing import ( - Dict, - Optional, - Type, - Union, -) +from typing import Dict, Optional, Type, Union import torch from jaxtyping import Int from torch import Tensor -from nerfstudio.configs.base_config import ( - InstantiateConfig, -) +from nerfstudio.configs.base_config import InstantiateConfig from nerfstudio.data.utils.pixel_sampling_utils import erode_mask diff --git a/nerfstudio/data/scene_box.py b/nerfstudio/data/scene_box.py index 835c2b1cfd..bf36702045 100644 --- a/nerfstudio/data/scene_box.py +++ b/nerfstudio/data/scene_box.py @@ -17,10 +17,10 @@ """ from dataclasses import dataclass -from typing import Union, Tuple -import viser.transforms as vtf +from typing import Tuple, Union import torch +import viser.transforms as vtf from jaxtyping import Float from torch import Tensor diff --git a/nerfstudio/data/utils/pixel_sampling_utils.py b/nerfstudio/data/utils/pixel_sampling_utils.py index 464381f4cd..5e10cf3b67 100644 --- a/nerfstudio/data/utils/pixel_sampling_utils.py +++ b/nerfstudio/data/utils/pixel_sampling_utils.py @@ -15,8 +15,8 @@ """ Pixel sampling utils such as eroding of valid masks that we sample from. """ import torch -from torch import Tensor from jaxtyping import Float +from torch import Tensor def dilate(tensor: Float[Tensor, "bs 1 H W"], kernel_size=3) -> Float[Tensor, "bs 1 H W"]: diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index 35bb428441..c4b867f9f5 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -28,6 +28,11 @@ from typing import DefaultDict, Dict, List, Literal, Optional, Tuple, Type, cast import torch +from rich import box, style +from rich.panel import Panel +from rich.table import Table +from torch.cuda.amp.grad_scaler import GradScaler + from nerfstudio.configs.experiment_config import ExperimentConfig from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes, TrainingCallbackLocation from nerfstudio.engine.optimizers import Optimizers @@ -39,10 +44,6 @@ from nerfstudio.utils.writer import EventName, TimeWriter from nerfstudio.viewer.viewer import Viewer as ViewerState from nerfstudio.viewer_legacy.server.viewer_state import ViewerLegacyState -from rich import box, style -from rich.panel import Panel -from rich.table import Table -from torch.cuda.amp.grad_scaler import GradScaler TRAIN_INTERATION_OUTPUT = Tuple[torch.Tensor, Dict[str, torch.Tensor], Dict[str, torch.Tensor]] TORCH_DEVICE = str diff --git a/nerfstudio/exporter/exporter_utils.py b/nerfstudio/exporter/exporter_utils.py index 56b4b64456..45f6ef6866 100644 --- a/nerfstudio/exporter/exporter_utils.py +++ b/nerfstudio/exporter/exporter_utils.py @@ -28,11 +28,11 @@ import pymeshlab import torch from jaxtyping import Float -from nerfstudio.cameras.rays import RayBundle from rich.progress import BarColumn, Progress, TaskProgressColumn, TextColumn, TimeRemainingColumn from torch import Tensor from nerfstudio.cameras.cameras import Cameras +from nerfstudio.cameras.rays import RayBundle from nerfstudio.data.datasets.base_dataset import InputDataset from nerfstudio.data.scene_box import OrientedBox from nerfstudio.pipelines.base_pipeline import Pipeline, VanillaPipeline diff --git a/nerfstudio/exporter/marching_cubes.py b/nerfstudio/exporter/marching_cubes.py index ee13422b1f..672bd0a4c2 100644 --- a/nerfstudio/exporter/marching_cubes.py +++ b/nerfstudio/exporter/marching_cubes.py @@ -169,9 +169,7 @@ def generate_mesh_with_multires_marching_cubes( A torch tensor with the SDF values evaluated at the given points. """ # Check if resolution is divisible by 512 - assert ( - resolution % 512 == 0 - ), f"""resolution must be divisible by 512, got {resolution}. + assert resolution % 512 == 0, f"""resolution must be divisible by 512, got {resolution}. This is important because the algorithm uses a multi-resolution approach to evaluate the SDF where the mimimum resolution is 512.""" # Prepare coarse mask if provided diff --git a/nerfstudio/field_components/__init__.py b/nerfstudio/field_components/__init__.py index d509df2220..5f99dbcc8b 100644 --- a/nerfstudio/field_components/__init__.py +++ b/nerfstudio/field_components/__init__.py @@ -14,6 +14,5 @@ """init field modules""" from .base_field_component import FieldComponent as FieldComponent -from .encodings import Encoding as Encoding -from .encodings import ScalingAndOffset as ScalingAndOffset +from .encodings import Encoding as Encoding, ScalingAndOffset as ScalingAndOffset from .mlp import MLP as MLP diff --git a/nerfstudio/field_components/encodings.py b/nerfstudio/field_components/encodings.py index dc410673b9..8d63cd9279 100644 --- a/nerfstudio/field_components/encodings.py +++ b/nerfstudio/field_components/encodings.py @@ -28,11 +28,7 @@ from nerfstudio.field_components.base_field_component import FieldComponent from nerfstudio.utils.external import TCNN_EXISTS, tcnn -from nerfstudio.utils.math import ( - components_from_spherical_harmonics, - expected_sin, - generate_polyhedron_basis, -) +from nerfstudio.utils.math import components_from_spherical_harmonics, expected_sin, generate_polyhedron_basis from nerfstudio.utils.printing import print_tcnn_speed_warning diff --git a/nerfstudio/field_components/mlp.py b/nerfstudio/field_components/mlp.py index 94e6b1eb49..f1df2baa33 100644 --- a/nerfstudio/field_components/mlp.py +++ b/nerfstudio/field_components/mlp.py @@ -23,11 +23,10 @@ from torch import Tensor, nn from nerfstudio.field_components.base_field_component import FieldComponent -from nerfstudio.utils.printing import print_tcnn_speed_warning from nerfstudio.field_components.encodings import HashEncoding - -from nerfstudio.utils.rich_utils import CONSOLE from nerfstudio.utils.external import TCNN_EXISTS, tcnn +from nerfstudio.utils.printing import print_tcnn_speed_warning +from nerfstudio.utils.rich_utils import CONSOLE def activation_to_tcnn_string(activation: Union[nn.Module, None]) -> str: diff --git a/nerfstudio/fields/semantic_nerf_field.py b/nerfstudio/fields/semantic_nerf_field.py index 5681dd8621..62b75094f4 100644 --- a/nerfstudio/fields/semantic_nerf_field.py +++ b/nerfstudio/fields/semantic_nerf_field.py @@ -22,12 +22,7 @@ from nerfstudio.cameras.rays import RaySamples from nerfstudio.field_components.encodings import Encoding, Identity -from nerfstudio.field_components.field_heads import ( - DensityFieldHead, - FieldHeadNames, - RGBFieldHead, - SemanticFieldHead, -) +from nerfstudio.field_components.field_heads import DensityFieldHead, FieldHeadNames, RGBFieldHead, SemanticFieldHead from nerfstudio.field_components.mlp import MLP from nerfstudio.fields.base_field import Field diff --git a/nerfstudio/fields/vanilla_nerf_field.py b/nerfstudio/fields/vanilla_nerf_field.py index 22e11a25b8..efb3589ebe 100644 --- a/nerfstudio/fields/vanilla_nerf_field.py +++ b/nerfstudio/fields/vanilla_nerf_field.py @@ -22,12 +22,7 @@ from nerfstudio.cameras.rays import RaySamples from nerfstudio.field_components.encodings import Encoding, Identity -from nerfstudio.field_components.field_heads import ( - DensityFieldHead, - FieldHead, - FieldHeadNames, - RGBFieldHead, -) +from nerfstudio.field_components.field_heads import DensityFieldHead, FieldHead, FieldHeadNames, RGBFieldHead from nerfstudio.field_components.mlp import MLP from nerfstudio.field_components.spatial_distortions import SpatialDistortion from nerfstudio.fields.base_field import Field diff --git a/nerfstudio/generative/deepfloyd.py b/nerfstudio/generative/deepfloyd.py index 72c6cc3730..03c02b5cdd 100644 --- a/nerfstudio/generative/deepfloyd.py +++ b/nerfstudio/generative/deepfloyd.py @@ -19,21 +19,16 @@ import torch import torch.nn.functional as F import tyro - from jaxtyping import Float from PIL import Image from torch import Generator, Tensor, nn from torch.cuda.amp.grad_scaler import GradScaler - from nerfstudio.generative.utils import CatchMissingPackages try: - from diffusers import IFPipeline as IFOrig - from diffusers.pipelines.deepfloyd_if import IFPipelineOutput as IFOutputOrig - - from diffusers import IFPipeline, DiffusionPipeline - from diffusers.pipelines.deepfloyd_if import IFPipelineOutput + from diffusers import DiffusionPipeline, IFPipeline, IFPipeline as IFOrig + from diffusers.pipelines.deepfloyd_if import IFPipelineOutput, IFPipelineOutput as IFOutputOrig from transformers import T5EncoderModel except ImportError: diff --git a/nerfstudio/generative/positional_text_embeddings.py b/nerfstudio/generative/positional_text_embeddings.py index 4ba62b024b..38cd3a52eb 100644 --- a/nerfstudio/generative/positional_text_embeddings.py +++ b/nerfstudio/generative/positional_text_embeddings.py @@ -14,14 +14,15 @@ """Utility helper functions for generative 3D models""" +from typing import Union + import torch -from torch import Tensor from jaxtyping import Float +from torch import Tensor from typing_extensions import Literal -from typing import Union -from nerfstudio.generative.stable_diffusion import StableDiffusion from nerfstudio.generative.deepfloyd import DeepFloyd +from nerfstudio.generative.stable_diffusion import StableDiffusion class PositionalTextEmbeddings: diff --git a/nerfstudio/generative/stable_diffusion.py b/nerfstudio/generative/stable_diffusion.py index 7e7f3dac44..7bf0fb96d8 100644 --- a/nerfstudio/generative/stable_diffusion.py +++ b/nerfstudio/generative/stable_diffusion.py @@ -18,7 +18,6 @@ from pathlib import Path from typing import List, Optional, Union -from nerfstudio.utils.rich_utils import CONSOLE import mediapy import numpy as np @@ -30,10 +29,10 @@ from torch.cuda.amp.grad_scaler import GradScaler from nerfstudio.generative.utils import CatchMissingPackages - +from nerfstudio.utils.rich_utils import CONSOLE try: - from diffusers import PNDMScheduler, StableDiffusionPipeline, DiffusionPipeline + from diffusers import DiffusionPipeline, PNDMScheduler, StableDiffusionPipeline except ImportError: PNDMScheduler = StableDiffusionPipeline = CatchMissingPackages() diff --git a/nerfstudio/generative/utils.py b/nerfstudio/generative/utils.py index ec5a728441..5e8ed02e58 100644 --- a/nerfstudio/generative/utils.py +++ b/nerfstudio/generative/utils.py @@ -15,6 +15,7 @@ """Utility helper functions for diffusion models""" import sys + from nerfstudio.utils.rich_utils import CONSOLE diff --git a/nerfstudio/models/base_model.py b/nerfstudio/models/base_model.py index 2a434a85e8..e1c9507ff0 100644 --- a/nerfstudio/models/base_model.py +++ b/nerfstudio/models/base_model.py @@ -31,7 +31,7 @@ from nerfstudio.cameras.rays import RayBundle from nerfstudio.configs.base_config import InstantiateConfig from nerfstudio.configs.config_utils import to_immutable_dict -from nerfstudio.data.scene_box import SceneBox, OrientedBox +from nerfstudio.data.scene_box import OrientedBox, SceneBox from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes from nerfstudio.model_components.scene_colliders import NearFarCollider diff --git a/nerfstudio/models/depth_nerfacto.py b/nerfstudio/models/depth_nerfacto.py index d925c0516a..e439ce6777 100644 --- a/nerfstudio/models/depth_nerfacto.py +++ b/nerfstudio/models/depth_nerfacto.py @@ -21,8 +21,8 @@ from dataclasses import dataclass, field from typing import Dict, Tuple, Type -import torch import numpy as np +import torch from nerfstudio.cameras.rays import RayBundle from nerfstudio.model_components import losses diff --git a/nerfstudio/models/gaussian_splatting.py b/nerfstudio/models/gaussian_splatting.py index 7edc064918..60563a96ae 100644 --- a/nerfstudio/models/gaussian_splatting.py +++ b/nerfstudio/models/gaussian_splatting.py @@ -719,9 +719,7 @@ def get_outputs(self, camera: Cameras) -> Dict[str, Union[torch.Tensor, List]]: H, W, torch.ones(3, device=self.device) * 10, - )[ - ..., 0:1 - ] # type: ignore + )[..., 0:1] # type: ignore return {"rgb": rgb, "depth": depth_im} # type: ignore diff --git a/nerfstudio/models/generfacto.py b/nerfstudio/models/generfacto.py index c07f904c7d..5c1316ab79 100644 --- a/nerfstudio/models/generfacto.py +++ b/nerfstudio/models/generfacto.py @@ -27,33 +27,16 @@ from typing_extensions import Literal from nerfstudio.cameras.rays import RayBundle -from nerfstudio.engine.callbacks import ( - TrainingCallback, - TrainingCallbackAttributes, - TrainingCallbackLocation, -) +from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes, TrainingCallbackLocation from nerfstudio.field_components.field_heads import FieldHeadNames from nerfstudio.fields.density_fields import HashMLPDensityField from nerfstudio.fields.generfacto_field import GenerfactoField -from nerfstudio.generative.stable_diffusion import StableDiffusion from nerfstudio.generative.deepfloyd import DeepFloyd from nerfstudio.generative.positional_text_embeddings import PositionalTextEmbeddings -from nerfstudio.model_components.losses import ( - MSELoss, - distortion_loss, - interlevel_loss, - orientation_loss, -) -from nerfstudio.model_components.ray_samplers import ( - ProposalNetworkSampler, - UniformSampler, -) -from nerfstudio.model_components.renderers import ( - AccumulationRenderer, - DepthRenderer, - NormalsRenderer, - RGBRenderer, -) +from nerfstudio.generative.stable_diffusion import StableDiffusion +from nerfstudio.model_components.losses import MSELoss, distortion_loss, interlevel_loss, orientation_loss +from nerfstudio.model_components.ray_samplers import ProposalNetworkSampler, UniformSampler +from nerfstudio.model_components.renderers import AccumulationRenderer, DepthRenderer, NormalsRenderer, RGBRenderer from nerfstudio.model_components.scene_colliders import AABBBoxCollider, SphereCollider from nerfstudio.model_components.shaders import LambertianShader, NormalsShader from nerfstudio.models.base_model import Model, ModelConfig @@ -275,22 +258,30 @@ def get_training_callbacks( ) -> List[TrainingCallback]: # the callback that we want to run every X iterations after the training iteration def taper_density( - self, training_callback_attributes: TrainingCallbackAttributes, step: int # pylint: disable=unused-argument + self, + training_callback_attributes: TrainingCallbackAttributes, + step: int, # pylint: disable=unused-argument ): self.density_strength = np.interp(step, self.config.taper_range, self.config.taper_strength) def start_training_normals( - self, training_callback_attributes: TrainingCallbackAttributes, step: int # pylint: disable=unused-argument + self, + training_callback_attributes: TrainingCallbackAttributes, + step: int, # pylint: disable=unused-argument ): self.train_normals = True def start_shaded_training( - self, training_callback_attributes: TrainingCallbackAttributes, step: int # pylint: disable=unused-argument + self, + training_callback_attributes: TrainingCallbackAttributes, + step: int, # pylint: disable=unused-argument ): self.train_shaded = True def update_orientation_loss_mult( - self, training_callback_attributes: TrainingCallbackAttributes, step: int # pylint: disable=unused-argument + self, + training_callback_attributes: TrainingCallbackAttributes, + step: int, # pylint: disable=unused-argument ): if step <= self.config.start_normals_training: self.orientation_loss_mult = 0 diff --git a/nerfstudio/models/instant_ngp.py b/nerfstudio/models/instant_ngp.py index 48cf1a95ab..7dbcc30ccf 100644 --- a/nerfstudio/models/instant_ngp.py +++ b/nerfstudio/models/instant_ngp.py @@ -29,21 +29,13 @@ from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity from nerfstudio.cameras.rays import RayBundle -from nerfstudio.engine.callbacks import ( - TrainingCallback, - TrainingCallbackAttributes, - TrainingCallbackLocation, -) +from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes, TrainingCallbackLocation from nerfstudio.field_components.field_heads import FieldHeadNames from nerfstudio.field_components.spatial_distortions import SceneContraction from nerfstudio.fields.nerfacto_field import NerfactoField from nerfstudio.model_components.losses import MSELoss, scale_gradients_by_distance_squared from nerfstudio.model_components.ray_samplers import VolumetricSampler -from nerfstudio.model_components.renderers import ( - AccumulationRenderer, - DepthRenderer, - RGBRenderer, -) +from nerfstudio.model_components.renderers import AccumulationRenderer, DepthRenderer, RGBRenderer from nerfstudio.models.base_model import Model, ModelConfig from nerfstudio.utils import colormaps diff --git a/nerfstudio/models/mipnerf.py b/nerfstudio/models/mipnerf.py index cf48b08a75..b7e792100f 100644 --- a/nerfstudio/models/mipnerf.py +++ b/nerfstudio/models/mipnerf.py @@ -31,11 +31,7 @@ from nerfstudio.fields.vanilla_nerf_field import NeRFField from nerfstudio.model_components.losses import MSELoss, scale_gradients_by_distance_squared from nerfstudio.model_components.ray_samplers import PDFSampler, UniformSampler -from nerfstudio.model_components.renderers import ( - AccumulationRenderer, - DepthRenderer, - RGBRenderer, -) +from nerfstudio.model_components.renderers import AccumulationRenderer, DepthRenderer, RGBRenderer from nerfstudio.models.base_model import Model from nerfstudio.models.vanilla_nerf import VanillaModelConfig from nerfstudio.utils import colormaps, misc diff --git a/nerfstudio/models/neus.py b/nerfstudio/models/neus.py index f46cc9839d..d754c63b48 100644 --- a/nerfstudio/models/neus.py +++ b/nerfstudio/models/neus.py @@ -22,11 +22,7 @@ from typing import Dict, List, Type from nerfstudio.cameras.rays import RayBundle -from nerfstudio.engine.callbacks import ( - TrainingCallback, - TrainingCallbackAttributes, - TrainingCallbackLocation, -) +from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes, TrainingCallbackLocation from nerfstudio.field_components.field_heads import FieldHeadNames from nerfstudio.model_components.ray_samplers import NeuSSampler from nerfstudio.models.base_surface_model import SurfaceModel, SurfaceModelConfig diff --git a/nerfstudio/models/neus_facto.py b/nerfstudio/models/neus_facto.py index 5c1a3ec9d9..39e4e7fedd 100644 --- a/nerfstudio/models/neus_facto.py +++ b/nerfstudio/models/neus_facto.py @@ -27,18 +27,11 @@ from torch.nn import Parameter from nerfstudio.cameras.rays import RayBundle -from nerfstudio.engine.callbacks import ( - TrainingCallback, - TrainingCallbackAttributes, - TrainingCallbackLocation, -) +from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes, TrainingCallbackLocation from nerfstudio.field_components.field_heads import FieldHeadNames from nerfstudio.fields.density_fields import HashMLPDensityField from nerfstudio.model_components.losses import interlevel_loss -from nerfstudio.model_components.ray_samplers import ( - ProposalNetworkSampler, - UniformSampler, -) +from nerfstudio.model_components.ray_samplers import ProposalNetworkSampler, UniformSampler from nerfstudio.models.neus import NeuSModel, NeuSModelConfig from nerfstudio.utils import colormaps diff --git a/nerfstudio/models/semantic_nerfw.py b/nerfstudio/models/semantic_nerfw.py index 87388d0e94..8d5be08497 100644 --- a/nerfstudio/models/semantic_nerfw.py +++ b/nerfstudio/models/semantic_nerfw.py @@ -30,11 +30,7 @@ from nerfstudio.cameras.rays import RayBundle from nerfstudio.data.dataparsers.base_dataparser import Semantics -from nerfstudio.engine.callbacks import ( - TrainingCallback, - TrainingCallbackAttributes, - TrainingCallbackLocation, -) +from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes, TrainingCallbackLocation from nerfstudio.field_components.field_heads import FieldHeadNames from nerfstudio.field_components.spatial_distortions import SceneContraction from nerfstudio.fields.density_fields import HashMLPDensityField diff --git a/nerfstudio/models/tensorf.py b/nerfstudio/models/tensorf.py index 0ca56f10de..8827f896f3 100644 --- a/nerfstudio/models/tensorf.py +++ b/nerfstudio/models/tensorf.py @@ -28,32 +28,19 @@ from torchmetrics.image import PeakSignalNoiseRatio from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity +from nerfstudio.cameras.camera_optimizers import CameraOptimizer, CameraOptimizerConfig from nerfstudio.cameras.rays import RayBundle from nerfstudio.configs.config_utils import to_immutable_dict -from nerfstudio.engine.callbacks import ( - TrainingCallback, - TrainingCallbackAttributes, - TrainingCallbackLocation, -) -from nerfstudio.field_components.encodings import ( - NeRFEncoding, - TensorCPEncoding, - TensorVMEncoding, - TriplaneEncoding, -) +from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes, TrainingCallbackLocation +from nerfstudio.field_components.encodings import NeRFEncoding, TensorCPEncoding, TensorVMEncoding, TriplaneEncoding from nerfstudio.field_components.field_heads import FieldHeadNames from nerfstudio.fields.tensorf_field import TensoRFField -from nerfstudio.model_components.losses import MSELoss, tv_loss, scale_gradients_by_distance_squared +from nerfstudio.model_components.losses import MSELoss, scale_gradients_by_distance_squared, tv_loss from nerfstudio.model_components.ray_samplers import PDFSampler, UniformSampler -from nerfstudio.model_components.renderers import ( - AccumulationRenderer, - DepthRenderer, - RGBRenderer, -) +from nerfstudio.model_components.renderers import AccumulationRenderer, DepthRenderer, RGBRenderer from nerfstudio.model_components.scene_colliders import AABBBoxCollider from nerfstudio.models.base_model import Model, ModelConfig from nerfstudio.utils import colormaps, colors, misc -from nerfstudio.cameras.camera_optimizers import CameraOptimizer, CameraOptimizerConfig @dataclass diff --git a/nerfstudio/models/vanilla_nerf.py b/nerfstudio/models/vanilla_nerf.py index 4f4dbd8f61..795bbb23c8 100644 --- a/nerfstudio/models/vanilla_nerf.py +++ b/nerfstudio/models/vanilla_nerf.py @@ -19,7 +19,7 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Any, Dict, List, Tuple, Type, Literal +from typing import Any, Dict, List, Literal, Tuple, Type import torch from torch.nn import Parameter @@ -35,11 +35,7 @@ from nerfstudio.fields.vanilla_nerf_field import NeRFField from nerfstudio.model_components.losses import MSELoss, scale_gradients_by_distance_squared from nerfstudio.model_components.ray_samplers import PDFSampler, UniformSampler -from nerfstudio.model_components.renderers import ( - AccumulationRenderer, - DepthRenderer, - RGBRenderer, -) +from nerfstudio.model_components.renderers import AccumulationRenderer, DepthRenderer, RGBRenderer from nerfstudio.models.base_model import Model, ModelConfig from nerfstudio.utils import colormaps, misc diff --git a/nerfstudio/pipelines/dynamic_batch.py b/nerfstudio/pipelines/dynamic_batch.py index 65679d0012..2d0eb37ce7 100644 --- a/nerfstudio/pipelines/dynamic_batch.py +++ b/nerfstudio/pipelines/dynamic_batch.py @@ -17,7 +17,7 @@ """ from dataclasses import dataclass, field -from typing import Literal, Type, Optional +from typing import Literal, Optional, Type import torch from torch.cuda.amp.grad_scaler import GradScaler diff --git a/nerfstudio/plugins/registry_dataparser.py b/nerfstudio/plugins/registry_dataparser.py index e35aeaf48b..6b1e34a0f3 100644 --- a/nerfstudio/plugins/registry_dataparser.py +++ b/nerfstudio/plugins/registry_dataparser.py @@ -16,8 +16,8 @@ Module that keeps all registered plugins and allows for plugin discovery. """ -import os import importlib +import os import sys import typing as t from dataclasses import dataclass diff --git a/nerfstudio/process_data/odm_utils.py b/nerfstudio/process_data/odm_utils.py index aa9b695410..8b8d69364d 100644 --- a/nerfstudio/process_data/odm_utils.py +++ b/nerfstudio/process_data/odm_utils.py @@ -15,11 +15,11 @@ """Helper utils for processing ODM data into the nerfstudio format.""" import json -from pathlib import Path -from typing import Dict, List +import math import os import sys -import math +from pathlib import Path +from typing import Dict, List import numpy as np diff --git a/nerfstudio/process_data/process_data_utils.py b/nerfstudio/process_data/process_data_utils.py index c1946305b3..768786996b 100644 --- a/nerfstudio/process_data/process_data_utils.py +++ b/nerfstudio/process_data/process_data_utils.py @@ -31,6 +31,7 @@ import newrawpy as rawpy # type: ignore import numpy as np + from nerfstudio.utils.rich_utils import CONSOLE, status from nerfstudio.utils.scripts import run_command diff --git a/nerfstudio/scripts/blender/nerfstudio_blender.py b/nerfstudio/scripts/blender/nerfstudio_blender.py index c5f6e32515..c079a6b2a4 100644 --- a/nerfstudio/scripts/blender/nerfstudio_blender.py +++ b/nerfstudio/scripts/blender/nerfstudio_blender.py @@ -135,7 +135,7 @@ def construct_json_obj(self): render_seconds = 1 / render_fps else: render_seconds = ( - ((bpy.context.scene.frame_end - bpy.context.scene.frame_start) // (bpy.context.scene.frame_step) + 1) + (bpy.context.scene.frame_end - bpy.context.scene.frame_start) // (bpy.context.scene.frame_step) + 1 ) / render_fps smoothness_value = 0 diff --git a/nerfstudio/scripts/completions/install.py b/nerfstudio/scripts/completions/install.py index d6d8b63b22..67e0189d5f 100644 --- a/nerfstudio/scripts/completions/install.py +++ b/nerfstudio/scripts/completions/install.py @@ -24,8 +24,7 @@ import stat import subprocess import sys -from typing import List, Literal, Optional, Union -from typing import get_args as typing_get_args +from typing import List, Literal, Optional, Union, get_args as typing_get_args import tyro from rich.prompt import Confirm diff --git a/nerfstudio/scripts/downloads/download_data.py b/nerfstudio/scripts/downloads/download_data.py index 9560beba8d..ec4ce9b158 100644 --- a/nerfstudio/scripts/downloads/download_data.py +++ b/nerfstudio/scripts/downloads/download_data.py @@ -27,10 +27,10 @@ import gdown import torch import tyro -from nerfstudio.process_data import process_data_utils from typing_extensions import Annotated from nerfstudio.configs.base_config import PrintableConfig +from nerfstudio.process_data import process_data_utils from nerfstudio.utils import install_checks from nerfstudio.utils.scripts import run_command diff --git a/nerfstudio/scripts/exporter.py b/nerfstudio/scripts/exporter.py index 1eef27dbf7..be4ad08248 100644 --- a/nerfstudio/scripts/exporter.py +++ b/nerfstudio/scripts/exporter.py @@ -122,7 +122,7 @@ class ExportPointCloud(Exporter): std_ratio: float = 10.0 """Threshold based on STD of the average distances across the point cloud to remove outliers.""" save_world_frame: bool = True - """If true, saves in the frame of the transform.json file, if false saves in the frame of the scaled + """If true, saves in the frame of the transform.json file, if false saves in the frame of the scaled dataparser transform""" def main(self) -> None: @@ -417,9 +417,7 @@ def main(self) -> None: CONSOLE.print("Extracting mesh with marching cubes... which may take a while") - assert ( - self.resolution % 512 == 0 - ), f"""resolution must be divisible by 512, got {self.resolution}. + assert self.resolution % 512 == 0, f"""resolution must be divisible by 512, got {self.resolution}. This is important because the algorithm uses a multi-resolution approach to evaluate the SDF where the minimum resolution is 512.""" diff --git a/nerfstudio/scripts/github/run_actions.py b/nerfstudio/scripts/github/run_actions.py index 1081cee571..5a966338c1 100644 --- a/nerfstudio/scripts/github/run_actions.py +++ b/nerfstudio/scripts/github/run_actions.py @@ -23,7 +23,7 @@ from nerfstudio.utils.rich_utils import CONSOLE -LOCAL_TESTS = ["Run license checks", "Run Ruff", "Run Black", "Run Pyright", "Test with pytest"] +LOCAL_TESTS = ["Run license checks", "Run Ruff Linter", "Run Ruff Formatter", "Run Pyright", "Test with pytest"] def run_command(command: str, continue_on_fail: bool = False) -> bool: @@ -57,7 +57,7 @@ def run_github_actions_file(filename: str, continue_on_fail: bool = False): for step in steps: if "name" in step and step["name"] in LOCAL_TESTS: compressed = step["run"].replace("\n", ";").replace("\\", "") - if "ruff" in compressed: + if "ruff check" in compressed: curr_command = f"{compressed} --fix" else: curr_command = compressed.replace("--check", "") diff --git a/nerfstudio/scripts/process_data.py b/nerfstudio/scripts/process_data.py index 2c2cd7a381..6890cd933a 100644 --- a/nerfstudio/scripts/process_data.py +++ b/nerfstudio/scripts/process_data.py @@ -28,11 +28,11 @@ from nerfstudio.process_data import ( metashape_utils, + odm_utils, polycam_utils, process_data_utils, realitycapture_utils, record3d_utils, - odm_utils, ) from nerfstudio.process_data.colmap_converter_to_nerfstudio_dataset import BaseConverterToNerfstudioDataset from nerfstudio.process_data.images_to_nerfstudio_dataset import ImagesToNerfstudioDataset diff --git a/nerfstudio/scripts/render.py b/nerfstudio/scripts/render.py index 114e4195ec..bcf050df21 100644 --- a/nerfstudio/scripts/render.py +++ b/nerfstudio/scripts/render.py @@ -33,38 +33,20 @@ import numpy as np import torch import tyro +import viser.transforms as tf from jaxtyping import Float from rich import box, style from rich.panel import Panel -from rich.progress import ( - BarColumn, - Progress, - TaskProgressColumn, - TextColumn, - TimeElapsedColumn, - TimeRemainingColumn, -) +from rich.progress import BarColumn, Progress, TaskProgressColumn, TextColumn, TimeElapsedColumn, TimeRemainingColumn from rich.table import Table from torch import Tensor from typing_extensions import Annotated -import viser.transforms as tf - -from nerfstudio.cameras.camera_paths import ( - get_interpolated_camera_path, - get_path_from_json, - get_spiral_path, -) - +from nerfstudio.cameras.camera_paths import get_interpolated_camera_path, get_path_from_json, get_spiral_path from nerfstudio.cameras.cameras import Cameras, CameraType, RayBundle -from nerfstudio.data.datamanagers.base_datamanager import ( - VanillaDataManager, - VanillaDataManagerConfig, -) +from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager, VanillaDataManagerConfig from nerfstudio.data.datamanagers.parallel_datamanager import ParallelDataManager -from nerfstudio.data.datamanagers.random_cameras_datamanager import ( - RandomCamerasDataManager, -) +from nerfstudio.data.datamanagers.random_cameras_datamanager import RandomCamerasDataManager from nerfstudio.data.datasets.base_dataset import Dataset from nerfstudio.data.scene_box import OrientedBox from nerfstudio.data.utils.dataloaders import FixedIndicesEvalDataloader diff --git a/nerfstudio/scripts/viewer/run_viewer.py b/nerfstudio/scripts/viewer/run_viewer.py index f071478d50..63c2007fbd 100644 --- a/nerfstudio/scripts/viewer/run_viewer.py +++ b/nerfstudio/scripts/viewer/run_viewer.py @@ -30,8 +30,8 @@ from nerfstudio.pipelines.base_pipeline import Pipeline from nerfstudio.utils import writer from nerfstudio.utils.eval_utils import eval_setup -from nerfstudio.viewer_legacy.server.viewer_state import ViewerLegacyState from nerfstudio.viewer.viewer import Viewer as ViewerState +from nerfstudio.viewer_legacy.server.viewer_state import ViewerLegacyState @dataclass diff --git a/nerfstudio/utils/eval_utils.py b/nerfstudio/utils/eval_utils.py index 0b7306791c..2c4bfcd154 100644 --- a/nerfstudio/utils/eval_utils.py +++ b/nerfstudio/utils/eval_utils.py @@ -20,7 +20,7 @@ import os import sys from pathlib import Path -from typing import Literal, Optional, Tuple, Callable +from typing import Callable, Literal, Optional, Tuple import torch import yaml diff --git a/nerfstudio/utils/misc.py b/nerfstudio/utils/misc.py index e67f63fc94..ae3c4d5573 100644 --- a/nerfstudio/utils/misc.py +++ b/nerfstudio/utils/misc.py @@ -17,11 +17,11 @@ """ -from inspect import currentframe -import typing import platform -from typing import Any, Callable, Dict, List, Optional, TypeVar, Union +import typing import warnings +from inspect import currentframe +from typing import Any, Callable, Dict, List, Optional, TypeVar, Union import torch diff --git a/nerfstudio/utils/profiler.py b/nerfstudio/utils/profiler.py index 4bfac28f0e..ff0ae5bb93 100644 --- a/nerfstudio/utils/profiler.py +++ b/nerfstudio/utils/profiler.py @@ -24,28 +24,13 @@ from collections import deque from contextlib import ContextDecorator, contextmanager from pathlib import Path -from typing import ( - Any, - Callable, - ContextManager, - Dict, - List, - Optional, - Tuple, - TypeVar, - Union, - overload, -) +from typing import Any, Callable, ContextManager, Dict, List, Optional, Tuple, TypeVar, Union, overload from torch.profiler import ProfilerActivity, profile, record_function from nerfstudio.configs import base_config as cfg from nerfstudio.utils import comms -from nerfstudio.utils.decorators import ( - check_main_thread, - check_profiler_enabled, - decorate_all, -) +from nerfstudio.utils.decorators import check_main_thread, check_profiler_enabled, decorate_all from nerfstudio.utils.rich_utils import CONSOLE PROFILER = [] diff --git a/nerfstudio/utils/rich_utils.py b/nerfstudio/utils/rich_utils.py index cb9fa23fee..6d4e7557ec 100644 --- a/nerfstudio/utils/rich_utils.py +++ b/nerfstudio/utils/rich_utils.py @@ -18,15 +18,7 @@ from typing import Optional from rich.console import Console -from rich.progress import ( - BarColumn, - Progress, - ProgressColumn, - Task, - TaskProgressColumn, - TextColumn, - TimeRemainingColumn, -) +from rich.progress import BarColumn, Progress, ProgressColumn, Task, TaskProgressColumn, TextColumn, TimeRemainingColumn from rich.text import Text CONSOLE = Console(width=120) diff --git a/nerfstudio/utils/writer.py b/nerfstudio/utils/writer.py index aacddbe4a9..2bc5bcfc60 100644 --- a/nerfstudio/utils/writer.py +++ b/nerfstudio/utils/writer.py @@ -28,12 +28,13 @@ import torch import wandb from jaxtyping import Float +from torch import Tensor +from torch.utils.tensorboard import SummaryWriter + from nerfstudio.configs import base_config as cfg from nerfstudio.utils.decorators import check_main_thread, decorate_all from nerfstudio.utils.printing import human_format from nerfstudio.utils.rich_utils import CONSOLE -from torch import Tensor -from torch.utils.tensorboard import SummaryWriter def to8b(x): diff --git a/nerfstudio/viewer/control_panel.py b/nerfstudio/viewer/control_panel.py index 7eba58f9ee..6524c53093 100644 --- a/nerfstudio/viewer/control_panel.py +++ b/nerfstudio/viewer/control_panel.py @@ -20,6 +20,8 @@ import torch import viser import viser.transforms as vtf +from viser import ViserServer + from nerfstudio.data.scene_box import OrientedBox from nerfstudio.utils.colormaps import ColormapOptions, Colormaps from nerfstudio.viewer.viewer_elements import ( # ViewerButtonGroup, @@ -32,7 +34,6 @@ ViewerSlider, ViewerVec3, ) -from viser import ViserServer class ControlPanel: diff --git a/nerfstudio/viewer/export_panel.py b/nerfstudio/viewer/export_panel.py index ceff4e41f0..d003076f1a 100644 --- a/nerfstudio/viewer/export_panel.py +++ b/nerfstudio/viewer/export_panel.py @@ -18,9 +18,10 @@ import viser import viser.transforms as vtf +from typing_extensions import Literal + from nerfstudio.data.scene_box import OrientedBox from nerfstudio.viewer.control_panel import ControlPanel -from typing_extensions import Literal def populate_export_tab(server: viser.ViserServer, control_panel: ControlPanel, config_path: Path) -> None: diff --git a/nerfstudio/viewer/render_panel.py b/nerfstudio/viewer/render_panel.py index 735499bd13..f5c06ad19a 100644 --- a/nerfstudio/viewer/render_panel.py +++ b/nerfstudio/viewer/render_panel.py @@ -20,15 +20,17 @@ import json import threading import time -from typing import Dict, List, Optional, Tuple, Literal from pathlib import Path -from nerfstudio.viewer.control_panel import ControlPanel +from typing import Dict, List, Literal, Optional, Tuple + +import numpy as np +import scipy import splines import splines.quaternion import viser import viser.transforms as tf -import numpy as np -import scipy + +from nerfstudio.viewer.control_panel import ControlPanel @dataclasses.dataclass diff --git a/nerfstudio/viewer/render_state_machine.py b/nerfstudio/viewer/render_state_machine.py index 05bd658bc6..74fde38336 100644 --- a/nerfstudio/viewer/render_state_machine.py +++ b/nerfstudio/viewer/render_state_machine.py @@ -24,6 +24,8 @@ import numpy as np import torch import torch.nn.functional as F +from viser import ClientHandle + from nerfstudio.cameras.cameras import Cameras from nerfstudio.model_components.renderers import background_color_override_context from nerfstudio.models.gaussian_splatting import GaussianSplattingModel @@ -32,8 +34,6 @@ from nerfstudio.viewer.utils import CameraState, get_camera from nerfstudio.viewer_legacy.server import viewer_utils -from viser import ClientHandle - if TYPE_CHECKING: from nerfstudio.viewer.viewer import Viewer diff --git a/nerfstudio/viewer/utils.py b/nerfstudio/viewer/utils.py index 8b9e237938..b58423be33 100644 --- a/nerfstudio/viewer/utils.py +++ b/nerfstudio/viewer/utils.py @@ -20,10 +20,11 @@ import numpy as np import torch from jaxtyping import Float +from torch import nn + from nerfstudio.cameras.cameras import Cameras, CameraType from nerfstudio.data.scene_box import SceneBox from nerfstudio.models.base_model import Model -from torch import nn @dataclass diff --git a/nerfstudio/viewer/viewer.py b/nerfstudio/viewer/viewer.py index 00d1925529..56472d5366 100644 --- a/nerfstudio/viewer/viewer.py +++ b/nerfstudio/viewer/viewer.py @@ -23,8 +23,11 @@ import numpy as np import torch import torchvision +import viser import viser.theme import viser.transforms as vtf +from typing_extensions import assert_never + from nerfstudio.cameras.camera_optimizers import CameraOptimizer from nerfstudio.cameras.cameras import CameraType from nerfstudio.configs import base_config as cfg @@ -40,9 +43,6 @@ from nerfstudio.viewer.utils import CameraState, parse_object from nerfstudio.viewer.viewer_elements import ViewerControl, ViewerElement from nerfstudio.viewer_legacy.server import viewer_utils -from typing_extensions import assert_never - -import viser if TYPE_CHECKING: from nerfstudio.engine.trainer import Trainer diff --git a/nerfstudio/viewer/viewer_elements.py b/nerfstudio/viewer/viewer_elements.py index 2a8b64efe0..c9851acb32 100644 --- a/nerfstudio/viewer/viewer_elements.py +++ b/nerfstudio/viewer/viewer_elements.py @@ -35,8 +35,8 @@ ViserServer, ) -from nerfstudio.viewer.utils import CameraState, get_camera from nerfstudio.cameras.cameras import Cameras, CameraType +from nerfstudio.viewer.utils import CameraState, get_camera if TYPE_CHECKING: from nerfstudio.viewer.viewer import Viewer diff --git a/nerfstudio/viewer_legacy/viser/__init__.py b/nerfstudio/viewer_legacy/viser/__init__.py index 683b145c35..7560a4cc80 100644 --- a/nerfstudio/viewer_legacy/viser/__init__.py +++ b/nerfstudio/viewer_legacy/viser/__init__.py @@ -15,7 +15,6 @@ """ Viser is used for the nerfstudio viewer backend """ -from .message_api import GuiHandle as GuiHandle -from .message_api import GuiSelectHandle as GuiSelectHandle +from .message_api import GuiHandle as GuiHandle, GuiSelectHandle as GuiSelectHandle from .messages import NerfstudioMessage as NerfstudioMessage from .server import ViserServer as ViserServer diff --git a/nerfstudio/viewer_legacy/viser/gui.py b/nerfstudio/viewer_legacy/viser/gui.py index 4b5ef0a733..75c18decae 100644 --- a/nerfstudio/viewer_legacy/viser/gui.py +++ b/nerfstudio/viewer_legacy/viser/gui.py @@ -21,28 +21,12 @@ import dataclasses import time -from typing import ( - TYPE_CHECKING, - Any, - Callable, - Dict, - Generic, - List, - Optional, - Type, - TypeVar, - Union, -) +from typing import TYPE_CHECKING, Any, Callable, Dict, Generic, List, Optional, Type, TypeVar, Union import numpy as onp from viser.infra import ClientId -from .messages import ( - GuiRemoveMessage, - GuiSetHiddenMessage, - GuiSetLevaConfMessage, - GuiSetValueMessage, -) +from .messages import GuiRemoveMessage, GuiSetHiddenMessage, GuiSetLevaConfMessage, GuiSetValueMessage if TYPE_CHECKING: from .message_api import MessageApi diff --git a/pyproject.toml b/pyproject.toml index db899745a9..de6bf497a3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -85,12 +85,11 @@ gen = [ # Development packages dev = [ - "black[jupyter]==23.3.0", "pre-commit==3.3.2", "pytest==7.1.2", "pytest-xdist==2.5.0", "typeguard==2.13.3", - "ruff==0.0.267", + "ruff==0.1.13", "sshconf==0.2.5", "pycolmap>=0.3.0", # NOTE: pycolmap==0.3.0 is not available on newer python versions "diffusers==0.16.1", @@ -142,9 +141,6 @@ include = ["nerfstudio*"] [tool.setuptools.package-data] "*" = ["*.cu", "*.json", "py.typed", "setup.bash", "setup.zsh"] -[tool.black] -line-length = 120 - [tool.pytest.ini_options] addopts = "-n=4 --typeguard-packages=nerfstudio --disable-warnings" testpaths = [ @@ -167,9 +163,11 @@ pythonPlatform = "Linux" [tool.ruff] line-length = 120 +respect-gitignore = false select = [ "E", # pycodestyle errors. "F", # Pyflakes rules. + "I", # isort formatting. "PLC", # Pylint convention warnings. "PLE", # Pylint errors. "PLR", # Pylint refactor recommendations. @@ -190,3 +188,8 @@ ignore = [ "PLW0603", # Globa statement updates are discouraged. "PLW2901", # For loop variable overwritten. ] + +[tool.ruff.lint.isort] +combine-as-imports = true +known-first-party = ["nerfstudio"] +split-on-trailing-comma = false diff --git a/tests/dataparsers/test_nerfstudio_dataparser.py b/tests/dataparsers/test_nerfstudio_dataparser.py index b5581d5f80..3f7bf06dbd 100644 --- a/tests/dataparsers/test_nerfstudio_dataparser.py +++ b/tests/dataparsers/test_nerfstudio_dataparser.py @@ -71,10 +71,7 @@ def test_nerfstudio_dataparser_split_filelist(mocked_dataset): f.truncate(0) json.dump(data, f) - from nerfstudio.data.dataparsers.nerfstudio_dataparser import ( - Nerfstudio, - NerfstudioDataParserConfig, - ) + from nerfstudio.data.dataparsers.nerfstudio_dataparser import Nerfstudio, NerfstudioDataParserConfig parser: Nerfstudio = NerfstudioDataParserConfig( data=mocked_dataset, diff --git a/tests/pipelines/test_vanilla_pipeline.py b/tests/pipelines/test_vanilla_pipeline.py index c27aa20691..3bb73e06af 100644 --- a/tests/pipelines/test_vanilla_pipeline.py +++ b/tests/pipelines/test_vanilla_pipeline.py @@ -7,14 +7,9 @@ from torch import nn from nerfstudio.cameras.cameras import Cameras -from nerfstudio.data.datasets.base_dataset import DataparserOutputs, InputDataset -from nerfstudio.pipelines.base_pipeline import ( - Model, - ModelConfig, - VanillaPipeline, - VanillaPipelineConfig, -) from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManagerConfig +from nerfstudio.data.datasets.base_dataset import DataparserOutputs, InputDataset +from nerfstudio.pipelines.base_pipeline import Model, ModelConfig, VanillaPipeline, VanillaPipelineConfig class MockedDataManager: diff --git a/tests/process_data/test_process_images.py b/tests/process_data/test_process_images.py index a3d0323551..eed454b11f 100644 --- a/tests/process_data/test_process_images.py +++ b/tests/process_data/test_process_images.py @@ -9,9 +9,13 @@ from PIL import Image from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig -from nerfstudio.data.utils.colmap_parsing_utils import Camera -from nerfstudio.data.utils.colmap_parsing_utils import Image as ColmapImage -from nerfstudio.data.utils.colmap_parsing_utils import qvec2rotmat, write_cameras_binary, write_images_binary +from nerfstudio.data.utils.colmap_parsing_utils import ( + Camera, + Image as ColmapImage, + qvec2rotmat, + write_cameras_binary, + write_images_binary, +) from nerfstudio.process_data.images_to_nerfstudio_dataset import ImagesToNerfstudioDataset From 3e2271a6563d76d06c0fc27c8e23440914e6020f Mon Sep 17 00:00:00 2001 From: Cyrus Vachha Date: Mon, 15 Jan 2024 14:04:22 -0800 Subject: [PATCH 4/4] Setting save_world_frame as False by default and Splat export panel in viewer (#2700) * Setting save_world_frame as False by default Setting save_world_frame as False by default for the exporter as well as in the viewer, since it is exports the point cloud in a different coordinate system from the Blender add-on and there is not a similar checkbox in the original viewer, so by default the exported point cloud would not be compatible. This flag can still be set true in the cmd command and in the new viewer Also updating the blender docs to mention support for gs scenes * Adding splat export panel in viewer Adding the splat option in the export panel which returns a basic splat export cmd. Also renaming the splat exports as splat.ply instead of point_cloud.ply * Showing export panels based on method If gaussian splatting is running, then the crop box, mesh, and point cloud exports are hidden/disabled. Else, all but splat export is visible. * Fix pyright type check * Changing viewer model argument for export panels * Update export_panel.py * Update export_panel.py * Update export_panel.py * Update exporter.py * Update exporter.py --------- Co-authored-by: Brent Yi --- docs/extensions/blender_addon.md | 4 +- nerfstudio/scripts/exporter.py | 8 +- nerfstudio/viewer/export_panel.py | 206 +++++++++++++++++++----------- nerfstudio/viewer/viewer.py | 2 +- 4 files changed, 136 insertions(+), 84 deletions(-) diff --git a/docs/extensions/blender_addon.md b/docs/extensions/blender_addon.md index cc56ff346c..677f8b005c 100644 --- a/docs/extensions/blender_addon.md +++ b/docs/extensions/blender_addon.md @@ -6,7 +6,7 @@ ## Overview -This Blender add-on allows for compositing with a Nerfstudio render as a background layer by generating a camera path JSON file from the Blender camera path, as well as a way to import Nerfstudio JSON files as a Blender camera baked with the Nerfstudio camera path. This add-on also allows compositing multiple NeRF objects into a NeRF scene. This is achieved by importing a mesh or point-cloud representation of the NeRF scene from Nerfstudio into Blender and getting the camera coordinates relative to the transformations of the NeRF representation. Dynamic FOV from the Blender camera is supported and will match the Nerfstudio render. Perspective, equirectangular, VR180, and omnidirectional stereo (VR 360) cameras are supported. +This Blender add-on allows for compositing with a Nerfstudio render as a background layer by generating a camera path JSON file from the Blender camera path, as well as a way to import Nerfstudio JSON files as a Blender camera baked with the Nerfstudio camera path. This add-on also allows compositing multiple NeRF objects into a NeRF scene. This is achieved by importing a mesh or point-cloud representation of the NeRF scene from Nerfstudio into Blender and getting the camera coordinates relative to the transformations of the NeRF representation. Dynamic FOV from the Blender camera is supported and will match the Nerfstudio render. Perspective, equirectangular, VR180, and omnidirectional stereo (VR 360) cameras are supported. This add-on also supports Gaussian Splatting scenes as well, however equirectangular and VR video rendering is not currently supported.
image @@ -30,7 +30,7 @@ This Blender add-on allows for compositing with a Nerfstudio render as a backgro ## Scene Setup -1. Export the mesh or point cloud representation of the NeRF from Nerfstudio, which will be used as reference for the actual NeRF in the Blender scene. Mesh export at a good quality is preferred, however, if the export is not clear or the NeRF is large, a detailed point cloud export will also work. +1. Export the mesh or point cloud representation of the NeRF from Nerfstudio, which will be used as reference for the actual NeRF in the Blender scene. Mesh export at a good quality is preferred, however, if the export is not clear or the NeRF is large, a detailed point cloud export will also work. Keep the `save_world_frame` flag as False or in the viewer, de-select the "Save in world frame" checkbox to keep the correct coordinate system for the add-on. 2. Import the mesh or point cloud representation of the NeRF into the scene. You may need to crop the mesh further. Since it is used as a reference and won't be visible in the final render, only the parts that the blender animation will interact with may be necessary to import. diff --git a/nerfstudio/scripts/exporter.py b/nerfstudio/scripts/exporter.py index be4ad08248..45be2e84cc 100644 --- a/nerfstudio/scripts/exporter.py +++ b/nerfstudio/scripts/exporter.py @@ -121,9 +121,9 @@ class ExportPointCloud(Exporter): """Number of rays to evaluate per batch. Decrease if you run out of memory.""" std_ratio: float = 10.0 """Threshold based on STD of the average distances across the point cloud to remove outliers.""" - save_world_frame: bool = True - """If true, saves in the frame of the transform.json file, if false saves in the frame of the scaled - dataparser transform""" + save_world_frame: bool = False + """If set, saves the point cloud in the same frame as the original dataset. Otherwise, uses the + scaled and reoriented coordinate space expected by the NeRF models.""" def main(self) -> None: """Export point cloud.""" @@ -492,7 +492,7 @@ def main(self) -> None: model: GaussianSplattingModel = pipeline.model - filename = self.output_dir / "point_cloud.ply" + filename = self.output_dir / "splat.ply" map_to_tensors = {} diff --git a/nerfstudio/viewer/export_panel.py b/nerfstudio/viewer/export_panel.py index d003076f1a..5d93a20dd8 100644 --- a/nerfstudio/viewer/export_panel.py +++ b/nerfstudio/viewer/export_panel.py @@ -21,23 +21,34 @@ from typing_extensions import Literal from nerfstudio.data.scene_box import OrientedBox +from nerfstudio.models.base_model import Model +from nerfstudio.models.gaussian_splatting import GaussianSplattingModel from nerfstudio.viewer.control_panel import ControlPanel -def populate_export_tab(server: viser.ViserServer, control_panel: ControlPanel, config_path: Path) -> None: - crop_output = server.add_gui_checkbox("Use Crop", False) +def populate_export_tab( + server: viser.ViserServer, + control_panel: ControlPanel, + config_path: Path, + viewer_model: Model, +) -> None: + viewing_gsplat = isinstance(viewer_model, GaussianSplattingModel) + if not viewing_gsplat: + crop_output = server.add_gui_checkbox("Use Crop", False) - @crop_output.on_update - def _(_) -> None: - control_panel.crop_viewport = crop_output.value + @crop_output.on_update + def _(_) -> None: + control_panel.crop_viewport = crop_output.value + with server.add_gui_folder("Splat"): + populate_splat_tab(server, control_panel, config_path, viewing_gsplat) with server.add_gui_folder("Point Cloud"): - populate_point_cloud_tab(server, control_panel, config_path) + populate_point_cloud_tab(server, control_panel, config_path, viewing_gsplat) with server.add_gui_folder("Mesh"): - populate_mesh_tab(server, control_panel, config_path) + populate_mesh_tab(server, control_panel, config_path, viewing_gsplat) -def show_command_modal(client: viser.ClientHandle, what: Literal["mesh", "point cloud"], command: str) -> None: +def show_command_modal(client: viser.ClientHandle, what: Literal["mesh", "point cloud", "splat"], command: str) -> None: """Show a modal to each currently connected client. In the future, we should only show the modal to the client that pushes the @@ -79,82 +90,123 @@ def populate_point_cloud_tab( server: viser.ViserServer, control_panel: ControlPanel, config_path: Path, + viewing_gsplat: bool, ) -> None: - server.add_gui_markdown("Render depth, project to an oriented point cloud, and filter. ") - num_points = server.add_gui_number("# Points", initial_value=1_000_000, min=1, max=None, step=1) - world_frame = server.add_gui_checkbox( - "Save in world frame", - True, - hint="Save the point cloud in the transforms.json frame, rather than scaled scene frame", - ) - remove_outliers = server.add_gui_checkbox("Remove outliers", True) - normals = server.add_gui_dropdown( - "Normals", - # TODO: options here could depend on what's available to the model. - ("open3d", "model_output"), - initial_value="open3d", - hint="Normal map source.", - ) - output_dir = server.add_gui_text("Output Directory", initial_value="exports/pcd/") - generate_command = server.add_gui_button("Generate Command", icon=viser.Icon.TERMINAL_2) - - @generate_command.on_click - def _(event: viser.GuiEvent) -> None: - assert event.client is not None - command = " ".join( - [ - "ns-export pointcloud", - f"--load-config {config_path}", - f"--output-dir {output_dir.value}", - f"--num-points {num_points.value}", - f"--remove-outliers {remove_outliers.value}", - f"--normal-method {normals.value}", - f"--use_bounding_box {control_panel.crop_viewport}", - f"--save-world-frame {world_frame.value}", - get_crop_string(control_panel.crop_obb), - ] + if not viewing_gsplat: + server.add_gui_markdown("Render depth, project to an oriented point cloud, and filter ") + num_points = server.add_gui_number("# Points", initial_value=1_000_000, min=1, max=None, step=1) + world_frame = server.add_gui_checkbox( + "Save in world frame", + False, + hint=( + "If checked, saves the point cloud in the same frame as the original dataset. Otherwise, uses the " + "scaled and reoriented coordinate space expected by the NeRF models." + ), + ) + remove_outliers = server.add_gui_checkbox("Remove outliers", True) + normals = server.add_gui_dropdown( + "Normals", + # TODO: options here could depend on what's available to the model. + ("open3d", "model_output"), + initial_value="open3d", + hint="Normal map source.", ) - show_command_modal(event.client, "point cloud", command) + output_dir = server.add_gui_text("Output Directory", initial_value="exports/pcd/") + generate_command = server.add_gui_button("Generate Command", icon=viser.Icon.TERMINAL_2) + + @generate_command.on_click + def _(event: viser.GuiEvent) -> None: + assert event.client is not None + command = " ".join( + [ + "ns-export pointcloud", + f"--load-config {config_path}", + f"--output-dir {output_dir.value}", + f"--num-points {num_points.value}", + f"--remove-outliers {remove_outliers.value}", + f"--normal-method {normals.value}", + f"--use_bounding_box {control_panel.crop_viewport}", + f"--save-world-frame {world_frame.value}", + get_crop_string(control_panel.crop_obb), + ] + ) + show_command_modal(event.client, "point cloud", command) + + else: + server.add_gui_markdown("Point cloud export is not currently supported with Gaussian Splatting") def populate_mesh_tab( server: viser.ViserServer, control_panel: ControlPanel, config_path: Path, + viewing_gsplat: bool, ) -> None: - server.add_gui_markdown( - "Render depth, project to an oriented point cloud, and run Poisson surface reconstruction." - ) - - normals = server.add_gui_dropdown( - "Normals", - ("open3d", "model_output"), - initial_value="open3d", - hint="Source for normal maps.", - ) - num_faces = server.add_gui_number("# Faces", initial_value=50_000, min=1) - texture_resolution = server.add_gui_number("Texture Resolution", min=8, initial_value=2048) - output_directory = server.add_gui_text("Output Directory", initial_value="exports/mesh/") - num_points = server.add_gui_number("# Points", initial_value=1_000_000, min=1, max=None, step=1) - remove_outliers = server.add_gui_checkbox("Remove outliers", True) - - generate_command = server.add_gui_button("Generate Command", icon=viser.Icon.TERMINAL_2) - - @generate_command.on_click - def _(event: viser.GuiEvent) -> None: - assert event.client is not None - command = " ".join( - [ - "ns-export poisson", - f"--load-config {config_path}", - f"--output-dir {output_directory.value}", - f"--target-num-faces {num_faces.value}", - f"--num-pixels-per-side {texture_resolution.value}", - f"--num-points {num_points.value}", - f"--remove-outliers {remove_outliers.value}", - f"--normal-method {normals.value}", - f"--use_bounding_box {control_panel.crop_viewport}", - get_crop_string(control_panel.crop_obb), - ] + if not viewing_gsplat: + server.add_gui_markdown( + "Render depth, project to an oriented point cloud, and run Poisson surface reconstruction" ) - show_command_modal(event.client, "mesh", command) + + normals = server.add_gui_dropdown( + "Normals", + ("open3d", "model_output"), + initial_value="open3d", + hint="Source for normal maps.", + ) + num_faces = server.add_gui_number("# Faces", initial_value=50_000, min=1) + texture_resolution = server.add_gui_number("Texture Resolution", min=8, initial_value=2048) + output_directory = server.add_gui_text("Output Directory", initial_value="exports/mesh/") + num_points = server.add_gui_number("# Points", initial_value=1_000_000, min=1, max=None, step=1) + remove_outliers = server.add_gui_checkbox("Remove outliers", True) + + generate_command = server.add_gui_button("Generate Command", icon=viser.Icon.TERMINAL_2) + + @generate_command.on_click + def _(event: viser.GuiEvent) -> None: + assert event.client is not None + command = " ".join( + [ + "ns-export poisson", + f"--load-config {config_path}", + f"--output-dir {output_directory.value}", + f"--target-num-faces {num_faces.value}", + f"--num-pixels-per-side {texture_resolution.value}", + f"--num-points {num_points.value}", + f"--remove-outliers {remove_outliers.value}", + f"--normal-method {normals.value}", + f"--use_bounding_box {control_panel.crop_viewport}", + get_crop_string(control_panel.crop_obb), + ] + ) + show_command_modal(event.client, "mesh", command) + + else: + server.add_gui_markdown("Mesh export is not currently supported with Gaussian Splatting") + + +def populate_splat_tab( + server: viser.ViserServer, + control_panel: ControlPanel, + config_path: Path, + viewing_gsplat: bool, +) -> None: + if viewing_gsplat: + server.add_gui_markdown("Generate ply export of Gaussian Splat") + + output_directory = server.add_gui_text("Output Directory", initial_value="exports/splat/") + generate_command = server.add_gui_button("Generate Command", icon=viser.Icon.TERMINAL_2) + + @generate_command.on_click + def _(event: viser.GuiEvent) -> None: + assert event.client is not None + command = " ".join( + [ + "ns-export gaussian-splat", + f"--load-config {config_path}", + f"--output-dir {output_directory.value}", + ] + ) + show_command_modal(event.client, "splat", command) + + else: + server.add_gui_markdown("Splat export is only supported with Gaussian Splatting methods") diff --git a/nerfstudio/viewer/viewer.py b/nerfstudio/viewer/viewer.py index 56472d5366..8cff2b9763 100644 --- a/nerfstudio/viewer/viewer.py +++ b/nerfstudio/viewer/viewer.py @@ -204,7 +204,7 @@ def __init__( ) with tabs.add_tab("Export", viser.Icon.PACKAGE_EXPORT): - populate_export_tab(self.viser_server, self.control_panel, config_path) + populate_export_tab(self.viser_server, self.control_panel, config_path, self.pipeline.model) # Keep track of the pointers to generated GUI folders, because each generated folder holds a unique ID. viewer_gui_folders = dict()