Merge branch 'nerfstudio-project:main' into main

nerfstudio-project · Mar 12, 2024 · 1dde17b · 1dde17b
2 parents 8878834 + 8e0c687
commit 1dde17b
Show file tree

Hide file tree

Showing 31 changed files with 600 additions and 316 deletions.
diff --git a/README.md b/README.md
@@ -1,6 +1,6 @@
 <p align="center">
     <!-- community badges -->
-    <a href="https://discord.gg/uMbNqcraFc"><img src="https://img.shields.io/badge/Join-Discord-blue.svg"/></a>
+    <a href="https://discord.gg/uMbNqcraFc"><img src="https://dcbadge.vercel.app/api/server/uMbNqcraFc?style=plastic"/></a>
     <!-- doc badges -->
     <a href='https://docs.nerf.studio/'>
         <img src='https://readthedocs.com/projects/plenoptix-nerfstudio/badge/?version=latest' alt='Documentation Status' /></a>

diff --git a/colab/demo.ipynb b/colab/demo.ipynb
@@ -74,7 +74,7 @@
         "\n",
         "%cd /content/\n",
         "!pip install --upgrade pip\n",
-        "!pip install torch==2.0.1+cu118 torchvision==0.15.2+cu118 --extra-index-url https://download.pytorch.org/whl/cu118\n",
+        "!pip install torch==2.0.1+cu118 torchvision==0.15.2+cu118 torchaudio==2.0.2+cu118 --extra-index-url https://download.pytorch.org/whl/cu118\n",
         "\n",
         "# Installing TinyCuda\n",
         "%cd /content/\n",
@@ -166,84 +166,6 @@
         "print(\"Data Processing Succeeded!\")"
       ]
     },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "cellView": "form",
-        "colab": {
-          "base_uri": "https://localhost:8080/",
-          "height": 928
-        },
-        "id": "VoKDxqEcjmfC",
-        "outputId": "d2919aa4-96dd-4e50-829f-289e4208882b"
-      },
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "/content\n",
-            "\u001b[K\u001b[?25h/tools/node/bin/lt -> /tools/node/lib/node_modules/localtunnel/bin/lt.js\n",
-            "\u001b[K\u001b[?25h+ [email protected]\n",
-            "added 22 packages from 22 contributors in 2.07s\n",
-            "https://viewer.nerf.studio/?websocket_url=wss://cyan-facts-matter-34-91-1-218.loca.lt\n",
-            "You may need to click Refresh Page after you start training!\n"
-          ]
-        },
-        {
-          "data": {
-            "text/html": [
-              "\n",
-              "        <iframe\n",
-              "            width=\"100%\"\n",
-              "            height=\"800\"\n",
-              "            src=\"https://viewer.nerf.studio/?websocket_url=wss://cyan-facts-matter-34-91-1-218.loca.lt\"\n",
-              "            frameborder=\"0\"\n",
-              "            allowfullscreen\n",
-              "        ></iframe>\n",
-              "        "
-            ],
-            "text/plain": [
-              "<IPython.lib.display.IFrame at 0x7f1d0da6f950>"
-            ]
-          },
-          "execution_count": 2,
-          "metadata": {},
-          "output_type": "execute_result"
-        }
-      ],
-      "source": [
-        "#@markdown <h1>Set up and Start Viewer</h1>\n",
-        "\n",
-        "%cd /content\n",
-        "\n",
-        "# Install localtunnel\n",
-        "# We are using localtunnel https://github.com/localtunnel/localtunnel but ngrok could also be used\n",
-        "!npm install -g localtunnel\n",
-        "\n",
-        "# Tunnel port 7007, the default for\n",
-        "!rm url.txt 2> /dev/null\n",
-        "get_ipython().system_raw('lt --port 7007 >> url.txt 2>&1 &')\n",
-        "\n",
-        "import time\n",
-        "time.sleep(3) # the previous command needs time to write to url.txt\n",
-        "\n",
-        "\n",
-        "with open('url.txt') as f:\n",
-        "  lines = f.readlines()\n",
-        "websocket_url = lines[0].split(\": \")[1].strip().replace(\"https\", \"wss\")\n",
-        "# from nerfstudio.utils.io import load_from_json\n",
-        "# from pathlib import Path\n",
-        "# json_filename = \"nerfstudio/nerfstudio/viewer/app/package.json\"\n",
-        "# version = load_from_json(Path(json_filename))[\"version\"]\n",
-        "url = f\"https://viewer.nerf.studio/?websocket_url={websocket_url}\"\n",
-        "print(url)\n",
-        "print(\"You may need to click Refresh Page after you start training!\")\n",
-        "from IPython import display\n",
-        "display.IFrame(src=url, height=800, width=\"100%\")"
-      ]
-    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -256,8 +178,16 @@
         "#@markdown <h1>Start Training</h1>\n",
         "\n",
         "%cd /content\n",
+        "!pip install colab-xterm\n",
+        "%load_ext colabxterm\n",
+        "%env TERM=xterm\n",
+        "from IPython.display import clear_output\n",
+        "clear_output(wait=True)\n",
         "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",
+        "    print(\"\\033[1m\" + \"Copy and paste the following command into the terminal window that pops up under this cell.\" + \"\\033[0m\")\n",
+        "    print(f\"ns-train nerfacto --viewer.websocket-port 7007 --viewer.make-share-url True nerfstudio-data --data data/nerfstudio/{scene} --downscale-factor 4\")\n",
+        "    print()\n",
+        "    %xterm\n",
         "else:\n",
         "    from IPython.core.display import display, HTML\n",
         "    display(HTML('<h3 style=\"color:red\">Error: Data processing did not complete</h3>'))\n",

diff --git a/docs/developer_guides/viewer/local_viewer.md b/docs/developer_guides/viewer/local_viewer.md
@@ -1,4 +1,8 @@
-# Local Server
+# (Legacy Viewer) Local Server
+
+**Note:** this doc only applies to the legacy version of the viewer, which was the default in in Nerfstudio versions `<=0.3.4`. It was deprecated starting Nerfstudio version `1.0.0`, where it needs to be opted into via the `--vis viewer_legacy` argument.
+
+---
 
 If you are unable to connect to `https://viewer.nerf.studio`, want to use Safari, or want to develop the viewer codebase, you can launch your own local viewer.
 

diff --git a/docs/index.md b/docs/index.md
@@ -156,6 +156,7 @@ This documentation is organized into 3 parts:
 - [Instruct-GS2GS](nerfology/methods/igs2gs.md): Editing 3DGS Scenes with Instructions
 - [PyNeRF](nerfology/methods/pynerf.md): Pyramidal Neural Radiance Fields
 - [SeaThru-NeRF](nerfology/methods/seathru_nerf.md): Neural Radiance Field for subsea scenes
+- [Zip-NeRF](nerfology/methods/zipnerf.md): Anti-Aliased Grid-Based Neural Radiance Fields
 
 **Eager to contribute a method?** We'd love to see you use nerfstudio in implementing new (or even existing) methods! Please view our {ref}`guide<own_method_docs>` for more details about how to add to this list!
 

diff --git a/docs/nerfology/methods/index.md b/docs/nerfology/methods/index.md
@@ -42,6 +42,7 @@ The following methods are supported in nerfstudio:
     Instruct-GS2GS<igs2gs.md>
     PyNeRF<pynerf.md>
     SeaThru-NeRF<seathru_nerf.md>
+    Zip-NeRF<zipnerf.md>
 ```
 
 (own_method_docs)=

diff --git a/docs/nerfology/methods/nerf.md b/docs/nerfology/methods/nerf.md
@@ -115,7 +115,7 @@ Rending RGB images is not the only type of output render supported. It is possib
 Associated nerfstudio code:
 
 ```python
-from nerfstudio.renderers.renderers import RGBRenderer
+from nerfstudio.model_components.renderers import RGBRenderer
 
 renderer_rgb = RGBRenderer(background_color=colors.WHITE)
 # Ray samples discussed in the next section

diff --git a/docs/nerfology/methods/splat.md b/docs/nerfology/methods/splat.md
@@ -31,6 +31,16 @@ Because the method trains on *full images* instead of bundles of rays, there is
 ### Running the Method
 To run splatfacto, run `ns-train splatfacto --data <data>`. Just like NeRF methods, the splat can be interactively viewed in the web-viewer, loaded from a checkpoint, rendered, and exported.
 
+We provide a few additional variants:
+
+| Method           | Description                    | Memory | Speed   |
+| ---------------- | ------------------------------ | ------ | ------- |
+| `splatfacto`     | Default Model                  | ~6GB   | Fast    |
+| `splatfacto-big` | More Gaussians, Higher Quality | ~12GB  | Slower  |
+
+
+A full evalaution of Nerfstudio's implementation of Gaussian Splatting against the original Inria method can be found [here](https://docs.gsplat.studio/tests/eval.html).
+
 #### Quality and Regularization
 The default settings provided maintain a balance between speed, quality, and splat file size, but if you care more about quality than training speed or size, you can decrease the alpha cull threshold 
 (threshold to delete translucent gaussians) and disable culling after 15k steps like so: `ns-train splatfacto --pipeline.model.cull_alpha_thresh=0.005 --pipeline.model.continue_cull_post_densification=False --data <data>`

diff --git a/docs/nerfology/methods/zipnerf.md b/docs/nerfology/methods/zipnerf.md
@@ -0,0 +1,33 @@
+# Zip-NeRF
+
+<h4>A pytorch implementation of "Zip-NeRF: Anti-Aliased Grid-Based Neural Radiance Fields"</h4>
+
+```{button-link} https://jonbarron.info/zipnerf/
+:color: primary
+:outline:
+Paper Website
+```
+```{button-link} https://github.com/SuLvXiangXin/zipnerf-pytorch
+:color: primary
+:outline:
+Code
+```
+### Installation
+First, install nerfstudio and its dependencies. Then run:
+```
+pip install git+https://github.com/SuLvXiangXin/zipnerf-pytorch#subdirectory=extensions/cuda
+pip install git+https://github.com/SuLvXiangXin/zipnerf-pytorch
+```
+Finally, install torch_scatter corresponding to your cuda version(https://pytorch-geometric.com/whl/torch-2.0.1%2Bcu118.html).
+
+
+### Running Model
+
+```bash
+ns-train zipnerf --data {DATA_DIR/SCENE}
+```
+
+## Overview
+Zipnerf combines mip-NeRF 360’s overall framework with iNGP’s featurization approach.
+Following mip-NeRF, zipnerf assume each pixel corresponds to a cone. Given an interval along the ray, it construct a set of multisamples that approximate the shape of that conical frustum.
+Also,  it present an alternative loss that, unlike mip-NeRF 360’s interlevel loss, is continuous and smooth with respect to distance along the ray to prevent z-aliasing.
diff --git a/nerfstudio/cameras/camera_optimizers.py b/nerfstudio/cameras/camera_optimizers.py
@@ -157,7 +157,7 @@ def apply_to_camera(self, camera: Cameras) -> None:
             assert camera.metadata is not None, "Must provide id of camera in its metadata"
             assert "cam_idx" in camera.metadata, "Must provide id of camera in its metadata"
             camera_idx = camera.metadata["cam_idx"]
-            adj = self([camera_idx])  # type: ignore
+            adj = self(torch.tensor([camera_idx], dtype=torch.long, device=camera.device))  # type: ignore
             adj = torch.cat([adj, torch.Tensor([0, 0, 0, 1])[None, None].to(adj)], dim=1)
             camera.camera_to_worlds = torch.bmm(camera.camera_to_worlds, adj)
 

diff --git a/nerfstudio/configs/external_methods.py b/nerfstudio/configs/external_methods.py
@@ -192,6 +192,22 @@ class ExternalMethod:
     )
 )
 
+# Zip-NeRF
+external_methods.append(
+    ExternalMethod(
+        """[bold yellow]Zip-NeRF[/bold yellow]
+For more information visit https://docs.nerf.studio/nerfology/methods/zipnerf.html
+
+To enable Zip-NeRF, you must install it first by running:
+  [grey]pip install git+https://github.com/SuLvXiangXin/zipnerf-pytorch#subdirectory=extensions/cuda 
+  and pip install git+https://github.com/SuLvXiangXin/zipnerf-pytorch[/grey]""",
+        configurations=[
+            ("zipnerf", "A pytorch implementation of 'Zip-NeRF: Anti-Aliased Grid-Based Neural Radiance Fields'")
+        ],
+        pip_package="pip install git+https://github.com/SuLvXiangXin/zipnerf-pytorch",
+    )
+)
+
 
 @dataclass
 class ExternalMethodDummyTrainerConfig:

diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py
@@ -599,7 +599,6 @@
     steps_per_eval_all_images=1000,
     max_num_iterations=30000,
     mixed_precision=False,
-    gradient_accumulation_steps={"camera_opt": 100},
     pipeline=VanillaPipelineConfig(
         datamanager=FullImageDatamanagerConfig(
             dataparser=NerfstudioDataParserConfig(load_3D_points=True),
@@ -608,7 +607,7 @@
         model=SplatfactoModelConfig(),
     ),
     optimizers={
-        "xyz": {
+        "means": {
             "optimizer": AdamOptimizerConfig(lr=1.6e-4, eps=1e-15),
             "scheduler": ExponentialDecaySchedulerConfig(
                 lr_final=1.6e-6,
@@ -623,15 +622,67 @@
             "optimizer": AdamOptimizerConfig(lr=0.0025 / 20, eps=1e-15),
             "scheduler": None,
         },
-        "opacity": {
+        "opacities": {
             "optimizer": AdamOptimizerConfig(lr=0.05, eps=1e-15),
             "scheduler": None,
         },
-        "scaling": {
+        "scales": {
             "optimizer": AdamOptimizerConfig(lr=0.005, eps=1e-15),
             "scheduler": None,
         },
-        "rotation": {"optimizer": AdamOptimizerConfig(lr=0.001, eps=1e-15), "scheduler": None},
+        "quats": {"optimizer": AdamOptimizerConfig(lr=0.001, eps=1e-15), "scheduler": None},
+        "camera_opt": {
+            "optimizer": AdamOptimizerConfig(lr=1e-3, eps=1e-15),
+            "scheduler": ExponentialDecaySchedulerConfig(lr_final=5e-5, max_steps=30000),
+        },
+    },
+    viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
+    vis="viewer",
+)
+
+method_configs["splatfacto-big"] = TrainerConfig(
+    method_name="splatfacto",
+    steps_per_eval_image=100,
+    steps_per_eval_batch=0,
+    steps_per_save=2000,
+    steps_per_eval_all_images=1000,
+    max_num_iterations=30000,
+    mixed_precision=False,
+    pipeline=VanillaPipelineConfig(
+        datamanager=FullImageDatamanagerConfig(
+            dataparser=NerfstudioDataParserConfig(load_3D_points=True),
+            cache_images_type="uint8",
+        ),
+        model=SplatfactoModelConfig(
+            cull_alpha_thresh=0.005,
+            continue_cull_post_densification=False,
+        ),
+    ),
+    optimizers={
+        "means": {
+            "optimizer": AdamOptimizerConfig(lr=1.6e-4, eps=1e-15),
+            "scheduler": ExponentialDecaySchedulerConfig(
+                lr_final=1.6e-6,
+                max_steps=30000,
+            ),
+        },
+        "features_dc": {
+            "optimizer": AdamOptimizerConfig(lr=0.0025, eps=1e-15),
+            "scheduler": None,
+        },
+        "features_rest": {
+            "optimizer": AdamOptimizerConfig(lr=0.0025 / 20, eps=1e-15),
+            "scheduler": None,
+        },
+        "opacities": {
+            "optimizer": AdamOptimizerConfig(lr=0.05, eps=1e-15),
+            "scheduler": None,
+        },
+        "scales": {
+            "optimizer": AdamOptimizerConfig(lr=0.005, eps=1e-15),
+            "scheduler": None,
+        },
+        "quats": {"optimizer": AdamOptimizerConfig(lr=0.001, eps=1e-15), "scheduler": None},
         "camera_opt": {
             "optimizer": AdamOptimizerConfig(lr=1e-3, eps=1e-15),
             "scheduler": ExponentialDecaySchedulerConfig(lr_final=5e-5, max_steps=30000),

diff --git a/nerfstudio/data/datamanagers/full_images_datamanager.py b/nerfstudio/data/datamanagers/full_images_datamanager.py
@@ -180,26 +180,30 @@ def process_eval_data(idx):
             return data
 
         CONSOLE.log("Caching / undistorting train images")
-        with ThreadPoolExecutor() as executor:
+        with ThreadPoolExecutor(max_workers=2) as executor:
             cached_train = list(
-                executor.map(
-                    process_train_data,
-                    track(
+                track(
+                    executor.map(
+                        process_train_data,
                         range(len(self.train_dataset)),
-                        description="Caching / undistorting train images",
-                        transient=True,
                     ),
+                    description="Caching / undistorting train images",
+                    transient=True,
+                    total=len(self.train_dataset),
                 )
             )
 
         CONSOLE.log("Caching / undistorting eval images")
-        with ThreadPoolExecutor() as executor:
+        with ThreadPoolExecutor(max_workers=2) as executor:
             cached_eval = list(
-                executor.map(
-                    process_eval_data,
-                    track(
-                        range(len(self.eval_dataset)), description="Caching / undistorting eval images", transient=True
+                track(
+                    executor.map(
+                        process_eval_data,
+                        range(len(self.eval_dataset)),
                     ),
+                    description="Caching / undistorting eval images",
+                    transient=True,
+                    total=len(self.eval_dataset),
                 )
             )
 
@@ -383,6 +387,8 @@ def _undistort_image(
                 mask = cv2.undistort(mask, K, distortion_params, None, newK)  # type: ignore
             mask = mask[y : y + h, x : x + w]
             mask = torch.from_numpy(mask).bool()
+            if len(mask.shape) == 2:
+                mask = mask[:, :, None]
         K = newK
 
     elif camera.camera_type.item() == CameraType.FISHEYE.value:
@@ -402,6 +408,8 @@ def _undistort_image(
             mask = mask.astype(np.uint8) * 255
             mask = cv2.fisheye.undistortImage(mask, K, distortion_params, None, newK)
             mask = torch.from_numpy(mask).bool()
+            if len(mask.shape) == 2:
+                mask = mask[:, :, None]
         K = newK
     elif camera.camera_type.item() == CameraType.FISHEYE624.value:
         fisheye624_params = torch.cat(
@@ -494,6 +502,8 @@ def _undistort_image(
             )
             / 255.0
         ).bool()[..., None]
+        if len(mask.shape) == 2:
+            mask = mask[:, :, None]
         assert mask.shape == (undist_h, undist_w, 1)
         K = undist_K.numpy()
     else: