diff --git a/README.md b/README.md index 3050b80dd..edef31622 100644 --- a/README.md +++ b/README.md @@ -184,14 +184,15 @@ ns-export pointcloud --help Using an existing dataset is great, but likely you want to use your own data! We support various methods for using your own data. Before it can be used in nerfstudio, the camera location and orientations must be determined and then converted into our format using `ns-process-data`. We rely on external tools for this, instructions and information can be found in the documentation. -| Data | Requirements | Preprocessing Speed | -| ------------------------------------------------------------------------------------------------- | ---------------- | ------------------- | -| ๐Ÿ“ท [Images](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#images-or-video) | COLMAP | ๐Ÿข | -| ๐Ÿ“น [Video](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#images-or-video) | COLMAP | ๐Ÿข | -| ๐Ÿ“ฑ [Polycam](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#polycam-capture) | LiDAR iOS Device | ๐Ÿ‡ | -| ๐Ÿ“ฑ [Record3D](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#record3d-capture) | LiDAR iOS Device | ๐Ÿ‡ | -| ๐Ÿ–ฅ [Metashape](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#metashape) | | ๐Ÿข | -| ๐Ÿ›  [Custom](https://docs.nerf.studio/en/latest/quickstart/data_conventions.html) | Poses | ๐Ÿ‡ | +| Data | Capture Device | Requirements | `ns-process-data` Speed | +| ------------------------------------------------------------------------------------------------- | -------------- | ----------------------------------------------- | ----------------------- | +| ๐Ÿ“ท [Images](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#images-and-video) | Any | [COLMAP](https://colmap.github.io/install.html) | ๐Ÿข | +| ๐Ÿ“น [Video](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#images-and-video) | Any | [COLMAP](https://colmap.github.io/install.html) | ๐Ÿข | +| ๐Ÿ“ฑ [Polycam](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#polycam-capture) | IOS with LiDAR | [Polycam App](https://poly.cam/) | ๐Ÿ‡ | +| ๐Ÿ“ฑ [KIRI Engine](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#kiri-capture) | IOS or Android | [KIRI Engine App](https://www.kiriengine.com/) | ๐Ÿ‡ | +| ๐Ÿ“ฑ [Record3D](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#record3d-capture) | IOS with LiDAR | [Record3D app](https://record3d.app/) | ๐Ÿ‡ | +| ๐Ÿ–ฅ [Metashape](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#metashape) | Any | [Metashape](https://www.agisoft.com/) | ๐Ÿ‡ | +| ๐Ÿ›  [Custom](https://docs.nerf.studio/en/latest/quickstart/data_conventions.html) | Any | Camera Poses | ๐Ÿ‡ | ## 5. Advanced Options diff --git a/colab/demo.ipynb b/colab/demo.ipynb index c47045b48..148da3b88 100644 --- a/colab/demo.ipynb +++ b/colab/demo.ipynb @@ -3,11 +3,11 @@ { "cell_type": "markdown", "metadata": { - "colab_type": "text", - "id": "view-in-github" + "id": "view-in-github", + "colab_type": "text" }, "source": [ - "\"Open" + "\"Open" ] }, { @@ -70,18 +70,18 @@ "base_uri": "https://localhost:8080/" }, "id": "RGr33zHaHak0", - "outputId": "2641dce6-cd57-4d31-d8b1-860bc1adf17f" + "outputId": "6ad38334-22d2-4675-bffb-0bae91b17577" }, "outputs": [ { - "name": "stdout", "output_type": "stream", + "name": "stdout", "text": [ "โฌ Downloading https://github.com/jaimergp/miniforge/releases/latest/download/Mambaforge-colab-Linux-x86_64.sh...\n", "๐Ÿ“ฆ Installing...\n", "๐Ÿ“Œ Adjusting configuration...\n", "๐Ÿฉน Patching environment...\n", - "โฒ Done in 0:00:27\n", + "โฒ Done in 0:00:18\n", "๐Ÿ” Restarting kernel...\n" ] } @@ -103,7 +103,7 @@ }, "outputs": [], "source": [ - "#@markdown

Install Nerfstudio and Dependencies (~10 min)

\n", + "#@markdown

Install Nerfstudio and Dependencies (~15 min)

\n", "\n", "%cd /content/\n", "!pip install --upgrade pip\n", @@ -111,8 +111,8 @@ "\n", "# Installing TinyCuda\n", "%cd /content/\n", - "!gdown \"https://drive.google.com/u/1/uc?id=1q8fuc-Mqiev5GTBTRA5UPgCaQDzuqKqj\" \n", - "!pip install tinycudann-1.6-cp37-cp37m-linux_x86_64.whl\n", + "!gdown \"https://drive.google.com/u/0/uc?id=1mB3IxcpmTVmQN-So0t2lxjxbHvqWr2zf\" \n", + "!pip install tinycudann-1.5-cp38-cp38-linux_x86_64.whl\n", "\n", "# Installing COLMAP\n", "%cd /content/\n", @@ -202,7 +202,7 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": null, "metadata": { "cellView": "form", "colab": { @@ -362,9 +362,8 @@ "metadata": { "accelerator": "GPU", "colab": { - "collapsed_sections": [], - "include_colab_link": true, - "provenance": [] + "provenance": [], + "include_colab_link": true }, "gpuClass": "standard", "kernelspec": { @@ -384,4 +383,4 @@ }, "nbformat": 4, "nbformat_minor": 0 -} +} \ No newline at end of file diff --git a/docs/quickstart/custom_dataset.md b/docs/quickstart/custom_dataset.md index e883560d8..ea508a661 100644 --- a/docs/quickstart/custom_dataset.md +++ b/docs/quickstart/custom_dataset.md @@ -11,14 +11,14 @@ ns-process-data {video,images,polycam,insta360,record3d} --data {DATA_PATH} --ou A full set of arguments can be found {doc}`here`. We Currently support the following custom data types: -| Data | Requirements | Preprocessing Speed | -| -------- | ----------------------- | ------------------- | -| ๐Ÿ“ท [Images](images_and_video) | COLMAP | ๐Ÿข | -| ๐Ÿ“น [Video](images_and_video) | COLMAP | ๐Ÿข | -| ๐Ÿ“ฑ [Polycam](polycam) | LiDAR iOS Device | ๐Ÿ‡ | -| ๐Ÿ“ฑ [Record3D](record3d) | LiDAR iOS Device | ๐Ÿ‡ | -| ๐Ÿ–ฅ [Metashape](metashape) | | ๐Ÿข | -| ๐Ÿ“ท Insta360 | 2 Sensor camera, COLMAP | ๐Ÿข | +| Data | Capture Device | Requirements | `ns-process-data` Speed | +| ----------------------------- | -------------- | ----------------------------------------------- | ----------------------- | +| ๐Ÿ“ท [Images](images_and_video) | Any | [COLMAP](https://colmap.github.io/install.html) | ๐Ÿข | +| ๐Ÿ“น [Video](images_and_video) | Any | [COLMAP](https://colmap.github.io/install.html) | ๐Ÿข | +| ๐Ÿ“ฑ [Polycam](polycam) | IOS with LiDAR | [Polycam App](https://poly.cam/) | ๐Ÿ‡ | +| ๐Ÿ“ฑ [KIRI Engine](kiri) | IOS or Android | [KIRI Engine App](https://www.kiriengine.com/) | ๐Ÿ‡ | +| ๐Ÿ“ฑ [Record3D](record3d) | IOS with LiDAR | [Record3D app](https://record3d.app/) | ๐Ÿ‡ | +| ๐Ÿ–ฅ [Metashape](metashape) | Any | [Metashape](https://www.agisoft.com/) | ๐Ÿ‡ | (images_and_video)= @@ -175,11 +175,58 @@ ns-process-data polycam --data {OUTPUT_FILE.zip} --output-dir {output directory} ns-train nerfacto --data {output directory} ``` +(kiri)= + +## KIRI Engine Capture + +Nerfstudio can trained from data processed by the [KIRI Engine app](https://www.kiriengine.com/). This works for both Android and iPhone and does not require a LiDAR supported device. + +:::{admonition} Note +:class: info +`ns-process-data` does not need to be run when using KIRI Engine. +::: + +### Setting up KIRI Engine + +```{image} imgs/kiri_setup.png +:width: 400 +:align: center +:alt: KIRI Engine setup +``` + +After downloading the app, `Developer Mode` needs to be enabled. A toggle can be found in the settings menu. + +### Process data + +```{image} imgs/kiri_capture.png +:width: 400 +:align: center +:alt: KIRI Engine setup +``` + +1. Navigate to captures window. + +2. Select `Dev.` tab. + +3. Tap the `+` button to create a new capture. + +4. Choose `Camera pose` as the capture option + +5. Capture the scene and provide a name. + +6. After processing is complete, export the scene. It will be sent to your email. + +7. Unzip the file and run the training script (`ns-process-data` is not necessary). + +``` +ns-train nerfacto --data {kiri output directory} +``` + (record3d)= ## Record3D Capture -Nerfstudio can also be trained directly from >=iPhone 12 Pro captures from the [Record3D app](https://record3d.app/). This uses the iPhone's LiDAR sensors to calculate camera poses, so COLMAP is not needed. +Nerfstudio can be trained directly from >=iPhone 12 Pro captures from the [Record3D app](https://record3d.app/). This uses the iPhone's LiDAR sensors to calculate camera poses, so COLMAP is not needed. Click on the image down below ๐Ÿ‘‡ for a 1-minute tutorial on how to run nerfstudio with Record3D from start to finish. diff --git a/docs/quickstart/imgs/kiri_capture.png b/docs/quickstart/imgs/kiri_capture.png new file mode 100644 index 000000000..64f1531b0 Binary files /dev/null and b/docs/quickstart/imgs/kiri_capture.png differ diff --git a/docs/quickstart/imgs/kiri_setup.png b/docs/quickstart/imgs/kiri_setup.png new file mode 100644 index 000000000..4203985ea Binary files /dev/null and b/docs/quickstart/imgs/kiri_setup.png differ diff --git a/nerfstudio/cameras/camera_utils.py b/nerfstudio/cameras/camera_utils.py index 30d755596..d3b0ca619 100644 --- a/nerfstudio/cameras/camera_utils.py +++ b/nerfstudio/cameras/camera_utils.py @@ -409,7 +409,7 @@ def rotation_matrix(a: TensorType[3], b: TensorType[3]) -> TensorType[3, 3]: def auto_orient_and_center_poses( poses: TensorType["num_poses":..., 4, 4], method: Literal["pca", "up", "none"] = "up", center_poses: bool = True -) -> TensorType["num_poses":..., 3, 4]: +) -> Tuple[TensorType["num_poses":..., 3, 4], TensorType[4, 4]]: """Orients and centers the poses. We provide two methods for orientation: pca and up. pca: Orient the poses so that the principal component of the points is aligned with the axes. @@ -424,7 +424,7 @@ def auto_orient_and_center_poses( center_poses: If True, the poses are centered around the origin. Returns: - The oriented poses. + Tuple of the oriented poses and the transform matrix. """ translation = poses[..., :3, 3] @@ -457,7 +457,9 @@ def auto_orient_and_center_poses( transform = torch.cat([rotation, rotation @ -translation[..., None]], dim=-1) oriented_poses = transform @ poses elif method == "none": - oriented_poses = poses - poses[:, :3, 3] -= translation + transform = torch.eye(4) + transform[:3, 3] = -translation + transform = transform[:3, :] + oriented_poses = transform @ poses - return oriented_poses + return oriented_poses, transform diff --git a/nerfstudio/configs/base_config.py b/nerfstudio/configs/base_config.py index e4828fab6..a74f07ed7 100644 --- a/nerfstudio/configs/base_config.py +++ b/nerfstudio/configs/base_config.py @@ -20,23 +20,14 @@ from __future__ import annotations import warnings from dataclasses import dataclass -from datetime import datetime from pathlib import Path -from typing import Any, Dict, List, Optional, Tuple, Type - -import yaml -from rich.console import Console -from typing_extensions import Literal - -from nerfstudio.configs.config_utils import to_immutable_dict +from typing import Any, List, Optional, Tuple, Type # model instances from nerfstudio.utils import writer warnings.filterwarnings("ignore", module="torchvision") -CONSOLE = Console(width=120) - # Pretty printing class class PrintableConfig: # pylint: disable=too-few-public-methods """Printable Config defining str function""" @@ -186,95 +177,3 @@ class ViewerConfig(PrintableConfig): skip_openrelay: bool = False """Avoid using openrelay to communicate with the viewer. Try disabling if you have trouble connecting to the viewer""" - - -from nerfstudio.engine.optimizers import OptimizerConfig -from nerfstudio.engine.schedulers import SchedulerConfig -from nerfstudio.pipelines.base_pipeline import VanillaPipelineConfig - - -@dataclass -class Config(PrintableConfig): - """Full config contents""" - - output_dir: Path = Path("outputs") - """relative or absolute output directory to save all checkpoints and logging""" - method_name: Optional[str] = None - """Method name. Required to set in python or via cli""" - experiment_name: Optional[str] = None - """Experiment name. If None, will automatically be set to dataset name""" - timestamp: str = "{timestamp}" - """Experiment timestamp.""" - machine: MachineConfig = MachineConfig() - """Machine configuration""" - logging: LoggingConfig = LoggingConfig() - """Logging configuration""" - viewer: ViewerConfig = ViewerConfig() - """Viewer configuration""" - trainer: TrainerConfig = TrainerConfig() - """Trainer configuration""" - pipeline: VanillaPipelineConfig = VanillaPipelineConfig() - """Pipeline configuration""" - optimizers: Dict[str, Any] = to_immutable_dict( - { - "fields": { - "optimizer": OptimizerConfig(), - "scheduler": SchedulerConfig(), - } - } - ) - """Dictionary of optimizer groups and their schedulers""" - vis: Literal["viewer", "wandb", "tensorboard"] = "wandb" - """Which visualizer to use.""" - data: Optional[Path] = None - """Alias for --pipeline.datamanager.dataparser.data""" - wandb_name: Optional[str] = None - """Alias for --logging.wandb.name""" - - def is_viewer_enabled(self) -> bool: - """Checks if a viewer is enabled.""" - return "viewer" == self.vis - - def is_wandb_enabled(self) -> bool: - """Checks if wandb is enabled.""" - return "wandb" == self.vis - - def is_tensorboard_enabled(self) -> bool: - """Checks if tensorboard is enabled.""" - return "tensorboard" == self.vis - - def set_timestamp(self) -> None: - """Dynamically set the experiment timestamp""" - if self.timestamp == "{timestamp}": - self.timestamp = datetime.now().strftime("%Y-%m-%d_%H%M%S") - - def set_experiment_name(self) -> None: - """Dynamically set the experiment name""" - if self.experiment_name is None: - self.experiment_name = str(self.pipeline.datamanager.dataparser.data).replace("../", "").replace("/", "-") - - def get_base_dir(self) -> Path: - """Retrieve the base directory to set relative paths""" - # check the experiment and method names - assert self.method_name is not None, "Please set method name in config or via the cli" - self.set_experiment_name() - return Path(f"{self.output_dir}/{self.experiment_name}/{self.method_name}/{self.timestamp}") - - def get_checkpoint_dir(self) -> Path: - """Retrieve the checkpoint directory""" - return Path(self.get_base_dir() / self.trainer.relative_model_dir) - - def print_to_terminal(self) -> None: - """Helper to pretty print config to terminal""" - CONSOLE.rule("Config") - CONSOLE.print(self) - CONSOLE.rule("") - - def save_config(self) -> None: - """Save config to base directory""" - base_dir = self.get_base_dir() - assert base_dir is not None - base_dir.mkdir(parents=True, exist_ok=True) - config_yaml_path = base_dir / "config.yml" - CONSOLE.log(f"Saving config to: {config_yaml_path}") - config_yaml_path.write_text(yaml.dump(self), "utf8") diff --git a/nerfstudio/configs/experiment_config.py b/nerfstudio/configs/experiment_config.py new file mode 100644 index 000000000..9db15af5e --- /dev/null +++ b/nerfstudio/configs/experiment_config.py @@ -0,0 +1,125 @@ +# Copyright 2022 The Nerfstudio Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Config used for running an experiment""" + +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from typing import Any, Dict, Optional + +import yaml +from rich.console import Console +from typing_extensions import Literal + +from nerfstudio.configs.base_config import ( + LoggingConfig, + MachineConfig, + PrintableConfig, + TrainerConfig, + ViewerConfig, +) +from nerfstudio.configs.config_utils import to_immutable_dict +from nerfstudio.engine.optimizers import OptimizerConfig +from nerfstudio.engine.schedulers import SchedulerConfig +from nerfstudio.pipelines.base_pipeline import VanillaPipelineConfig + +CONSOLE = Console(width=120) + + +@dataclass +class ExperimentConfig(PrintableConfig): + """Full config contents""" + + output_dir: Path = Path("outputs") + """relative or absolute output directory to save all checkpoints and logging""" + method_name: Optional[str] = None + """Method name. Required to set in python or via cli""" + experiment_name: Optional[str] = None + """Experiment name. If None, will automatically be set to dataset name""" + timestamp: str = "{timestamp}" + """Experiment timestamp.""" + machine: MachineConfig = MachineConfig() + """Machine configuration""" + logging: LoggingConfig = LoggingConfig() + """Logging configuration""" + viewer: ViewerConfig = ViewerConfig() + """Viewer configuration""" + trainer: TrainerConfig = TrainerConfig() + """Trainer configuration""" + pipeline: VanillaPipelineConfig = VanillaPipelineConfig() + """Pipeline configuration""" + optimizers: Dict[str, Any] = to_immutable_dict( + { + "fields": { + "optimizer": OptimizerConfig(), + "scheduler": SchedulerConfig(), + } + } + ) + """Dictionary of optimizer groups and their schedulers""" + vis: Literal["viewer", "wandb", "tensorboard"] = "wandb" + """Which visualizer to use.""" + data: Optional[Path] = None + """Alias for --pipeline.datamanager.dataparser.data""" + wandb_name: Optional[str] = None + """Alias for --logging.wandb.name""" + + def is_viewer_enabled(self) -> bool: + """Checks if a viewer is enabled.""" + return "viewer" == self.vis + + def is_wandb_enabled(self) -> bool: + """Checks if wandb is enabled.""" + return "wandb" == self.vis + + def is_tensorboard_enabled(self) -> bool: + """Checks if tensorboard is enabled.""" + return "tensorboard" == self.vis + + def set_timestamp(self) -> None: + """Dynamically set the experiment timestamp""" + if self.timestamp == "{timestamp}": + self.timestamp = datetime.now().strftime("%Y-%m-%d_%H%M%S") + + def set_experiment_name(self) -> None: + """Dynamically set the experiment name""" + if self.experiment_name is None: + self.experiment_name = str(self.pipeline.datamanager.dataparser.data).replace("../", "").replace("/", "-") + + def get_base_dir(self) -> Path: + """Retrieve the base directory to set relative paths""" + # check the experiment and method names + assert self.method_name is not None, "Please set method name in config or via the cli" + self.set_experiment_name() + return Path(f"{self.output_dir}/{self.experiment_name}/{self.method_name}/{self.timestamp}") + + def get_checkpoint_dir(self) -> Path: + """Retrieve the checkpoint directory""" + return Path(self.get_base_dir() / self.trainer.relative_model_dir) + + def print_to_terminal(self) -> None: + """Helper to pretty print config to terminal""" + CONSOLE.rule("Config") + CONSOLE.print(self) + CONSOLE.rule("") + + def save_config(self) -> None: + """Save config to base directory""" + base_dir = self.get_base_dir() + assert base_dir is not None + base_dir.mkdir(parents=True, exist_ok=True) + config_yaml_path = base_dir / "config.yml" + CONSOLE.log(f"Saving config to: {config_yaml_path}") + config_yaml_path.write_text(yaml.dump(self), "utf8") diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index f666b5ed1..3d4b694d9 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -21,14 +21,11 @@ from __future__ import annotations from typing import Dict import tyro +from nerfacc import ContractionType from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig -from nerfstudio.configs.base_config import ( - Config, - SchedulerConfig, - TrainerConfig, - ViewerConfig, -) +from nerfstudio.configs.base_config import TrainerConfig, ViewerConfig +from nerfstudio.configs.experiment_config import ExperimentConfig from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManagerConfig from nerfstudio.data.datamanagers.semantic_datamanager import SemanticDataManagerConfig from nerfstudio.data.datamanagers.variable_res_datamanager import ( @@ -37,11 +34,15 @@ from nerfstudio.data.datamanagers.variable_res_datamanager import ( from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig from nerfstudio.data.dataparsers.friends_dataparser import FriendsDataParserConfig +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.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig +from nerfstudio.engine.schedulers import SchedulerConfig from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind from nerfstudio.models.instant_ngp import InstantNGPModelConfig from nerfstudio.models.mipnerf import MipNerfModel @@ -52,10 +53,11 @@ from nerfstudio.models.vanilla_nerf import NeRFModel, VanillaModelConfig from nerfstudio.pipelines.base_pipeline import VanillaPipelineConfig from nerfstudio.pipelines.dynamic_batch import DynamicBatchPipelineConfig -method_configs: Dict[str, Config] = {} +method_configs: Dict[str, ExperimentConfig] = {} descriptions = { "nerfacto": "Recommended real-time model tuned for real captures. This model will be continually updated.", - "instant-ngp": "Implementation of Instant-NGP. Recommended real-time model for bounded synthetic data.", + "instant-ngp": "Implementation of Instant-NGP. Recommended real-time model for unbounded scenes.", + "instant-ngp-bounded": "Implementation of Instant-NGP. Recommended for bounded real and synthetic scenes", "mipnerf": "High quality model for bounded scenes. (slow)", "semantic-nerfw": "Predicts semantic segmentations and filters out transient objects.", "vanilla-nerf": "Original NeRF model. (slow)", @@ -64,7 +66,7 @@ descriptions = { "phototourism": "Uses the Phototourism data.", } -method_configs["nerfacto"] = Config( +method_configs["nerfacto"] = ExperimentConfig( method_name="nerfacto", trainer=TrainerConfig( steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True @@ -94,7 +96,7 @@ method_configs["nerfacto"] = Config( vis="viewer", ) -method_configs["instant-ngp"] = Config( +method_configs["instant-ngp"] = ExperimentConfig( method_name="instant-ngp", trainer=TrainerConfig( steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True @@ -113,16 +115,44 @@ method_configs["instant-ngp"] = Config( vis="viewer", ) -method_configs["mipnerf"] = Config( + +method_configs["instant-ngp-bounded"] = ExperimentConfig( + method_name="instant-ngp-bounded", + trainer=TrainerConfig( + steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True + ), + pipeline=DynamicBatchPipelineConfig( + datamanager=VanillaDataManagerConfig(dataparser=InstantNGPDataParserConfig(), train_num_rays_per_batch=8192), + model=InstantNGPModelConfig( + eval_num_rays_per_chunk=8192, + contraction_type=ContractionType.AABB, + render_step_size=0.001, + max_num_samples_per_ray=48, + near_plane=0.01, + background_color="black", + ), + ), + optimizers={ + "fields": { + "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": None, + } + }, + viewer=ViewerConfig(num_rays_per_chunk=64000), + vis="viewer", +) + + +method_configs["mipnerf"] = ExperimentConfig( method_name="mipnerf", pipeline=VanillaPipelineConfig( - datamanager=VanillaDataManagerConfig(dataparser=BlenderDataParserConfig(), train_num_rays_per_batch=8192), + datamanager=VanillaDataManagerConfig(dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=1024), model=VanillaModelConfig( _target=MipNerfModel, loss_coefficients={"rgb_loss_coarse": 0.1, "rgb_loss_fine": 1.0}, num_coarse_samples=128, num_importance_samples=128, - eval_num_rays_per_chunk=8192, + eval_num_rays_per_chunk=1024, ), ), optimizers={ @@ -133,7 +163,7 @@ method_configs["mipnerf"] = Config( }, ) -method_configs["semantic-nerfw"] = Config( +method_configs["semantic-nerfw"] = ExperimentConfig( method_name="semantic-nerfw", trainer=TrainerConfig( steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True @@ -158,7 +188,7 @@ method_configs["semantic-nerfw"] = Config( vis="viewer", ) -method_configs["vanilla-nerf"] = Config( +method_configs["vanilla-nerf"] = ExperimentConfig( method_name="vanilla-nerf", pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig( @@ -178,7 +208,7 @@ method_configs["vanilla-nerf"] = Config( }, ) -method_configs["tensorf"] = Config( +method_configs["tensorf"] = ExperimentConfig( method_name="tensorf", trainer=TrainerConfig(mixed_precision=False), pipeline=VanillaPipelineConfig( @@ -199,7 +229,7 @@ method_configs["tensorf"] = Config( }, ) -method_configs["dnerf"] = Config( +method_configs["dnerf"] = ExperimentConfig( method_name="dnerf", pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig(dataparser=DNeRFDataParserConfig()), @@ -221,7 +251,7 @@ method_configs["dnerf"] = Config( }, ) -method_configs["phototourism"] = Config( +method_configs["phototourism"] = ExperimentConfig( method_name="phototourism", trainer=TrainerConfig( steps_per_eval_batch=500, steps_per_save=2000, max_num_iterations=30000, mixed_precision=True diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 93d17ed70..b17e4fdeb 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -35,6 +35,7 @@ from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig from nerfstudio.cameras.cameras import CameraType from nerfstudio.cameras.rays import RayBundle from nerfstudio.configs.base_config import InstantiateConfig +from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig from nerfstudio.data.dataparsers.friends_dataparser import FriendsDataParserConfig @@ -46,7 +47,6 @@ from nerfstudio.data.dataparsers.nuscenes_dataparser import NuScenesDataParserCo from nerfstudio.data.dataparsers.phototourism_dataparser import ( PhototourismDataParserConfig, ) -from nerfstudio.data.dataparsers.record3d_dataparser import Record3DDataParserConfig from nerfstudio.data.datasets.base_dataset import InputDataset from nerfstudio.data.pixel_samplers import EquirectangularPixelSampler, PixelSampler from nerfstudio.data.utils.dataloaders import CacheDataloader @@ -66,7 +66,6 @@ AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[ # Omit prefixes of "friends-data": FriendsDataParserConfig(), "instant-ngp-data": InstantNGPDataParserConfig(), "nuscenes-data": NuScenesDataParserConfig(), - "record3d-data": Record3DDataParserConfig(), "dnerf-data": DNeRFDataParserConfig(), "phototourism-data": PhototourismDataParserConfig(), }, @@ -293,6 +292,7 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method config: VanillaDataManagerConfig train_dataset: InputDataset eval_dataset: InputDataset + train_dataparser_outputs: DataparserOutputs def __init__( self, @@ -311,6 +311,7 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method self.test_mode = test_mode self.test_split = "test" if test_mode in ["test", "inference"] else "val" self.dataparser = self.config.dataparser.setup() + self.train_dataparser_outputs = self.dataparser.get_dataparser_outputs(split="train") self.train_dataset = self.create_train_dataset() self.eval_dataset = self.create_eval_dataset() @@ -319,7 +320,7 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method def create_train_dataset(self) -> InputDataset: """Sets up the data loaders for training""" return InputDataset( - dataparser_outputs=self.dataparser.get_dataparser_outputs(split="train"), + dataparser_outputs=self.train_dataparser_outputs, scale_factor=self.config.camera_res_scale_factor, ) diff --git a/nerfstudio/data/datamanagers/semantic_datamanager.py b/nerfstudio/data/datamanagers/semantic_datamanager.py index b3cc223ae..188c7f6c3 100644 --- a/nerfstudio/data/datamanagers/semantic_datamanager.py +++ b/nerfstudio/data/datamanagers/semantic_datamanager.py @@ -42,7 +42,7 @@ class SemanticDataManager(VanillaDataManager): # pylint: disable=abstract-metho def create_train_dataset(self) -> SemanticDataset: return SemanticDataset( - dataparser_outputs=self.dataparser.get_dataparser_outputs(split="train"), + dataparser_outputs=self.train_dataparser_outputs, scale_factor=self.config.camera_res_scale_factor, ) diff --git a/nerfstudio/data/dataparsers/base_dataparser.py b/nerfstudio/data/dataparsers/base_dataparser.py index 3b99f274e..89c7c4e21 100644 --- a/nerfstudio/data/dataparsers/base_dataparser.py +++ b/nerfstudio/data/dataparsers/base_dataparser.py @@ -16,6 +16,7 @@ from __future__ import annotations +import json from abc import abstractmethod from dataclasses import dataclass, field from pathlib import Path @@ -63,11 +64,31 @@ class DataparserOutputs: """Dictionary of any metadata that be required for the given experiment. Will be processed by the InputDataset to create any additional tensors that may be required. """ + dataparser_transform: TensorType[3, 4] = torch.eye(4)[:3, :] + """Transform applied by the dataparser.""" + dataparser_scale: float = 1.0 + """Scale applied by the dataparser.""" def as_dict(self) -> dict: """Returns the dataclass as a dictionary.""" return vars(self) + def save_dataparser_transform(self, path: Path): + """Save dataparser transform to json file. Some dataparsers will apply a transform to the poses, + this method allows the transform to be saved so that it can be used in other applications. + + Args: + path: path to save transform to + """ + data = { + "transform": self.dataparser_transform.tolist(), + "scale": float(self.dataparser_scale), + } + if not path.parent.exists(): + path.parent.mkdir(parents=True) + with open(path, "w", encoding="UTF-8") as file: + json.dump(data, file, indent=4) + @dataclass class DataParserConfig(cfg.InstantiateConfig): diff --git a/nerfstudio/data/dataparsers/blender_dataparser.py b/nerfstudio/data/dataparsers/blender_dataparser.py index 0f31acff1..f1044e826 100644 --- a/nerfstudio/data/dataparsers/blender_dataparser.py +++ b/nerfstudio/data/dataparsers/blender_dataparser.py @@ -104,6 +104,7 @@ class Blender(DataParser): cameras=cameras, alpha_color=alpha_color_tensor, scene_box=scene_box, + dataparser_scale=self.scale_factor, ) return dataparser_outputs diff --git a/nerfstudio/data/dataparsers/dnerf_dataparser.py b/nerfstudio/data/dataparsers/dnerf_dataparser.py index 56a9f232b..c26621e92 100644 --- a/nerfstudio/data/dataparsers/dnerf_dataparser.py +++ b/nerfstudio/data/dataparsers/dnerf_dataparser.py @@ -102,7 +102,11 @@ class DNeRF(DataParser): ) dataparser_outputs = DataparserOutputs( - image_filenames=image_filenames, cameras=cameras, alpha_color=alpha_color_tensor, scene_box=scene_box + image_filenames=image_filenames, + cameras=cameras, + alpha_color=alpha_color_tensor, + scene_box=scene_box, + dataparser_scale=self.scale_factor, ) return dataparser_outputs diff --git a/nerfstudio/data/dataparsers/friends_dataparser.py b/nerfstudio/data/dataparsers/friends_dataparser.py index 06fceef3c..604dbd7c0 100644 --- a/nerfstudio/data/dataparsers/friends_dataparser.py +++ b/nerfstudio/data/dataparsers/friends_dataparser.py @@ -146,5 +146,6 @@ class Friends(DataParser): cameras=cameras, scene_box=scene_box, metadata={"semantics": semantics} if self.config.include_semantics else {}, + dataparser_scale=scale, ) return dataparser_outputs diff --git a/nerfstudio/data/dataparsers/instant_ngp_dataparser.py b/nerfstudio/data/dataparsers/instant_ngp_dataparser.py index 1e5a69b99..6a3478749 100644 --- a/nerfstudio/data/dataparsers/instant_ngp_dataparser.py +++ b/nerfstudio/data/dataparsers/instant_ngp_dataparser.py @@ -20,6 +20,7 @@ from dataclasses import dataclass, field from pathlib import Path from typing import Dict, Tuple, Type +import imageio import numpy as np import torch from rich.console import Console @@ -44,10 +45,10 @@ class InstantNGPDataParserConfig(DataParserConfig): _target: Type = field(default_factory=lambda: InstantNGP) """target class to instantiate""" data: Path = Path("data/ours/posterv2") - """Directory specifying location of data.""" + """Directory or explicit json file path specifying location of data.""" scale_factor: float = 1.0 """How much to scale the camera origins by.""" - scene_scale: float = 0.33 + scene_scale: float = 0.3333 """How much to scale the scene.""" @@ -58,16 +59,34 @@ class InstantNGP(DataParser): config: InstantNGPDataParserConfig def _generate_dataparser_outputs(self, split="train"): + if self.config.data.suffix == ".json": + meta = load_from_json(self.config.data) + data_dir = self.config.data.parent + else: + meta = load_from_json(self.config.data / "transforms.json") + data_dir = self.config.data - meta = load_from_json(self.config.data / "transforms.json") image_filenames = [] poses = [] num_skipped_image_filenames = 0 for frame in meta["frames"]: - fname = self.config.data / Path(frame["file_path"]) - if not fname: + fname = data_dir / Path(frame["file_path"]) + # search for png file + if not fname.exists(): + fname = data_dir / Path(frame["file_path"] + ".png") + if not fname.exists(): + CONSOLE.log(f"couldnt find {fname} image") num_skipped_image_filenames += 1 else: + if "w" not in meta: + img_0 = imageio.imread(fname) + h, w = img_0.shape[:2] + meta["w"] = w + if "h" in meta: + meta_h = meta["h"] + assert meta_h == h, f"height of image dont not correspond metadata {h} != {meta_h}" + else: + meta["h"] = h image_filenames.append(fname) poses.append(np.array(frame["transform_matrix"])) if num_skipped_image_filenames >= 0: @@ -84,12 +103,18 @@ class InstantNGP(DataParser): camera_to_world = torch.from_numpy(poses[:, :3]) # camera to world transform distortion_params = camera_utils.get_distortion_params( - k1=float(meta["k1"]), k2=float(meta["k2"]), p1=float(meta["p1"]), p2=float(meta["p2"]) + k1=float(meta.get("k1", 0)), + k2=float(meta.get("k2", 0)), + k3=float(meta.get("k3", 0)), + k4=float(meta.get("k4", 0)), + p1=float(meta.get("p1", 0)), + p2=float(meta.get("p2", 0)), ) # in x,y,z order # assumes that the scene is centered at the origin - aabb_scale = meta["aabb_scale"] + aabb_scale = 0.5 * meta.get("aabb_scale", 1) + scene_box = SceneBox( aabb=torch.tensor( [[-aabb_scale, -aabb_scale, -aabb_scale], [aabb_scale, aabb_scale, aabb_scale]], dtype=torch.float32 @@ -98,14 +123,16 @@ class InstantNGP(DataParser): fl_x, fl_y = InstantNGP.get_focal_lengths(meta) + w, h = meta["w"], meta["h"] + cameras = Cameras( fx=float(fl_x), fy=float(fl_y), - cx=float(meta["cx"]), - cy=float(meta["cy"]), + cx=float(meta.get("cx", 0.5 * w)), + cy=float(meta.get("cy", 0.5 * h)), distortion_params=distortion_params, - height=int(meta["h"]), - width=int(meta["w"]), + height=int(h), + width=int(w), camera_to_worlds=camera_to_world, camera_type=CameraType.PERSPECTIVE, ) @@ -115,6 +142,7 @@ class InstantNGP(DataParser): image_filenames=image_filenames, cameras=cameras, scene_box=scene_box, + dataparser_scale=self.config.scene_scale, ) return dataparser_outputs @@ -139,12 +167,15 @@ class InstantNGP(DataParser): elif "camera_angle_x" in meta: fl_x = fov_to_focal_length(meta["camera_angle_x"], meta["w"]) - if "fl_y" in meta: - fl_y = meta["fl_y"] - elif "y_fov" in meta: - fl_y = fov_to_focal_length(np.deg2rad(meta["y_fov"]), meta["h"]) - elif "camera_angle_y" in meta: - fl_y = fov_to_focal_length(meta["camera_angle_y"], meta["h"]) + if "camera_angle_y" not in meta or "y_fov" not in meta: + fl_y = fl_x + else: + if "fl_y" in meta: + fl_y = meta["fl_y"] + elif "y_fov" in meta: + fl_y = fov_to_focal_length(np.deg2rad(meta["y_fov"]), meta["h"]) + elif "camera_angle_y" in meta: + fl_y = fov_to_focal_length(meta["camera_angle_y"], meta["h"]) if fl_x == 0 or fl_y == 0: raise AttributeError("Focal length cannot be calculated from transforms.json (missing fields).") diff --git a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py index e2b86bd6d..e311a08cc 100644 --- a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py +++ b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py @@ -47,7 +47,7 @@ class NerfstudioDataParserConfig(DataParserConfig): _target: Type = field(default_factory=lambda: Nerfstudio) """target class to instantiate""" data: Path = Path("data/nerfstudio/poster") - """Directory specifying location of data.""" + """Directory or explicit json file path specifying location of data.""" scale_factor: float = 1.0 """How much to scale the camera origins by.""" downscale_factor: Optional[int] = None @@ -74,7 +74,13 @@ class Nerfstudio(DataParser): def _generate_dataparser_outputs(self, split="train"): # pylint: disable=too-many-statements - meta = load_from_json(self.config.data / "transforms.json") + if self.config.data.suffix == ".json": + meta = load_from_json(self.config.data) + data_dir = self.config.data.parent + else: + meta = load_from_json(self.config.data / "transforms.json") + data_dir = self.config.data + image_filenames = [] mask_filenames = [] poses = [] @@ -101,7 +107,7 @@ class Nerfstudio(DataParser): for frame in meta["frames"]: filepath = PurePath(frame["file_path"]) - fname = self._get_fname(filepath) + fname = self._get_fname(filepath, data_dir) if not fname.exists(): num_skipped_image_filenames += 1 continue @@ -140,7 +146,11 @@ class Nerfstudio(DataParser): poses.append(np.array(frame["transform_matrix"])) if "mask_path" in frame: mask_filepath = PurePath(frame["mask_path"]) - mask_fname = self._get_fname(mask_filepath, downsample_folder_prefix="masks_") + mask_fname = self._get_fname( + mask_filepath, + data_dir, + downsample_folder_prefix="masks_", + ) mask_filenames.append(mask_fname) if num_skipped_image_filenames >= 0: CONSOLE.log(f"Skipping {num_skipped_image_filenames} files in dataset split {split}.") @@ -181,7 +191,7 @@ class Nerfstudio(DataParser): orientation_method = self.config.orientation_method poses = torch.from_numpy(np.array(poses).astype(np.float32)) - poses = camera_utils.auto_orient_and_center_poses( + poses, transform_matrix = camera_utils.auto_orient_and_center_poses( poses, method=orientation_method, center_poses=self.config.center_poses, @@ -190,12 +200,14 @@ class Nerfstudio(DataParser): # Scale poses scale_factor = 1.0 if self.config.auto_scale_poses: - scale_factor /= torch.max(torch.abs(poses[:, :3, 3])) + scale_factor /= float(torch.max(torch.abs(poses[:, :3, 3]))) + scale_factor *= self.config.scale_factor - poses[:, :3, 3] *= scale_factor * self.config.scale_factor + poses[:, :3, 3] *= scale_factor # Choose image_filenames and poses based on split, but after auto orient and scaling the poses. image_filenames = [image_filenames[i] for i in indices] + mask_filenames = [mask_filenames[i] for i in indices] if len(mask_filenames) > 0 else [] poses = poses[indices] # in x,y,z order @@ -251,24 +263,30 @@ class Nerfstudio(DataParser): cameras=cameras, scene_box=scene_box, mask_filenames=mask_filenames if len(mask_filenames) > 0 else None, + dataparser_scale=scale_factor, + dataparser_transform=transform_matrix, ) return dataparser_outputs - def _get_fname(self, filepath: PurePath, downsample_folder_prefix="images_") -> Path: + def _get_fname(self, filepath: PurePath, data_dir: PurePath, downsample_folder_prefix="images_") -> Path: """Get the filename of the image file. downsample_folder_prefix can be used to point to auxillary image data, e.g. masks + + filepath: the base file name of the transformations. + data_dir: the directory of the data that contains the transform file + downsample_folder_prefix: prefix of the newly generated downsampled images """ if self.downscale_factor is None: if self.config.downscale_factor is None: - test_img = Image.open(self.config.data / filepath) + test_img = Image.open(data_dir / filepath) h, w = test_img.size max_res = max(h, w) df = 0 while True: if (max_res / 2 ** (df)) < MAX_AUTO_RESOLUTION: break - if not (self.config.data / f"{downsample_folder_prefix}{2**(df+1)}" / filepath.name).exists(): + if not (data_dir / f"{downsample_folder_prefix}{2**(df+1)}" / filepath.name).exists(): break df += 1 @@ -278,5 +296,5 @@ class Nerfstudio(DataParser): self.downscale_factor = self.config.downscale_factor if self.downscale_factor > 1: - return self.config.data / f"{downsample_folder_prefix}{self.downscale_factor}" / filepath.name - return self.config.data / filepath + return data_dir / f"{downsample_folder_prefix}{self.downscale_factor}" / filepath.name + return data_dir / filepath diff --git a/nerfstudio/data/dataparsers/phototourism_dataparser.py b/nerfstudio/data/dataparsers/phototourism_dataparser.py index f0f6987e0..bb401d9dc 100644 --- a/nerfstudio/data/dataparsers/phototourism_dataparser.py +++ b/nerfstudio/data/dataparsers/phototourism_dataparser.py @@ -137,16 +137,17 @@ class Phototourism(DataParser): else: raise ValueError(f"Unknown dataparser split {split}") - poses = camera_utils.auto_orient_and_center_poses( + poses, transform_matrix = camera_utils.auto_orient_and_center_poses( poses, method=self.config.orientation_method, center_poses=self.config.center_poses ) # Scale poses scale_factor = 1.0 if self.config.auto_scale_poses: - scale_factor /= torch.max(torch.abs(poses[:, :3, 3])) + scale_factor /= float(torch.max(torch.abs(poses[:, :3, 3]))) + scale_factor *= self.config.scale_factor - poses[:, :3, 3] *= scale_factor * self.config.scale_factor + poses[:, :3, 3] *= scale_factor # in x,y,z order # assumes that the scene is centered at the origin @@ -175,6 +176,8 @@ class Phototourism(DataParser): image_filenames=image_filenames, cameras=cameras, scene_box=scene_box, + dataparser_scale=scale_factor, + dataparser_transform=transform_matrix, ) return dataparser_outputs diff --git a/nerfstudio/data/dataparsers/record3d_dataparser.py b/nerfstudio/data/dataparsers/record3d_dataparser.py deleted file mode 100644 index 099470d01..000000000 --- a/nerfstudio/data/dataparsers/record3d_dataparser.py +++ /dev/null @@ -1,160 +0,0 @@ -# Copyright 2022 The Nerfstudio Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Data parser for record3d dataset""" -from __future__ import annotations - -from dataclasses import dataclass, field -from pathlib import Path -from typing import Type - -import numpy as np -import torch -from rich.console import Console -from scipy.spatial.transform import Rotation -from typing_extensions import Literal - -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.scene_box import SceneBox -from nerfstudio.utils import poses as pose_utils -from nerfstudio.utils.io import load_from_json - -CONSOLE = Console(width=120) - - -@dataclass -class Record3DDataParserConfig(DataParserConfig): - """Record3D dataset config""" - - _target: Type = field(default_factory=lambda: Record3D) - """target class to instantiate""" - data: Path = Path("data/record3d/bear") - """Location of data""" - val_skip: int = 8 - """1/val_skip images to use for validation.""" - aabb_scale: float = 4.0 - """Scene scale.""" - orientation_method: Literal["pca", "up"] = "up" - """The method to use for orientation""" - max_dataset_size: int = 300 - """Max number of images to train on. If the dataset has more, images will be sampled approximately evenly. If -1, - use all images.""" - - -@dataclass -class Record3D(DataParser): - """Record3D Dataset""" - - config: Record3DDataParserConfig - - def _generate_dataparser_outputs(self, split: str = "train") -> DataparserOutputs: - - CONSOLE.print( - "[bold red]DEPRECATION WARNING: The Record3D dataparser will be deprecated in future versions. " - "Use `ns-data-process record3d` to convert the data into the nerfstudio format instead." - ) - - image_dir = self.config.data / "rgb" - - if not image_dir.exists(): - raise ValueError(f"Image directory {image_dir} doesn't exist") - - image_filenames = [] - for f in image_dir.iterdir(): - if f.stem.isdigit(): # removes possible duplicate images (for example, 123(3).jpg) - image_filenames.append(f) - - image_filenames = sorted(image_filenames, key=lambda fn: int(fn.stem)) - image_filenames = np.array(image_filenames) - num_images = len(image_filenames) - - metadata_path = self.config.data / "metadata.json" - metadata_dict = load_from_json(metadata_path) - - poses_data = np.array(metadata_dict["poses"]) - # (N, 3, 4) - poses = np.concatenate( - [Rotation.from_quat(poses_data[:, :4]).as_matrix(), poses_data[:, 4:, None]], - axis=-1, - ).astype(np.float32) - - if self.config.max_dataset_size != -1 and num_images > self.config.max_dataset_size: - # Evenly select max_dataset_size images from dataset, including first - # and last indices. - idx = np.round(np.linspace(0, num_images - 1, self.config.max_dataset_size)).astype(int) - poses = poses[idx] - image_filenames = image_filenames[idx] - num_images = len(image_filenames) - - idx_test = np.arange(num_images)[:: self.config.val_skip] - idx_train = np.array([i for i in np.arange(num_images) if i not in idx_test]) - idx = idx_train if split == "train" else idx_test - if num_images != poses.shape[0]: - raise RuntimeError(f"Different number of images ({num_images}), and poses ({poses.shape[0]})") - - image_filenames = image_filenames[idx] - poses = poses[idx] - - # convert to Tensors - poses = torch.from_numpy(poses[:, :3, :4]) - - poses = camera_utils.auto_orient_and_center_poses( - pose_utils.to4x4(poses), method=self.config.orientation_method - )[:, :3, :4] - - # Centering poses - poses[:, :3, 3] = poses[:, :3, 3] - torch.mean(poses[:, :3, 3], dim=0) - poses = pose_utils.normalize(poses) - - # Camera intrinsics - K = np.array(metadata_dict["K"]).reshape((3, 3)).T - focal_length = K[0, 0] - - H = metadata_dict["h"] - W = metadata_dict["w"] - - # TODO(akristoffersen): The metadata dict comes with principle points, - # but caused errors in image coord indexing. Should update once that is fixed. - cx, cy = W / 2, H / 2 - - num_cameras = len(image_filenames) - num_intrinsics_params = 3 - intrinsics = torch.ones((num_cameras, num_intrinsics_params), dtype=torch.float32) - intrinsics *= torch.tensor([cx, cy, focal_length]) - - aabb = torch.tensor([[-1, -1, -1], [1, 1, 1]], dtype=torch.float32) * self.config.aabb_scale - scene_box = SceneBox(aabb=aabb) - - cameras = Cameras( - fx=focal_length, - fy=focal_length, - cx=cx, - cy=cy, - camera_to_worlds=poses, - camera_type=CameraType.PERSPECTIVE, - ) - - dataparser_outputs = DataparserOutputs( - image_filenames=image_filenames, - cameras=cameras, - scene_box=scene_box, - ) - - return dataparser_outputs diff --git a/nerfstudio/data/pixel_samplers.py b/nerfstudio/data/pixel_samplers.py index 761485572..2ff1f1c96 100644 --- a/nerfstudio/data/pixel_samplers.py +++ b/nerfstudio/data/pixel_samplers.py @@ -49,7 +49,7 @@ def collate_image_dataset_batch(batch: Dict, num_rays_per_batch: int, sample_all else: # only sample within the mask, if the mask is in the batch if "mask" in batch: - nonzero_indices = torch.nonzero(batch["mask"][..., 0].to(device), as_tuple=False) + nonzero_indices = torch.nonzero(batch["mask"][..., 0], as_tuple=False) chosen_indices = random.sample(range(len(nonzero_indices)), k=num_rays_per_batch) indices = nonzero_indices[chosen_indices] else: diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index ae9da1ec2..904aa7577 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -28,13 +28,13 @@ from rich.console import Console from torch.cuda.amp.grad_scaler import GradScaler from typing_extensions import Literal -from nerfstudio.configs import base_config as cfg +from nerfstudio.configs.experiment_config import ExperimentConfig from nerfstudio.engine.callbacks import ( TrainingCallback, TrainingCallbackAttributes, TrainingCallbackLocation, ) -from nerfstudio.engine.optimizers import Optimizers, setup_optimizers +from nerfstudio.engine.optimizers import Optimizers from nerfstudio.pipelines.base_pipeline import VanillaPipeline from nerfstudio.utils import profiler, writer from nerfstudio.utils.decorators import ( @@ -71,7 +71,7 @@ class Trainer: optimizers: Optimizers callbacks: List[TrainingCallback] - def __init__(self, config: cfg.Config, local_rank: int = 0, world_size: int = 1): + def __init__(self, config: ExperimentConfig, local_rank: int = 0, world_size: int = 1): self.config = config self.local_rank = local_rank self.world_size = world_size @@ -96,7 +96,7 @@ class Trainer: self._check_viewer_warnings() # set up writers/profilers if enabled writer_log_path = self.base_dir / config.logging.relative_log_dir - writer.setup_event_writer(config, log_dir=writer_log_path) + writer.setup_event_writer(config.is_wandb_enabled(), config.is_tensorboard_enabled(), log_dir=writer_log_path, wandb_name=config.wandb_name) writer.setup_local_writer( config.logging, max_iter=config.trainer.max_num_iterations, banner_messages=banner_messages ) @@ -115,7 +115,7 @@ class Trainer: self.pipeline = self.config.pipeline.setup( device=self.device, test_mode=test_mode, world_size=self.world_size, local_rank=self.local_rank ) - self.optimizers = setup_optimizers(self.config, self.pipeline.get_param_groups()) + self.optimizers = self.setup_optimizers() self._load_checkpoint() @@ -128,10 +128,31 @@ class Trainer: ) ) + def setup_optimizers(self) -> Optimizers: + """Helper to set up the optimizers + + Returns: + The optimizers object given the experiment config. + """ + optimizer_config = self.config.optimizers.copy() + camera_optimizer_config = self.config.pipeline.datamanager.camera_optimizer + param_groups = self.pipeline.get_param_groups() + if camera_optimizer_config.mode != "off": + assert camera_optimizer_config.param_group not in optimizer_config + optimizer_config[camera_optimizer_config.param_group] = { + "optimizer": camera_optimizer_config.optimizer, + "scheduler": camera_optimizer_config.scheduler, + } + return Optimizers(optimizer_config, param_groups) + def train(self) -> None: """Train the model.""" assert self.pipeline.datamanager.train_dataset is not None, "Missing DatsetInputs" + self.pipeline.datamanager.train_dataparser_outputs.save_dataparser_transform( + self.base_dir / "dataparser_transforms.json" + ) + self._init_viewer_state() with TimeWriter(writer, EventName.TOTAL_TRAIN_TIME): num_iterations = self.config.trainer.max_num_iterations diff --git a/nerfstudio/exporter/exporter_utils.py b/nerfstudio/exporter/exporter_utils.py index c3db219ed..3e5f4425e 100644 --- a/nerfstudio/exporter/exporter_utils.py +++ b/nerfstudio/exporter/exporter_utils.py @@ -39,7 +39,6 @@ from rich.progress import ( from torchtyping import TensorType from nerfstudio.cameras.cameras import Cameras -from nerfstudio.configs.base_config import Config # pylint: disable=unused-import from nerfstudio.pipelines.base_pipeline import Pipeline from nerfstudio.utils.rich_utils import ItersPerSecColumn diff --git a/nerfstudio/field_components/spatial_distortions.py b/nerfstudio/field_components/spatial_distortions.py index 449f4b5b9..02da548a3 100644 --- a/nerfstudio/field_components/spatial_distortions.py +++ b/nerfstudio/field_components/spatial_distortions.py @@ -65,12 +65,8 @@ class SceneContraction(SpatialDistortion): def forward(self, positions): def contract(x): - mag = torch.linalg.norm(x, ord=self.order, dim=-1) - mask = mag >= 1 - x_new = x.clone() - x_new[mask] = (2 - (1 / mag[mask][..., None])) * (x[mask] / mag[mask][..., None]) - - return x_new + mag = torch.linalg.norm(x, ord=self.order, dim=-1)[..., None] + return torch.where(mag < 1, x, (2 - (1 / mag)) * (x / mag)) if isinstance(positions, Gaussians): means = contract(positions.mean.clone()) diff --git a/nerfstudio/fields/base_field.py b/nerfstudio/fields/base_field.py index 043a29e01..c24d9659c 100644 --- a/nerfstudio/fields/base_field.py +++ b/nerfstudio/fields/base_field.py @@ -74,10 +74,15 @@ class Field(nn.Module): self._sample_locations.shape[:-1] == self._density_before_activation.shape[:-1] ), "Sample locations and density must have the same shape besides the last dimension." - self._density_before_activation.backward( - gradient=torch.ones_like(self._density_before_activation), inputs=self._sample_locations, retain_graph=True - ) - normals = -torch.nn.functional.normalize(self._sample_locations.grad, dim=-1) + normals = torch.autograd.grad( + self._density_before_activation, + self._sample_locations, + grad_outputs=torch.ones_like(self._density_before_activation), + retain_graph=True, + )[0] + + normals = -torch.nn.functional.normalize(normals, dim=-1) + return normals @abstractmethod diff --git a/nerfstudio/fields/nerfacto_field.py b/nerfstudio/fields/nerfacto_field.py index ea215b517..1a8c425d6 100644 --- a/nerfstudio/fields/nerfacto_field.py +++ b/nerfstudio/fields/nerfacto_field.py @@ -91,7 +91,7 @@ class TCNNNerfactoField(Field): hidden_dim: int = 64, geo_feat_dim: int = 15, num_levels: int = 16, - max_res: int = 1024, + max_res: int = 2048, log2_hashmap_size: int = 19, num_layers_color: int = 3, num_layers_transient: int = 2, diff --git a/nerfstudio/model_components/ray_samplers.py b/nerfstudio/model_components/ray_samplers.py index a65bd2128..4bc16fc05 100644 --- a/nerfstudio/model_components/ray_samplers.py +++ b/nerfstudio/model_components/ray_samplers.py @@ -460,7 +460,8 @@ class VolumetricSampler(Sampler): rays_d=rays_d, scene_aabb=self.scene_aabb, grid=self.occupancy_grid, - sigma_fn=self.get_sigma_fn(rays_o, rays_d), + # this is a workaround - using density causes crash and damage quality. should be fixed + sigma_fn=None, # self.get_sigma_fn(rays_o, rays_d), render_step_size=render_step_size, near_plane=near_plane, far_plane=far_plane, diff --git a/nerfstudio/models/instant_ngp.py b/nerfstudio/models/instant_ngp.py index 2b9f2e071..cff57963e 100644 --- a/nerfstudio/models/instant_ngp.py +++ b/nerfstudio/models/instant_ngp.py @@ -28,6 +28,7 @@ from torch.nn import Parameter from torchmetrics import PeakSignalNoiseRatio from torchmetrics.functional import structural_similarity_index_measure from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity +from typing_extensions import Literal from nerfstudio.cameras.rays import RayBundle from nerfstudio.engine.callbacks import ( @@ -76,8 +77,8 @@ class InstantNGPModelConfig(ModelConfig): """How far along ray to stop sampling.""" use_appearance_embedding: bool = False """Whether to use an appearance embedding.""" - randomize_background: bool = True - """Whether to randomize the background color.""" + background_color: Literal["random", "black", "white"] = "random" + """The color that is given to untrained areas.""" class NGPModel(Model): @@ -122,7 +123,10 @@ class NGPModel(Model): ) # renderers - background_color = "random" if self.config.randomize_background else colors.WHITE + background_color = "random" + if self.config.background_color in ["white", "black"]: + background_color = colors.COLORS_DICT[self.config.background_color] + self.renderer_rgb = RGBRenderer(background_color=background_color) self.renderer_accumulation = AccumulationRenderer() self.renderer_depth = DepthRenderer(method="expected") diff --git a/nerfstudio/process_data/process_data_utils.py b/nerfstudio/process_data/process_data_utils.py index e291d898a..e028bf1c0 100644 --- a/nerfstudio/process_data/process_data_utils.py +++ b/nerfstudio/process_data/process_data_utils.py @@ -167,6 +167,10 @@ def copy_images(data: Path, image_dir: Path, verbose) -> int: allowed_exts = [".jpg", ".jpeg", ".png", ".tif", ".tiff"] image_paths = sorted([p for p in data.glob("[!.]*") if p.suffix.lower() in allowed_exts]) + if len(image_paths) == 0: + CONSOLE.log("[bold red]:skull: No usable images in the data folder.") + sys.exit(1) + num_frames = len(copy_images_list(image_paths, image_dir, verbose)) return num_frames diff --git a/nerfstudio/utils/eval_utils.py b/nerfstudio/utils/eval_utils.py index 7401a916c..8d780cc82 100644 --- a/nerfstudio/utils/eval_utils.py +++ b/nerfstudio/utils/eval_utils.py @@ -27,13 +27,14 @@ import yaml from rich.console import Console from typing_extensions import Literal -from nerfstudio.configs import base_config as cfg +from nerfstudio.configs.base_config import TrainerConfig +from nerfstudio.configs.experiment_config import ExperimentConfig from nerfstudio.pipelines.base_pipeline import Pipeline CONSOLE = Console(width=120) -def eval_load_checkpoint(config: cfg.TrainerConfig, pipeline: Pipeline) -> Path: +def eval_load_checkpoint(config: TrainerConfig, pipeline: Pipeline) -> Path: ## TODO: ideally eventually want to get this to be the same as whatever is used to load train checkpoint too """Helper function to load checkpointed pipeline @@ -68,7 +69,7 @@ def eval_setup( config_path: Path, eval_num_rays_per_chunk: Optional[int] = None, test_mode: Literal["test", "val", "inference"] = "test", -) -> Tuple[cfg.Config, Pipeline, Path]: +) -> Tuple[ExperimentConfig, Pipeline, Path]: """Shared setup for loading a saved pipeline for evaluation. Args: @@ -85,7 +86,7 @@ def eval_setup( """ # load save config config = yaml.load(config_path.read_text(), Loader=yaml.Loader) - assert isinstance(config, cfg.Config) + assert isinstance(config, ExperimentConfig) if eval_num_rays_per_chunk: config.pipeline.model.eval_num_rays_per_chunk = eval_num_rays_per_chunk diff --git a/nerfstudio/utils/writer.py b/nerfstudio/utils/writer.py index 077c99b3f..a057a335c 100644 --- a/nerfstudio/utils/writer.py +++ b/nerfstudio/utils/writer.py @@ -188,7 +188,7 @@ def setup_local_writer(config: cfg.LoggingConfig, max_iter: int, banner_messages @check_main_thread -def setup_event_writer(config: cfg.Config, log_dir: Path) -> None: +def setup_event_writer(is_wandb_enabled: bool, is_tensorboard_enabled: bool, log_dir: Path, wandb_name: Optional[str] = None) -> None: """Initialization of all event writers specified in config Args: @@ -197,11 +197,11 @@ def setup_event_writer(config: cfg.Config, log_dir: Path) -> None: banner_messages: list of messages to always display at bottom of screen """ using_event_writer = False - if config.is_wandb_enabled(): - curr_writer = WandbWriter(log_dir=log_dir, wandb_name=config.wandb_name) + if is_wandb_enabled: + curr_writer = WandbWriter(log_dir=log_dir, wandb_name=wandb_name) EVENT_WRITERS.append(curr_writer) using_event_writer = True - if config.is_tensorboard_enabled(): + if is_tensorboard_enabled: curr_writer = TensorboardWriter(log_dir=log_dir) EVENT_WRITERS.append(curr_writer) using_event_writer = True @@ -385,12 +385,12 @@ class LocalWriter: Args: step: current train step """ - valid_step = step > 0 and step % GLOBAL_BUFFER["steps_per_log"] == 0 + valid_step = step % GLOBAL_BUFFER["steps_per_log"] == 0 if valid_step: if not self.has_printed and self.config.max_log_size: CONSOLE.log( f"Printing max of {self.config.max_log_size} lines. " - "Set flag รงรงรงรง[yellow]--logging.local-writer.max-log-size=0[/yellow] " + "Set flag [yellow]--logging.local-writer.max-log-size=0[/yellow] " "to disable line wrapping." ) latest_map, new_key = self._consolidate_events() diff --git a/nerfstudio/viewer/app/src/index.scss b/nerfstudio/viewer/app/src/index.scss index 194c5ac2e..fab8b8a68 100644 --- a/nerfstudio/viewer/app/src/index.scss +++ b/nerfstudio/viewer/app/src/index.scss @@ -137,7 +137,10 @@ $sidebar-width: 100% !default; } .StatusPanel .StatusPanel-hide-scene-button { - margin: 12px; + margin-top: 4px; + margin-bottom: 4px; + margin-left: 6px; + margin-right: 6px; } .StatusPanel .WebSocketUrlField { @@ -242,18 +245,37 @@ $sidebar-width: 100% !default; // ------------------------------------ // Camera Pane -.CameraPanel { +.CameraPanel-props { + margin-top: 10px; + margin-bottom: 10px; + margin-left: 20px; + margin-right: 20px; +} + +.CameraPanel-path-row { + display: inline-block; + width: 45%; + margin-top: 8px; + margin-bottom: 8px; + margin-left: 4px; + margin-right: 4px; +} + +.CameraPanel-row { + display: inline-block; + margin-top: 8px; + margin-bottom: 8px; + margin-left: 4px; + margin-right: 4px; } .CameraPanel-top-button { - display: inline-block; - margin: 4px; - margin-top: 8px; + width: 100%; } .CameraPanel-render-button { display: inline-block; - width: 70%; + width: 92%; margin: 4px; margin-top: 8px; } diff --git a/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPanel.jsx b/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPanel.jsx index b71266d4a..9958ea7b1 100644 --- a/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPanel.jsx +++ b/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPanel.jsx @@ -37,6 +37,7 @@ import { useDispatch, useSelector } from 'react-redux'; import { CSS2DObject } from 'three/examples/jsm/renderers/CSS2DRenderer'; import AddAPhotoIcon from '@mui/icons-material/AddAPhoto'; import FileDownloadOutlinedIcon from '@mui/icons-material/FileDownloadOutlined'; +import { LevaPanel, LevaStoreProvider, useCreateStore } from 'leva'; import FileUploadOutlinedIcon from '@mui/icons-material/FileUploadOutlined'; import { Stack } from '@mui/system'; import TextField from '@mui/material/TextField'; @@ -46,6 +47,8 @@ import { CameraHelper } from './CameraHelper'; import { get_curve_object_from_cameras, get_transform_matrix } from './curve'; import { WebSocketContext } from '../../WebSocket/WebSocket'; import RenderModal from '../../RenderModal'; +import CameraPropPanel from './CameraPropPanel'; +import LevaTheme from '../../../themes/leva_theme.json'; const msgpack = require('msgpack-lite'); @@ -460,6 +463,9 @@ export default function CameraPanel(props) { const [animate, setAnimate] = React.useState(new Set()); const [globalFov, setGlobalFov] = React.useState(DEFAULT_FOV); + // leva store + const cameraPropsStore = useCreateStore(); + const scene_state = sceneTree.get_scene_state(); // Template for sharing state between Vanilla JS Three.js and React components @@ -496,10 +502,6 @@ export default function CameraPanel(props) { }; // ui state - const [ui_render_height, setUIRenderHeight] = React.useState(render_height); - const [ui_render_width, setUIRenderWidth] = React.useState(render_width); - const [ui_seconds, setUISeconds] = React.useState(seconds); - const [ui_fps, setUIfps] = React.useState(fps); const [fovLabel, setFovLabel] = React.useState(FOV_LABELS.FOV); // nonlinear render option @@ -612,6 +614,11 @@ export default function CameraPanel(props) { sceneTree.delete(['Camera Path', 'Cameras']); // delete old cameras, which is important if (cameras.length < 1) { + dispatch({ + type: 'write', + path: 'renderingState/camera_choice', + data: 'Main Camera', + }); camera_render_helper.set_visibility(false); } else { camera_render_helper.set_visibility(true); @@ -625,7 +632,7 @@ export default function CameraPanel(props) { labelDiv.className = 'label'; labelDiv.textContent = camera.properties.get('NAME'); labelDiv.style.color = 'black'; - labelDiv.style.backgroundColor = "rgba(255, 255, 255, 0.61)"; + labelDiv.style.backgroundColor = 'rgba(255, 255, 255, 0.61)'; labelDiv.style.backdropFilter = 'blur(5px)'; labelDiv.style.padding = '6px'; labelDiv.style.borderRadius = '6px'; @@ -863,15 +870,10 @@ export default function CameraPanel(props) { const new_properties = new Map(cameraProperties); setRenderHeight(camera_path_object.render_height); - setUIRenderHeight(camera_path_object.render_height); setRenderWidth(camera_path_object.render_width); - setUIRenderWidth(camera_path_object.render_width); setFps(camera_path_object.fps); - setUIfps(camera_path_object.fps); - setSeconds(camera_path_object.seconds); - setUISeconds(camera_path_object.seconds); set_smoothness_value(camera_path_object.smoothness_value); setIsCycle(camera_path_object.is_cycle); @@ -963,7 +965,7 @@ export default function CameraPanel(props) { return (
-
+
-
+
-
- { - if (e.target.validity.valid) { - setUIRenderHeight(e.target.value); - } - }} - onBlur={(e) => { - if (e.target.validity.valid) { - if (e.target.value !== '') { - setRenderHeight(parseInt(e.target.value, 10)); - } else { - setUIRenderHeight(render_height); - } - } - }} - value={ui_render_height} - error={ui_render_height <= 0} - helperText={ui_render_height <= 0 ? 'Required' : ''} - variant="standard" - /> - { - if (e.target.validity.valid) { - setUIRenderWidth(e.target.value); - } - }} - onBlur={(e) => { - if (e.target.validity.valid) { - if (e.target.value !== '') { - setRenderWidth(parseInt(e.target.value, 10)); - } else { - setUIRenderWidth(render_width); - } - } - }} - value={ui_render_width} - error={ui_render_width <= 0} - helperText={ui_render_width <= 0 ? 'Required' : ''} - variant="standard" - /> -
-
- { - if (e.target.validity.valid) { - setUISeconds(e.target.value); - } - }} - onBlur={(e) => { - if (e.target.validity.valid) { - if (e.target.value !== '') { - setSeconds(parseInt(e.target.value, 10)); - } else { - setUISeconds(seconds); - } - } - }} - value={ui_seconds} - error={ui_seconds <= 0} - helperText={ui_seconds <= 0 ? 'Required' : ''} - variant="standard" - /> - { - if (e.target.validity.valid) { - setUIfps(e.target.value); - } - }} - onBlur={(e) => { - if (e.target.validity.valid) { - if (e.target.value !== '') { - setFps(parseInt(e.target.value, 10)); - } else { - setUIfps(fps); - } - } - }} - value={ui_fps} - error={ui_fps <= 0} - helperText={ui_fps <= 0 ? 'Required' : ''} - variant="standard" +
+ + + +
@@ -1156,7 +1071,7 @@ export default function CameraPanel(props) { />
-
+
-
+
-
+
@@ -115,7 +119,8 @@ export default function StatusPanel(props: StatusPanelProps) { set_is_images_visible(!is_images_visible); }} style={{ textTransform: 'none' }} - startIcon={} + startIcon={} + disabled={camera_choice === 'Render Camera'} > {cameras_button} diff --git a/nerfstudio/viewer/app/src/modules/ViewerWindow/ViewerWindow.jsx b/nerfstudio/viewer/app/src/modules/ViewerWindow/ViewerWindow.jsx index 90fe975e0..94990f2d9 100644 --- a/nerfstudio/viewer/app/src/modules/ViewerWindow/ViewerWindow.jsx +++ b/nerfstudio/viewer/app/src/modules/ViewerWindow/ViewerWindow.jsx @@ -9,6 +9,8 @@ import OpenWithIcon from '@mui/icons-material/OpenWith'; import PublicOffSharpIcon from '@mui/icons-material/PublicOffSharp'; import PublicSharpIcon from '@mui/icons-material/PublicSharp'; import SyncOutlinedIcon from '@mui/icons-material/SyncOutlined'; +import ThreeDRotationIcon from '@mui/icons-material/ThreeDRotation'; +import VideoCameraBackIcon from '@mui/icons-material/VideoCameraBackOutlined'; import WebRtcWindow from '../WebRtcWindow/WebRtcWindow'; import { WebSocketContext } from '../WebSocket/WebSocket'; @@ -19,27 +21,36 @@ function CameraToggle() { const camera_choice = useSelector( (state) => state.renderingState.camera_choice, ); - const set_camera_choice = (event: SelectChangeEvent) => { - const value = event.target.value; - dispatch({ - type: 'write', - path: 'renderingState/camera_choice', - data: value, - }); + const set_camera_choice = (event: SelectChangeEvent, value: string[]) => { + if (value != null) { + dispatch({ + type: 'write', + path: 'renderingState/camera_choice', + data: value, + }); + } }; return ( - Viewport - Render View + + + Viewport + + + + Render View + ); } diff --git a/nerfstudio/viewer/app/src/themes/theme.ts b/nerfstudio/viewer/app/src/themes/theme.ts index 554445561..4a973209e 100644 --- a/nerfstudio/viewer/app/src/themes/theme.ts +++ b/nerfstudio/viewer/app/src/themes/theme.ts @@ -126,10 +126,13 @@ export const appTheme = createTheme({ styleOverrides: { root: { color: '#eeeeee', - backgroundColor: '#393e46', + backgroundColor: 'FFD369', '&.Mui-selected': { - color: '#FFD369', - backgroundColor: '#555555', + color: '#222831', + backgroundColor: '#FFD369', + ':hover': { + backgroundColor: '#FFD369', + }, }, '&.Mui-disabled': { color: '#999999', diff --git a/nerfstudio/viewer/server/utils.py b/nerfstudio/viewer/server/utils.py index 92ccc0ef7..6d13f9378 100644 --- a/nerfstudio/viewer/server/utils.py +++ b/nerfstudio/viewer/server/utils.py @@ -54,6 +54,9 @@ def three_js_perspective_camera_focal_length(fov: float, image_height: int): fov: the field of view of the camera in degrees. image_height: the height of the image in pixels. """ + if fov is None: + print("Warning: fov is None, using default value") + return 50 pp_h = image_height / 2.0 focal_length = pp_h / np.tan(fov * (np.pi / 180.0) / 2.0) return focal_length diff --git a/nerfstudio/viewer/server/viewer_utils.py b/nerfstudio/viewer/server/viewer_utils.py index 05be19ea0..c3ebb950e 100644 --- a/nerfstudio/viewer/server/viewer_utils.py +++ b/nerfstudio/viewer/server/viewer_utils.py @@ -454,6 +454,19 @@ class ViewerState: else: self.prev_camera_matrix = camera_object["matrix"] self.camera_moving = True + + output_type = self.vis["renderingState/output_choice"].read() + if output_type is None: + output_type = OutputTypes.INIT + if self.prev_output_type != output_type: + self.camera_moving = True + + colormap_type = self.vis["renderingState/colormap_choice"].read() + if colormap_type is None: + colormap_type = ColormapTypes.INIT + if self.prev_colormap_type != colormap_type: + self.camera_moving = True + return camera_object def _apply_colormap(self, outputs: Dict[str, Any], colors: torch.Tensor = None, eps=1e-6): @@ -672,9 +685,12 @@ class ViewerState: aspect_ratio = camera_object["aspect"] - image_height = (num_vis_rays / aspect_ratio) ** 0.5 - image_height = int(round(image_height, -1)) - image_height = min(self.max_resolution, image_height) + if not self.camera_moving and not is_training: + image_height = self.max_resolution + else: + image_height = (num_vis_rays / aspect_ratio) ** 0.5 + image_height = int(round(image_height, -1)) + image_height = min(self.max_resolution, image_height) image_width = int(image_height * aspect_ratio) if image_width > self.max_resolution: image_width = self.max_resolution diff --git a/pyproject.toml b/pyproject.toml index 064ff1cf8..fe908f8c6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "nerfstudio" -version = "0.1.12" +version = "0.1.13" description = "All-in-one repository for state-of-the-art NeRFs" readme = "README.md" license = { text="Apache 2.0"} diff --git a/scripts/benchmarking/launch_train_blender.sh b/scripts/benchmarking/launch_train_blender.sh index e57561e02..332ec95ef 100755 --- a/scripts/benchmarking/launch_train_blender.sh +++ b/scripts/benchmarking/launch_train_blender.sh @@ -62,6 +62,13 @@ timestamp=$(date "+%Y-%m-%d_%H%M%S") # kill all the background jobs if terminated: trap "trap - SIGTERM && kill -- -$$" SIGINT SIGTERM EXIT +dataparser="blender-data" +trans_file="" +if [ "$method_name" = "instant-ngp-bounded" ]; then + dataparser="" + trans_file="/transforms_train.json" +fi + for dataset in "${DATASETS[@]}"; do if "$single" && [ -n "${GPU_PID[$idx]+x}" ]; then echo "Waiting for GPU ${GPU_IDX[$idx]}" @@ -70,7 +77,7 @@ for dataset in "${DATASETS[@]}"; do fi export CUDA_VISIBLE_DEVICES="${GPU_IDX[$idx]}" ns-train "${method_name}" "${method_opts[@]}" \ - --data="data/blender/${dataset}" \ + --data="data/blender/${dataset}${trans_file}" \ --experiment-name="blender_${dataset}_${tag}" \ --trainer.relative-model-dir=nerfstudio_models/ \ --trainer.steps-per-save=1000 \ @@ -79,7 +86,7 @@ for dataset in "${DATASETS[@]}"; do --logging.enable-profiler=False \ --vis "${vis}" \ --timestamp "$timestamp" \ - blender-data & GPU_PID[$idx]=$! + ${dataparser} & GPU_PID[$idx]=$! echo "Launched ${method_name} ${dataset} on gpu ${GPU_IDX[$idx]}, ${tag}" # update gpu diff --git a/scripts/downloads/download_data.py b/scripts/downloads/download_data.py index b26886e89..3926caa00 100755 --- a/scripts/downloads/download_data.py +++ b/scripts/downloads/download_data.py @@ -15,6 +15,7 @@ from rich.console import Console from typing_extensions import Annotated from nerfstudio.configs.base_config import PrintableConfig +from nerfstudio.utils.scripts import run_command CONSOLE = Console(width=120) @@ -84,17 +85,39 @@ def grab_file_id(zip_url: str) -> str: return s.split("/")[0] +nerfstudio_dataset = [ + "Egypt", + "person", + "kitchen", + "plane", + "dozer", + "floating-tree", + "aspen", + "stump", + "sculpture", + "Giannini-Hall", +] nerfstudio_file_ids = { "bww_entrance": grab_file_id("https://drive.google.com/file/d/1ylkRHtfB3n3IRLf2wplpfxzPTq7nES9I/view?usp=sharing"), "campanile": grab_file_id("https://drive.google.com/file/d/13aOfGJRRH05pOOk9ikYGTwqFc2L1xskU/view?usp=sharing"), "desolation": grab_file_id("https://drive.google.com/file/d/14IzOOQm9KBJ3kPbunQbUTHPnXnmZus-f/view?usp=sharing"), - "dozer": grab_file_id("https://drive.google.com/file/d/1-OR5F_V5S4s-yzxohbwTylaXjzYLu8ZR/view?usp=sharing"), "library": grab_file_id("https://drive.google.com/file/d/1Hjbh_-BuaWETQExn2x2qGD74UwrFugHx/view?usp=sharing"), "poster": grab_file_id("https://drive.google.com/file/d/1dmjWGXlJnUxwosN6MVooCDQe970PkD-1/view?usp=sharing"), "redwoods2": grab_file_id("https://drive.google.com/file/d/1rg-4NoXT8p6vkmbWxMOY6PSG4j3rfcJ8/view?usp=sharing"), "storefront": grab_file_id("https://drive.google.com/file/d/16b792AguPZWDA_YC4igKCwXJqW0Tb21o/view?usp=sharing"), "vegetation": grab_file_id("https://drive.google.com/file/d/1wBhLQ2odycrtU39y2akVurXEAt9SsVI3/view?usp=sharing"), + "Egypt": "https://data.nerf.studio/nerfstudio/Egypt.zip", + "person": "https://data.nerf.studio/nerfstudio/person.zip", + "kitchen": "https://data.nerf.studio/nerfstudio/kitchen.zip", + "plane": "https://data.nerf.studio/nerfstudio/plane.zip", + "dozer": "https://data.nerf.studio/nerfstudio/dozer.zip", + "floating-tree": "https://data.nerf.studio/nerfstudio/floating-tree.zip", + "aspen": "https://data.nerf.studio/nerfstudio/aspen.zip", + "stump": "https://data.nerf.studio/nerfstudio/stump.zip", + "sculpture": "https://data.nerf.studio/nerfstudio/sculpture.zip", + "Giannini-Hall": "https://data.nerf.studio/nerfstudio/Giannini-Hall.zip", "all": None, + "nerfstudio-dataset": nerfstudio_dataset, } NerfstudioCaptureName = tyro.extras.literal_type_from_choices(nerfstudio_file_ids.keys()) @@ -103,14 +126,33 @@ NerfstudioCaptureName = tyro.extras.literal_type_from_choices(nerfstudio_file_id def download_capture_name(save_dir: Path, dataset_name: str, capture_name: str, capture_name_to_file_id: dict): """Download specific captures a given dataset and capture name.""" - url = f"https://drive.google.com/uc?id={capture_name_to_file_id[capture_name]}" - target_path = str(save_dir / f"{dataset_name}/{capture_name}") - os.makedirs(target_path, exist_ok=True) - download_path = Path(f"{target_path}.zip") - tmp_path = str(save_dir / ".temp") - shutil.rmtree(tmp_path, ignore_errors=True) - os.makedirs(tmp_path, exist_ok=True) - gdown.download(url, output=str(download_path)) + file_id_or_zip_url = capture_name_to_file_id[capture_name] + if file_id_or_zip_url.endswith(".zip"): + url = file_id_or_zip_url # zip url + target_path = str(save_dir / f"{dataset_name}/{capture_name}") + os.makedirs(target_path, exist_ok=True) + download_path = Path(f"{target_path}.zip") + tmp_path = str(save_dir / ".temp") + shutil.rmtree(tmp_path, ignore_errors=True) + os.makedirs(tmp_path, exist_ok=True) + try: + os.remove(download_path) + except OSError: + pass + run_command(f"wget {url} -O {download_path}", verbose=True) + else: + url = f"https://drive.google.com/uc?id={file_id_or_zip_url}" # file id + target_path = str(save_dir / f"{dataset_name}/{capture_name}") + os.makedirs(target_path, exist_ok=True) + download_path = Path(f"{target_path}.zip") + tmp_path = str(save_dir / ".temp") + shutil.rmtree(tmp_path, ignore_errors=True) + os.makedirs(tmp_path, exist_ok=True) + try: + os.remove(download_path) + except OSError: + pass + gdown.download(url, output=str(download_path)) with zipfile.ZipFile(download_path, "r") as zip_ref: zip_ref.extractall(tmp_path) inner_folders = os.listdir(tmp_path) @@ -132,7 +174,13 @@ class NerfstudioDownload(DatasetDownload): """Download the nerfstudio dataset.""" if self.capture_name == "all": for capture_name in nerfstudio_file_ids: - if capture_name != "all": + if capture_name not in ("all", "nerfstudio-dataset"): + download_capture_name(save_dir, "nerfstudio", capture_name, nerfstudio_file_ids) + return + + if self.capture_name == "nerfstudio-dataset": + for capture_name in nerfstudio_dataset: + if capture_name not in ("all", "nerfstudio-dataset"): download_capture_name(save_dir, "nerfstudio", capture_name, nerfstudio_file_ids) return diff --git a/scripts/render.py b/scripts/render.py index 60154f8b7..c8bb4fdae 100644 --- a/scripts/render.py +++ b/scripts/render.py @@ -26,7 +26,6 @@ from typing_extensions import Literal, assert_never from nerfstudio.cameras.camera_paths import get_path_from_json, get_spiral_path from nerfstudio.cameras.cameras import Cameras -from nerfstudio.configs.base_config import Config # pylint: disable=unused-import from nerfstudio.pipelines.base_pipeline import Pipeline from nerfstudio.utils import install_checks from nerfstudio.utils.eval_utils import eval_setup diff --git a/scripts/train.py b/scripts/train.py index 37717a11c..910dbff1c 100755 --- a/scripts/train.py +++ b/scripts/train.py @@ -45,8 +45,8 @@ import tyro import yaml from rich.console import Console -from nerfstudio.configs import base_config as cfg from nerfstudio.configs.config_utils import convert_markup_to_ansi +from nerfstudio.configs.experiment_config import ExperimentConfig from nerfstudio.configs.method_configs import AnnotatedBaseConfigUnion from nerfstudio.engine.trainer import Trainer from nerfstudio.utils import comms, profiler @@ -74,7 +74,7 @@ def _set_random_seed(seed) -> None: torch.manual_seed(seed) -def train_loop(local_rank: int, world_size: int, config: cfg.Config, global_rank: int = 0): +def train_loop(local_rank: int, world_size: int, config: ExperimentConfig, global_rank: int = 0): """Main training function that sets up and runs the trainer per process Args: @@ -95,7 +95,7 @@ def _distributed_worker( num_gpus_per_machine: int, machine_rank: int, dist_url: str, - config: cfg.Config, + config: ExperimentConfig, timeout: timedelta = DEFAULT_TIMEOUT, ) -> Any: """Spawned distributed worker that handles the initialization of process group and handles the @@ -110,7 +110,7 @@ def _distributed_worker( dist_url: URL to connect to for distributed jobs, including protocol E.g., "tcp://127.0.0.1:8686". It can be set to "auto" to automatically select a free port on localhost. - config: Config specifying training regimen. + config: ExperimentConfig specifying training regimen. timeout: Timeout of the distributed workers. Raises: @@ -150,7 +150,7 @@ def launch( num_machines: int = 1, machine_rank: int = 0, dist_url: str = "auto", - config: Optional[cfg.Config] = None, + config: Optional[ExperimentConfig] = None, timeout: timedelta = DEFAULT_TIMEOUT, ) -> None: """Function that spawns muliple processes to call on main_func @@ -161,7 +161,7 @@ def launch( num_machines (int, optional): total number of machines machine_rank (int, optional): rank of this machine. dist_url (str, optional): url to connect to for distributed jobs. - config (Config, optional): config file specifying training regimen. + config (ExperimentConfig, optional): config file specifying training regimen. timeout (timedelta, optional): timeout of the distributed workers. """ assert config is not None @@ -215,7 +215,7 @@ def launch( profiler.flush_profiler(config.logging) -def main(config: cfg.Config) -> None: +def main(config: ExperimentConfig) -> None: """Main function.""" config.set_timestamp() diff --git a/tests/test_train.py b/tests/test_train.py index 6332ba260..f7e860275 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -9,15 +9,15 @@ from pathlib import Path import pytest -from nerfstudio.configs.base_config import Config +from nerfstudio.configs.experiment_config import ExperimentConfig from nerfstudio.configs.method_configs import method_configs from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig from scripts.train import train_loop -BLACKLIST = ["base", "semantic-nerfw", "instant-ngp", "nerfacto", "phototourism"] +BLACKLIST = ["base", "semantic-nerfw", "instant-ngp", "instant-ngp-bounded", "nerfacto", "phototourism"] -def set_reduced_config(config: Config): +def set_reduced_config(config: ExperimentConfig): """Reducing the config settings to speedup test""" config.machine.num_gpus = 0 config.trainer.max_num_iterations = 2