diff --git a/README.md b/README.md index ca2bd48b4..7ea21deb8 100644 --- a/README.md +++ b/README.md @@ -95,13 +95,25 @@ python -m pip install --upgrade pip ### Dependencies -Install pytorch with CUDA (this repo has been tested with CUDA 11.3) and [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn) +Install pytorch with CUDA (this repo has been tested with CUDA 11.3 and CUDA 11.7) and [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn) + +For CUDA 11.3: ```bash pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch ``` +For CUDA 11.7: + +```bash +pip install torch torchvision functorch --extra-index-url https://download.pytorch.org/whl/cu117 +pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch +``` + +See [Dependencies](https://github.com/nerfstudio-project/nerfstudio/blob/main/docs/quickstart/installation.md#dependencies) +in the Installation documentation for more. + ### Installing nerfstudio Easy option: @@ -293,13 +305,13 @@ We provide the following support structures to make life easier for getting star If you use this library or find the documentation useful for your research, please consider citing: ``` -@misc{nerfstudio, - title={Nerfstudio: A Framework for Neural Radiance Field Development}, - author={Matthew Tancik*, Ethan Weber*, Evonne Ng*, Ruilong Li, Brent Yi, - Terrance Wang, Alexander Kristoffersen, Jake Austin, Kamyar Salahi, - Abhik Ahuja, David McAllister, Angjoo Kanazawa}, - year={2022}, - url={https://github.com/nerfstudio-project/nerfstudio}, +@article{nerfstudio, + author = {Tancik, Matthew and Weber, Ethan and Ng, Evonne and Li, Ruilong and Yi, + Brent and Kerr, Justin and Wang, Terrance and Kristoffersen, Alexander and Austin, + Jake and Salahi, Kamyar and Ahuja, Abhik and McAllister, David and Kanazawa, Angjoo}, + title = {Nerfstudio: A Modular Framework for Neural Radiance Field Development}, + journal = {arXiv preprint arXiv:2302.04264}, + year = {2023}, } ``` diff --git a/colab/demo.ipynb b/colab/demo.ipynb index e92f8b04d..d1c58ad6c 100644 --- a/colab/demo.ipynb +++ b/colab/demo.ipynb @@ -292,9 +292,9 @@ "if os.path.exists(f\"data/nerfstudio/{scene}/transforms.json\"):\n", " !ns-train nerfacto --viewer.websocket-port 7007 nerfstudio-data --data data/nerfstudio/$scene --downscale-factor 4\n", "else:\n", - " import IPython\n", - " display(IPython.display.HTML('

Error: Data processing did not complete

'))\n", - " display(IPython.display.HTML('

Please re-run `Downloading and Processing Data`, or view the FAQ for more info.

'))" + " from IPython.core.display import display, HTML\n", + " display(HTML('

Error: Data processing did not complete

'))\n", + " display(HTML('

Please re-run `Downloading and Processing Data`, or view the FAQ for more info.

'))" ] }, { @@ -383,4 +383,4 @@ }, "nbformat": 4, "nbformat_minor": 0 -} \ No newline at end of file +} diff --git a/docs/developer_guides/pipelines/datamanagers.md b/docs/developer_guides/pipelines/datamanagers.md index bbe3413d9..32c55b63e 100644 --- a/docs/developer_guides/pipelines/datamanagers.md +++ b/docs/developer_guides/pipelines/datamanagers.md @@ -87,7 +87,7 @@ def next_train(self, step: int) -> Tuple[RayBundle, Dict]: You can see our code for more details. -```{button-link} https://github.com/nerfstudio-project/nerfstudio/blob/master/nerfstudio/data/datamanagers.py +```{button-link} https://github.com/nerfstudio-project/nerfstudio/blob/main/nerfstudio/data/datamanagers/base_datamanager.py :color: primary :outline: See the code! diff --git a/docs/index.md b/docs/index.md index 53d45784d..023b87cc8 100644 --- a/docs/index.md +++ b/docs/index.md @@ -158,13 +158,13 @@ We'll be constantly growing this list! So make sure to check back in to see our If you use this library or find the documentation useful for your research, please consider citing: ```none -@misc{nerfstudio, - title={Nerfstudio: A Framework for Neural Radiance Field Development}, - author={Matthew Tancik*, Ethan Weber*, Evonne Ng*, Ruilong Li, Brent Yi, - Terrance Wang, Alexander Kristoffersen, Jake Austin, Kamyar Salahi, - Abhik Ahuja, David McAllister, Angjoo Kanazawa}, - year={2022}, - url={https://github.com/nerfstudio-project/nerfstudio}, +@article{nerfstudio, + author = {Tancik, Matthew and Weber, Ethan and Ng, Evonne and Li, Ruilong and Yi, + Brent and Kerr, Justin and Wang, Terrance and Kristoffersen, Alexander and Austin, + Jake and Salahi, Kamyar and Ahuja, Abhik and McAllister, David and Kanazawa, Angjoo}, + title = {Nerfstudio: A Modular Framework for Neural Radiance Field Development}, + journal = {arXiv preprint arXiv:2302.04264}, + year = {2023}, } ``` diff --git a/docs/quickstart/first_nerf.md b/docs/quickstart/first_nerf.md index e7071c329..73e40a9d8 100644 --- a/docs/quickstart/first_nerf.md +++ b/docs/quickstart/first_nerf.md @@ -30,7 +30,7 @@ Navigating to the link at the end of the terminal will load the webviewer. If yo - All data configurations must go at the end. In this case, `nerfstudio-data` and all of its corresponding configurations come at the end after the model and viewer specification. ::: -## Resume from checkpoint / visualize existing run +## Resume from checkpoint It is possible to load a pretrained model by running @@ -40,6 +40,14 @@ ns-train nerfacto --data data/nerfstudio/poster --load-dir {outputs/.../nerfstud This will automatically start training. If you do not want it to train, add `--viewer.start-train False` to your training command. +## Visualize existing run + +Given a pretrained model checkpoint, you can start the viewer by running + +```bash +ns-viewer --load-config {outputs/.../config.yml} +``` + ## Exporting Results Once you have a NeRF model you can either render out a video or export a point cloud. diff --git a/docs/quickstart/installation.md b/docs/quickstart/installation.md index 97c63aac8..20ab8859c 100644 --- a/docs/quickstart/installation.md +++ b/docs/quickstart/installation.md @@ -17,12 +17,38 @@ python -m pip install --upgrade pip ## Dependencies -Install pytorch with CUDA (this repo has been tested with CUDA 11.3) and [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn) +(pytorch)= +### pytorch + +Either pytorch 1.12.1 (with CUDA up to 11.3) or 1.13.1 (with CUDA 11.6/7/8) can be used. + +- To install 1.12.1 with CUDA 11.3: ```bash pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html -pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch +``` +- To install 1.13.1 with CUDA 11.7: + +Note that if a pytorch version prior to 1.13 is installed, +it should be uninstalled first to avoid upgrade issues (e.g. with functorch) + +```bash +pip uninstall torch torchvision functorch +``` + +Install pytorch 1.13.1 with CUDA (this repo has been tested with CUDA 11.7.1) and [tiny-cuda-nn](https://github.com/NVlabs/tiny-cuda-nn) + +```bash +pip install torch torchvision functorch --extra-index-url https://download.pytorch.org/whl/cu117 +``` + +### tinycudann + +After pytorch, install the torch bindings for tinycudann: + +```bash +pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch ``` ## Installing nerfstudio @@ -134,9 +160,8 @@ While installing tiny-cuda, you run into: `The detected CUDA version mismatches **Solution**: -``` -pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html -``` +Reinstall pytorch with the correct CUDA version. +See [pytorch](pytorch) under Dependencies, above.
diff --git a/nerfstudio/cameras/camera_paths.py b/nerfstudio/cameras/camera_paths.py index 61c489edf..a43504694 100644 --- a/nerfstudio/cameras/camera_paths.py +++ b/nerfstudio/cameras/camera_paths.py @@ -100,12 +100,16 @@ def get_spiral_path( new_c2ws.append(c2wh[:3, :4]) new_c2ws = torch.stack(new_c2ws, dim=0) + times = None + if camera.times is not None: + times = torch.linspace(0, 1, steps)[:, None] return Cameras( fx=camera.fx[0], fy=camera.fy[0], cx=camera.cx[0], cy=camera.cy[0], camera_to_worlds=new_c2ws, + times=times, ) diff --git a/nerfstudio/cameras/cameras.py b/nerfstudio/cameras/cameras.py index d7d228bae..80b7f6487 100644 --- a/nerfstudio/cameras/cameras.py +++ b/nerfstudio/cameras/cameras.py @@ -230,7 +230,7 @@ class Cameras(TensorDataclass): c_x_y: cx or cy for when h_w == None """ if isinstance(h_w, int): - h_w = torch.Tensor([h_w]).to(torch.int64).to(self.device) + h_w = torch.as_tensor([h_w]).to(torch.int64).to(self.device) elif isinstance(h_w, torch.Tensor): assert not torch.is_floating_point(h_w), f"height and width tensor must be of type int, not: {h_w.dtype}" h_w = h_w.to(torch.int64).to(self.device) @@ -238,7 +238,7 @@ class Cameras(TensorDataclass): h_w = h_w.unsqueeze(-1) # assert torch.all(h_w == h_w.view(-1)[0]), "Batched cameras of different h, w will be allowed in the future." elif h_w is None: - h_w = torch.Tensor((c_x_y * 2).to(torch.int64).to(self.device)) + h_w = torch.as_tensor((c_x_y * 2)).to(torch.int64).to(self.device) else: raise ValueError("Height must be an int, tensor, or None, received: " + str(type(h_w))) return h_w diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index d032c4d48..1c8c3e08d 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -49,6 +49,7 @@ from nerfstudio.data.dataparsers.nuscenes_dataparser import NuScenesDataParserCo from nerfstudio.data.dataparsers.phototourism_dataparser import ( PhototourismDataParserConfig, ) +from nerfstudio.data.dataparsers.sdfstudio_dataparser import SDFStudioDataParserConfig from nerfstudio.data.datasets.base_dataset import InputDataset from nerfstudio.data.pixel_samplers import EquirectangularPixelSampler, PixelSampler from nerfstudio.data.utils.dataloaders import ( @@ -75,6 +76,7 @@ AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[ # Omit prefixes of "dnerf-data": DNeRFDataParserConfig(), "phototourism-data": PhototourismDataParserConfig(), "dycheck-data": DycheckDataParserConfig(), + "sdfstudio-data": SDFStudioDataParserConfig(), }, prefix_names=False, # Omit prefixes in subcommands themselves. ) diff --git a/nerfstudio/data/dataparsers/instant_ngp_dataparser.py b/nerfstudio/data/dataparsers/instant_ngp_dataparser.py index d09aa86f9..2fdc2df85 100644 --- a/nerfstudio/data/dataparsers/instant_ngp_dataparser.py +++ b/nerfstudio/data/dataparsers/instant_ngp_dataparser.py @@ -123,6 +123,10 @@ class InstantNGP(DataParser): w, h = meta["w"], meta["h"] + camera_type = CameraType.PERSPECTIVE + if meta.get("is_fisheye", False): + camera_type = CameraType.FISHEYE + cameras = Cameras( fx=float(fl_x), fy=float(fl_y), @@ -132,7 +136,7 @@ class InstantNGP(DataParser): height=int(h), width=int(w), camera_to_worlds=camera_to_world, - camera_type=CameraType.PERSPECTIVE, + camera_type=camera_type, ) # TODO(ethan): add alpha background color diff --git a/nerfstudio/data/dataparsers/sdfstudio_dataparser.py b/nerfstudio/data/dataparsers/sdfstudio_dataparser.py new file mode 100644 index 000000000..f25c8f3e7 --- /dev/null +++ b/nerfstudio/data/dataparsers/sdfstudio_dataparser.py @@ -0,0 +1,158 @@ +# 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. +"""Datapaser for sdfstudio formatted data""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Type + +import torch +from rich.console import Console + +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.io import load_from_json + +CONSOLE = Console() + + +@dataclass +class SDFStudioDataParserConfig(DataParserConfig): + """Scene dataset parser config""" + + _target: Type = field(default_factory=lambda: SDFStudio) + """target class to instantiate""" + data: Path = Path("data/DTU/scan65") + """Directory specifying location of data.""" + include_mono_prior: bool = False + """whether or not to load monocular depth and normal """ + include_foreground_mask: bool = False + """whether or not to load foreground mask""" + downscale_factor: int = 1 + scene_scale: float = 2.0 + """ + Sets the bounding cube to have edge length of this size. + The longest dimension of the Friends axis-aligned bbox will be scaled to this value. + """ + skip_every_for_val_split: int = 1 + """sub sampling validation images""" + auto_orient: bool = False + + +@dataclass +class SDFStudio(DataParser): + """SDFStudio Dataset""" + + config: SDFStudioDataParserConfig + + def _generate_dataparser_outputs(self, split="train"): # pylint: disable=unused-argument,too-many-statements + # load meta data + meta = load_from_json(self.config.data / "meta_data.json") + + indices = list(range(len(meta["frames"]))) + # subsample to avoid out-of-memory for validation set + if split != "train" and self.config.skip_every_for_val_split >= 1: + indices = indices[:: self.config.skip_every_for_val_split] + + image_filenames = [] + depth_filenames = [] + normal_filenames = [] + transform = None + fx = [] + fy = [] + cx = [] + cy = [] + camera_to_worlds = [] + for i, frame in enumerate(meta["frames"]): + if i not in indices: + continue + + image_filename = self.config.data / frame["rgb_path"] + depth_filename = self.config.data / frame["mono_depth_path"] + normal_filename = self.config.data / frame["mono_normal_path"] + + intrinsics = torch.tensor(frame["intrinsics"]) + camtoworld = torch.tensor(frame["camtoworld"]) + + # append data + image_filenames.append(image_filename) + depth_filenames.append(depth_filename) + normal_filenames.append(normal_filename) + fx.append(intrinsics[0, 0]) + fy.append(intrinsics[1, 1]) + cx.append(intrinsics[0, 2]) + cy.append(intrinsics[1, 2]) + camera_to_worlds.append(camtoworld) + + fx = torch.stack(fx) + fy = torch.stack(fy) + cx = torch.stack(cx) + cy = torch.stack(cy) + camera_to_worlds = torch.stack(camera_to_worlds) + + # Convert from COLMAP's/OPENCV's camera coordinate system to nerfstudio + camera_to_worlds[:, 0:3, 1:3] *= -1 + + if self.config.auto_orient: + camera_to_worlds, transform = camera_utils.auto_orient_and_center_poses( + camera_to_worlds, + method="up", + center_poses=False, + ) + + # scene box from meta data + meta_scene_box = meta["scene_box"] + aabb = torch.tensor(meta_scene_box["aabb"], dtype=torch.float32) + scene_box = SceneBox( + aabb=aabb, + ) + + height, width = meta["height"], meta["width"] + cameras = Cameras( + fx=fx, + fy=fy, + cx=cx, + cy=cy, + height=height, + width=width, + camera_to_worlds=camera_to_worlds[:, :3, :4], + camera_type=CameraType.PERSPECTIVE, + ) + + # TODO supports downsample + # cameras.rescale_output_resolution(scaling_factor=1.0 / self.config.downscale_factor) + + assert meta["has_mono_prior"] == self.config.include_mono_prior, f"no mono prior in {self.config.data}" + + dataparser_outputs = DataparserOutputs( + image_filenames=image_filenames, + cameras=cameras, + scene_box=scene_box, + metadata={ + "depth_filenames": depth_filenames if len(depth_filenames) > 0 else None, + "normal_filenames": normal_filenames if len(normal_filenames) > 0 else None, + "transform": transform, + "camera_to_worlds": camera_to_worlds if len(camera_to_worlds) > 0 else None, + "include_mono_prior": self.config.include_mono_prior, + }, + ) + return dataparser_outputs diff --git a/nerfstudio/data/datasets/sdf_dataset.py b/nerfstudio/data/datasets/sdf_dataset.py new file mode 100644 index 000000000..48e0cff6d --- /dev/null +++ b/nerfstudio/data/datasets/sdf_dataset.py @@ -0,0 +1,97 @@ +# 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. + +""" +SDFStudio dataset. +""" + +from pathlib import Path +from typing import Dict + +import numpy as np +import torch + +from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs +from nerfstudio.data.datasets.base_dataset import InputDataset + + +class SDFDataset(InputDataset): + """Dataset that returns images and depths. + + Args: + dataparser_outputs: description of where and how to read input images. + scale_factor: The scaling factor for the dataparser outputs. + """ + + def __init__(self, dataparser_outputs: DataparserOutputs, scale_factor: float = 1.0): + super().__init__(dataparser_outputs, scale_factor) + + # can be none if monoprior not included + self.depth_filenames = self.metadata["depth_filenames"] + self.normal_filenames = self.metadata["normal_filenames"] + self.camera_to_worlds = self.metadata["camera_to_worlds"] + # can be none if auto orient not enabled in dataparser + self.transform = self.metadata["normal_filenames"] + self.include_mono_prior = self.metadata["include_mono_prior"] + + def get_metadata(self, data: Dict) -> Dict: + # TODO supports foreground_masks + metadata = {} + if self.include_mono_prior: + depth_filepath = self.depth_filenames[data["image_idx"]] + normal_filepath = self.normal_filenames[data["image_idx"]] + camtoworld = self.camera_to_worlds[data["image_idx"]] + + # Scale depth images to meter units and also by scaling applied to cameras + depth_image, normal_image = self.get_depths_and_normals( + depth_filepath=depth_filepath, normal_filename=normal_filepath, camtoworld=camtoworld + ) + metadata["depth_image"] = depth_image + metadata["normal_image"] = normal_image + + return metadata + + def get_depths_and_normals(self, depth_filepath: Path, normal_filename: Path, camtoworld: np.ndarray): + """function to process additional depths and normal information + Args: + depth_filepath: path to depth file + normal_filename: path to normal file + camtoworld: camera to world transformation matrix + """ + + # load mono depth + depth = np.load(depth_filepath) + depth = torch.from_numpy(depth).float() + + # load mono normal + normal = np.load(normal_filename) + + # transform normal to world coordinate system + normal = normal * 2.0 - 1.0 # omnidata output is normalized so we convert it back to normal here + normal = torch.from_numpy(normal).float() + + rot = camtoworld[:3, :3] + + normal_map = normal.reshape(3, -1) + normal_map = torch.nn.functional.normalize(normal_map, p=2, dim=0) + + normal_map = rot @ normal_map + normal = normal_map.permute(1, 0).reshape(*normal.shape[1:], 3) + + if self.transform is not None: + h, w, _ = normal.shape + normal = self.transform[:3, :3] @ normal.reshape(-1, 3).permute(1, 0) + normal = normal.permute(1, 0).reshape(h, w, 3) + + return depth, normal diff --git a/nerfstudio/field_components/temporal_grid.py b/nerfstudio/field_components/temporal_grid.py index 6bccc110a..3bbe274b7 100644 --- a/nerfstudio/field_components/temporal_grid.py +++ b/nerfstudio/field_components/temporal_grid.py @@ -116,7 +116,6 @@ class TemporalGridEncodeFunc(Function): @staticmethod @custom_bwd def backward(ctx, grad): - inputs, temporal_row_index, embeddings, offsets, dy_dx = ctx.saved_tensors B, D, grid_channel, C, L, S, H, gridtype = ctx.dims align_corners = ctx.align_corners @@ -336,7 +335,7 @@ class TemporalGridEncoder(nn.Module): """ outputs = TemporalGridEncodeFunc.apply( xyz, - self.get_temporal_index(time[:, 0]), + self.get_temporal_index(time[:, 0].float()), self.embeddings, self.offsets, self.per_level_scale, diff --git a/nerfstudio/models/instant_ngp.py b/nerfstudio/models/instant_ngp.py index 736cd0870..9c0579031 100644 --- a/nerfstudio/models/instant_ngp.py +++ b/nerfstudio/models/instant_ngp.py @@ -137,7 +137,7 @@ class NGPModel(Model): # metrics self.psnr = PeakSignalNoiseRatio(data_range=1.0) self.ssim = structural_similarity_index_measure - self.lpips = LearnedPerceptualImagePatchSimilarity() + self.lpips = LearnedPerceptualImagePatchSimilarity(normalize=True) def get_training_callbacks( self, training_callback_attributes: TrainingCallbackAttributes diff --git a/nerfstudio/models/mipnerf.py b/nerfstudio/models/mipnerf.py index b8c3bcc4f..31a337d02 100644 --- a/nerfstudio/models/mipnerf.py +++ b/nerfstudio/models/mipnerf.py @@ -86,7 +86,7 @@ class MipNerfModel(Model): # metrics self.psnr = PeakSignalNoiseRatio(data_range=1.0) self.ssim = structural_similarity_index_measure - self.lpips = LearnedPerceptualImagePatchSimilarity() + self.lpips = LearnedPerceptualImagePatchSimilarity(normalize=True) def get_param_groups(self) -> Dict[str, List[Parameter]]: param_groups = {} @@ -173,8 +173,8 @@ class MipNerfModel(Model): image = torch.moveaxis(image, -1, 0)[None, ...] rgb_coarse = torch.moveaxis(rgb_coarse, -1, 0)[None, ...] rgb_fine = torch.moveaxis(rgb_fine, -1, 0)[None, ...] - rgb_coarse = torch.clip(rgb_coarse, min=-1, max=1) - rgb_fine = torch.clip(rgb_fine, min=-1, max=1) + rgb_coarse = torch.clip(rgb_coarse, min=0, max=1) + rgb_fine = torch.clip(rgb_fine, min=0, max=1) coarse_psnr = self.psnr(image, rgb_coarse) fine_psnr = self.psnr(image, rgb_fine) diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 7442bfe8b..4aa58d468 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -197,7 +197,7 @@ class NerfactoModel(Model): # metrics self.psnr = PeakSignalNoiseRatio(data_range=1.0) self.ssim = structural_similarity_index_measure - self.lpips = LearnedPerceptualImagePatchSimilarity() + self.lpips = LearnedPerceptualImagePatchSimilarity(normalize=True) def get_param_groups(self) -> Dict[str, List[Parameter]]: param_groups = {} diff --git a/nerfstudio/models/nerfplayer_nerfacto.py b/nerfstudio/models/nerfplayer_nerfacto.py index f35a748b1..3811aba67 100644 --- a/nerfstudio/models/nerfplayer_nerfacto.py +++ b/nerfstudio/models/nerfplayer_nerfacto.py @@ -168,7 +168,7 @@ class NerfplayerNerfactoModel(NerfactoModel): # metrics self.psnr = PeakSignalNoiseRatio(data_range=1.0) self.ssim = structural_similarity_index_measure - self.lpips = LearnedPerceptualImagePatchSimilarity() + self.lpips = LearnedPerceptualImagePatchSimilarity(normalize=True) self.temporal_distortion = True # for viewer def get_outputs(self, ray_bundle: RayBundle): diff --git a/nerfstudio/models/nerfplayer_ngp.py b/nerfstudio/models/nerfplayer_ngp.py index 39507a538..878694468 100644 --- a/nerfstudio/models/nerfplayer_ngp.py +++ b/nerfstudio/models/nerfplayer_ngp.py @@ -145,7 +145,7 @@ class NerfplayerNGPModel(NGPModel): # metrics self.psnr = PeakSignalNoiseRatio(data_range=1.0) self.ssim = structural_similarity_index_measure - self.lpips = LearnedPerceptualImagePatchSimilarity() + self.lpips = LearnedPerceptualImagePatchSimilarity(normalize=True) self.temporal_distortion = True # for viewer def get_outputs(self, ray_bundle: RayBundle): diff --git a/nerfstudio/models/semantic_nerfw.py b/nerfstudio/models/semantic_nerfw.py index 941716dc8..22957d268 100644 --- a/nerfstudio/models/semantic_nerfw.py +++ b/nerfstudio/models/semantic_nerfw.py @@ -137,7 +137,7 @@ class SemanticNerfWModel(Model): # metrics self.psnr = PeakSignalNoiseRatio(data_range=1.0) self.ssim = structural_similarity_index_measure - self.lpips = LearnedPerceptualImagePatchSimilarity() + self.lpips = LearnedPerceptualImagePatchSimilarity(normalize=True) def get_param_groups(self) -> Dict[str, List[Parameter]]: param_groups = {} diff --git a/nerfstudio/models/tensorf.py b/nerfstudio/models/tensorf.py index 60fec9706..d5af1579d 100644 --- a/nerfstudio/models/tensorf.py +++ b/nerfstudio/models/tensorf.py @@ -201,7 +201,7 @@ class TensoRFModel(Model): # metrics self.psnr = PeakSignalNoiseRatio(data_range=1.0) self.ssim = structural_similarity_index_measure - self.lpips = LearnedPerceptualImagePatchSimilarity() + self.lpips = LearnedPerceptualImagePatchSimilarity(normalize=True) # colliders if self.config.enable_collider: diff --git a/nerfstudio/models/vanilla_nerf.py b/nerfstudio/models/vanilla_nerf.py index 3e731524f..53e838891 100644 --- a/nerfstudio/models/vanilla_nerf.py +++ b/nerfstudio/models/vanilla_nerf.py @@ -118,7 +118,7 @@ class NeRFModel(Model): # metrics self.psnr = PeakSignalNoiseRatio(data_range=1.0) self.ssim = structural_similarity_index_measure - self.lpips = LearnedPerceptualImagePatchSimilarity() + self.lpips = LearnedPerceptualImagePatchSimilarity(normalize=True) if getattr(self.config, "enable_temporal_distortion", False): params = self.config.temporal_distortion_params diff --git a/pyproject.toml b/pyproject.toml index ba57c01d4..8901aa282 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "nerfstudio" -version = "0.1.16" +version = "0.1.17" description = "All-in-one repository for state-of-the-art NeRFs" readme = "README.md" license = { text="Apache 2.0"} @@ -22,7 +22,7 @@ dependencies = [ "tyro>=0.3.31", "gdown>=4.6.0", "ninja>=1.10", - "functorch==0.2.1", + "functorch>=0.2.1", "h5py>=2.9.0", "imageio>=2.21.1", 'importlib-metadata>=6.0.0; python_version < "3.10"', @@ -32,7 +32,7 @@ dependencies = [ "mediapy>=1.1.0", "msgpack>=1.0.4", "msgpack_numpy>=0.4.8", - "nerfacc==0.3.4", + "nerfacc==0.3.3", "open3d>=0.16.0", "opencv-python==4.6.0.66", "Pillow>=9.3.0", @@ -45,7 +45,7 @@ dependencies = [ "rich>=12.5.1", "scikit-image>=0.19.3", "tensorboard==2.9.0", - "torch==1.12.1", + "torch>=1.12.1", "torchmetrics[image]>=0.9.3", "torchtyping>=0.1.4", "torchvision>=0.13.0", @@ -100,6 +100,7 @@ ns-install-cli = "scripts.completions.install:entrypoint" ns-process-data = "scripts.process_data:entrypoint" ns-download-data = "scripts.downloads.download_data:entrypoint" ns-train = "scripts.train:entrypoint" +ns-viewer = "scripts.viewer.run_viewer:entrypoint" ns-eval = "scripts.eval:entrypoint" ns-render = "scripts.render:entrypoint" ns-export = "scripts.exporter:entrypoint" diff --git a/scripts/downloads/download_data.py b/scripts/downloads/download_data.py index 89c77391f..b6775a339 100755 --- a/scripts/downloads/download_data.py +++ b/scripts/downloads/download_data.py @@ -287,6 +287,75 @@ class PhototourismDownload(DatasetDownload): os.remove(download_path) +# credit to https://autonomousvision.github.io/sdfstudio/ +# pylint: disable=line-too-long +sdfstudio_downloads = { + "sdfstudio-demo-data": "https://s3.eu-central-1.amazonaws.com/avg-projects/monosdf/data/sdfstudio-demo-data.tar", + "dtu": "https://s3.eu-central-1.amazonaws.com/avg-projects/monosdf/data/DTU.tar", + "replica": "https://s3.eu-central-1.amazonaws.com/avg-projects/monosdf/data/Replica.tar", + "scannet": "https://s3.eu-central-1.amazonaws.com/avg-projects/monosdf/data/scannet.tar", + "tanks-and-temple": "https://s3.eu-central-1.amazonaws.com/avg-projects/monosdf/data/tnt_advanced.tar", + "tanks-and-temple-highres": "https://s3.eu-central-1.amazonaws.com/avg-projects/monosdf/data/highresTNT.tar", + "heritage": "https://s3.eu-central-1.amazonaws.com/avg-projects/monosdf/data/Heritage-Recon.tar", + "neural-rgbd-data": "http://kaldir.vc.in.tum.de/neural_rgbd/neural_rgbd_data.zip", + "all": None, +} + +SDFstudioCaptureName = tyro.extras.literal_type_from_choices(sdfstudio_downloads.keys()) + + +@dataclass +class SDFstudioDemoDownload(DatasetDownload): + """Download the sdfstudio dataset.""" + + dataset_name: SDFstudioCaptureName = "sdfstudio-demo-data" + + def download(self, save_dir: Path): + """Download the D-NeRF dataset (https://github.com/albertpumarola/D-NeRF).""" + # TODO: give this code the same structure as download_nerfstudio + + if self.dataset_name == "all": + for dataset_name in sdfstudio_downloads: + if dataset_name != "all": + SDFstudioDemoDownload(dataset_name=dataset_name).download(save_dir) + return + + assert ( + self.dataset_name in sdfstudio_downloads + ), f"Capture name {self.dataset_name} not found in {sdfstudio_downloads.keys()}" + + url = sdfstudio_downloads[self.dataset_name] + + target_path = str(save_dir / self.dataset_name) + os.makedirs(target_path, exist_ok=True) + + file_format = url[-4:] + + download_path = Path(f"{target_path}{file_format}") + tmp_path = str(save_dir / ".temp") + shutil.rmtree(tmp_path, ignore_errors=True) + os.makedirs(tmp_path, exist_ok=True) + + os.system(f"curl -L {url} > {download_path}") + if file_format == ".tar": + with tarfile.open(download_path, "r") as tar_ref: + tar_ref.extractall(str(tmp_path)) + elif file_format == ".zip": + with zipfile.ZipFile(download_path, "r") as zip_ref: + zip_ref.extractall(str(target_path)) + return + else: + raise NotImplementedError + + inner_folders = os.listdir(tmp_path) + assert len(inner_folders) == 1, "There is more than one folder inside this zip file." + folder = os.path.join(tmp_path, inner_folders[0]) + shutil.rmtree(target_path) + shutil.move(folder, target_path) + shutil.rmtree(tmp_path) + os.remove(download_path) + + Commands = Union[ Annotated[BlenderDownload, tyro.conf.subcommand(name="blender")], Annotated[FriendsDownload, tyro.conf.subcommand(name="friends")], @@ -294,6 +363,7 @@ Commands = Union[ Annotated[Record3dDownload, tyro.conf.subcommand(name="record3d")], Annotated[DNerfDownload, tyro.conf.subcommand(name="dnerf")], Annotated[PhototourismDownload, tyro.conf.subcommand(name="phototourism")], + Annotated[SDFstudioDemoDownload, tyro.conf.subcommand(name="sdfstudio")], ] diff --git a/scripts/process_data.py b/scripts/process_data.py index f6f9a9a7f..752e1d15c 100755 --- a/scripts/process_data.py +++ b/scripts/process_data.py @@ -737,7 +737,7 @@ class ProcessRealityCapture: """Process RealityCapture data into a nerfstudio dataset. This script assumes that cameras have been aligned using RealityCapture. After alignment, it is necessary to - export the camera poses as a `.csv` file. + export the camera poses as a `.csv` file using the `Internal/External camera parameters` option. This script does the following: diff --git a/scripts/viewer/__init__.py b/scripts/viewer/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/scripts/viewer/run_viewer.py b/scripts/viewer/run_viewer.py new file mode 100644 index 000000000..9a96c6703 --- /dev/null +++ b/scripts/viewer/run_viewer.py @@ -0,0 +1,115 @@ +#!/usr/bin/env python +""" +Starts viewer in eval mode. +""" +from __future__ import annotations + +import time +from dataclasses import dataclass, field, fields +from pathlib import Path + +import tyro +from rich.console import Console + +from nerfstudio.configs.base_config import ViewerConfig +from nerfstudio.engine.trainer import TrainerConfig +from nerfstudio.pipelines.base_pipeline import Pipeline +from nerfstudio.utils import writer +from nerfstudio.utils.eval_utils import eval_setup +from nerfstudio.utils.writer import EventName, TimeWriter +from nerfstudio.viewer.server import viewer_utils + +CONSOLE = Console(width=120) + + +@dataclass +class ViewerConfigWithoutNumRays(ViewerConfig): + """Configuration for viewer instantiation""" + + num_rays_per_chunk: tyro.conf.Suppress[int] = -1 + start_train: tyro.conf.Suppress[bool] = False + + def as_viewer_config(self): + """Converts the instance to ViewerConfig""" + return ViewerConfig(**{x.name: getattr(self, x.name) for x in fields(self)}) + + +@dataclass +class RunViewer: + """Load a checkpoint and start the viewer.""" + + load_config: Path + """Path to config YAML file.""" + viewer: ViewerConfigWithoutNumRays = field(default_factory=ViewerConfigWithoutNumRays) + """Viewer configuration""" + + def main(self) -> None: + """Main function.""" + config, pipeline, _ = eval_setup( + self.load_config, + eval_num_rays_per_chunk=None, + test_mode="test", + ) + num_rays_per_chunk = config.viewer.num_rays_per_chunk + assert self.viewer.num_rays_per_chunk == -1 + config.vis = "viewer" + config.viewer = self.viewer.as_viewer_config() + config.viewer.num_rays_per_chunk = num_rays_per_chunk + + self._start_viewer(config, pipeline) + + def _start_viewer(self, config: TrainerConfig, pipeline): + base_dir = config.get_base_dir() + viewer_log_path = base_dir / config.viewer.relative_log_filename + viewer_state, banner_messages = viewer_utils.setup_viewer( + config.viewer, log_filename=viewer_log_path, datapath=config.pipeline.datamanager.dataparser.data + ) + + # We don't need logging, but writer.GLOBAL_BUFFER needs to be populated + config.logging.local_writer.enable = False + writer.setup_local_writer(config.logging, max_iter=config.max_num_iterations, banner_messages=banner_messages) + + assert viewer_state and pipeline.datamanager.train_dataset + viewer_state.init_scene( + dataset=pipeline.datamanager.train_dataset, + start_train=False, + ) + while True: + viewer_state.vis["renderingState/isTraining"].write(False) + self._update_viewer_state(viewer_state, config, pipeline) + + def _update_viewer_state(self, viewer_state: viewer_utils.ViewerState, config: TrainerConfig, pipeline: Pipeline): + """Updates the viewer state by rendering out scene with current pipeline + Returns the time taken to render scene. + + """ + # NOTE: step must be > 0 otherwise the rendering would not happen + step = 1 + num_rays_per_batch = config.pipeline.datamanager.train_num_rays_per_batch + with TimeWriter(writer, EventName.ITER_VIS_TIME) as _: + try: + viewer_state.update_scene(self, step, pipeline.model, num_rays_per_batch) + except RuntimeError: + time.sleep(0.03) # sleep to allow buffer to reset + assert viewer_state.vis is not None + viewer_state.vis["renderingState/log_errors"].write( + "Error: GPU out of memory. Reduce resolution to prevent viewer from crashing." + ) + + def save_checkpoint(self, *args, **kwargs): + """ + Mock method because we pass this instance to viewer_state.update_scene + """ + + +def entrypoint(): + """Entrypoint for use with pyproject scripts.""" + tyro.extras.set_accent_color("bright_yellow") + tyro.cli(RunViewer).main() + + +if __name__ == "__main__": + entrypoint() + +# For sphinx docs +get_parser_fn = lambda: tyro.extras.get_parser(RunViewer) # noqa diff --git a/scripts/viewer/view_dataset.py b/scripts/viewer/view_dataset.py index 1f5aa2201..4bd526552 100644 --- a/scripts/viewer/view_dataset.py +++ b/scripts/viewer/view_dataset.py @@ -3,22 +3,21 @@ view_dataset.py """ -from rich.console import Console - -CONSOLE = Console(width=120) import time from datetime import timedelta from pathlib import Path import torch import tyro +from rich.console import Console from nerfstudio.configs.base_config import ViewerConfig -from nerfstudio.data.datamanagers import AnnotatedDataParserUnion +from nerfstudio.data.datamanagers.base_datamanager import AnnotatedDataParserUnion from nerfstudio.data.datasets.base_dataset import InputDataset from nerfstudio.viewer.server import viewer_utils DEFAULT_TIMEOUT = timedelta(minutes=30) +CONSOLE = Console(width=120) # speedup for when input size to model doesn't change (much) torch.backends.cudnn.benchmark = True # type: ignore @@ -30,9 +29,10 @@ def main( log_base_dir: Path = Path("/tmp/nerfstudio_viewer_logs"), ) -> None: """Main function.""" - viewer_state = viewer_utils.ViewerState( + viewer_state, _ = viewer_utils.setup_viewer( viewer, log_filename=log_base_dir / viewer.relative_log_filename, + datapath=dataparser.data, ) dataset = InputDataset(dataparser.setup().get_dataparser_outputs(split="train")) viewer_state.init_scene(dataset=dataset, start_train=False)