diff --git a/.devcontainer/docker-compose.yml b/.devcontainer/docker-compose.yml
index de862fc1a..5cbf79cc3 100644
--- a/.devcontainer/docker-compose.yml
+++ b/.devcontainer/docker-compose.yml
@@ -3,10 +3,10 @@ version: "3.8"
services:
nerfstudio:
shm_size: '12gb'
- image: dromni/nerfstudio:0.1.18
- # build:
- # context: ..
- # dockerfile: ./Dockerfile
+ image: nerfstudio:latest
+ build:
+ context: ..
+ dockerfile: ./Dockerfile
stdin_open: true
tty: true
environment:
diff --git a/.vscode/launch.json b/.vscode/launch.json
index 935d7ae8b..118282928 100644
--- a/.vscode/launch.json
+++ b/.vscode/launch.json
@@ -92,6 +92,25 @@
"program": "scripts/train.py",
"console": "integratedTerminal",
"args": ["semantic_nerf"]
- }
+ },
+ {
+ "name": "Python: NeuS on Replica",
+ "type": "python",
+ "request": "launch",
+ "program": "scripts/train.py",
+ "console": "integratedTerminal",
+ "justMyCode": false,
+ "env": {"CUDA_VISIBLE_DEVICES":"0"},
+ "args": [
+ "neus",
+ "--pipeline.model.sdf-field.inside-outside=True",
+ "--pipeline.model.mono-depth-loss-mult=0.1",
+ "--pipeline.model.mono-normal-loss-mult=0.05",
+ "--vis=wandb",
+ "sdfstudio-data",
+ "--data=data/sdfstudio-demo-data/replica-room0",
+ "--include_mono_prior=True"
+ ]
+ },
]
}
diff --git a/Dockerfile b/Dockerfile
index 2bdb5328b..92af70932 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -34,6 +34,7 @@ RUN apt-get update && \
libboost-program-options-dev \
libboost-system-dev \
libboost-test-dev \
+ libhdf5-dev \
libcgal-dev \
libeigen3-dev \
libflann-dev \
@@ -132,6 +133,19 @@ RUN git clone --branch master --recursive https://github.com/cvg/Hierarchical-Lo
python3.10 -m pip install -e . && \
cd ..
+# Install pyceres from source
+RUN git clone --branch main --recursive https://github.com/cvg/pyceres.git && \
+ cd pyceres && \
+ python3.10 -m pip install -e . && \
+ cd ..
+
+# Install pixel perfect sfm.
+RUN git clone --branch main --recursive https://github.com/cvg/pixel-perfect-sfm && \
+ cd pixel-perfect-sfm && \
+ python3.10 -m pip install -e . && \
+ cd ..
+
+RUN python3.10 -m pip install omegaconf
# Copy nerfstudio folder and give ownership to user.
ADD . /home/user/nerfstudio
USER root
diff --git a/README.md b/README.md
index 0b0d8592c..8730cb20a 100644
--- a/README.md
+++ b/README.md
@@ -200,6 +200,7 @@ Using an existing dataset is great, but likely you want to use your own data! We
| ------------------------------------------------------------------------------------------------- | -------------- | ----------------------------------------------- | ----------------------- |
| 📷 [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) | 🐢 |
+| 🌎 [360 Data](https://docs.nerf.studio/en/latest/quickstart/custom_dataset.html#360_data) | 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/) | 🐇 |
diff --git a/colab/demo.ipynb b/colab/demo.ipynb
index f88f3f494..4a87f1718 100644
--- a/colab/demo.ipynb
+++ b/colab/demo.ipynb
@@ -7,7 +7,7 @@
"colab_type": "text"
},
"source": [
- "![\"Open](\"https://colab.research.google.com/assets/colab-badge.svg\")
"
+ ""
]
},
{
@@ -65,43 +65,12 @@
"cell_type": "code",
"execution_count": null,
"metadata": {
- "colab": {
- "base_uri": "https://localhost:8080/"
- },
- "id": "RGr33zHaHak0",
- "outputId": "6ad38334-22d2-4675-bffb-0bae91b17577"
- },
- "outputs": [
- {
- "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:18\n",
- "🔁 Restarting kernel...\n"
- ]
- }
- ],
- "source": [
- "#@markdown Install Conda (requires runtime restart)
\n",
- "\n",
- "!pip install -q condacolab\n",
- "import condacolab\n",
- "condacolab.install()"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "9oyLHl8QfYwP"
+ "id": "9oyLHl8QfYwP",
+ "cellView": "form"
},
"outputs": [],
"source": [
- "#@markdown Install Nerfstudio and Dependencies (~15 min)
\n",
+ "#@markdown Install Nerfstudio and Dependencies (~8 min)
\n",
"\n",
"%cd /content/\n",
"!pip install --upgrade pip\n",
@@ -109,18 +78,37 @@
"\n",
"# Installing TinyCuda\n",
"%cd /content/\n",
- "!gdown \"https://drive.google.com/u/1/uc?id=12RL_NVgE9WGvr_fEsXEiuaJ1QESvQCPl&confirm=t\" \n",
- "!pip install tinycudann-1.7-cp38-cp38-linux_x86_64.whl\n",
+ "!gdown \"https://drive.google.com/u/1/uc?id=1LlL2ofrViALsqE789MFGZeJyxkNu84cT&confirm=t\" \n",
+ "!pip install tinycudann-1.7-cp39-cp39-linux_x86_64.whl\n",
"\n",
"# Installing COLMAP\n",
"%cd /content/\n",
- "!conda install -c conda-forge colmap\n",
+ "!gdown \"https://drive.google.com/u/0/uc?id=15WngFRNar_b8CaPR5R-hvQ3eAnlyk_SL&confirm=t\"\n",
+ "!sudo apt-get install \\\n",
+ " build-essential \\\n",
+ " libboost-program-options-dev \\\n",
+ " libboost-filesystem-dev \\\n",
+ " libboost-graph-dev \\\n",
+ " libboost-system-dev \\\n",
+ " libboost-test-dev \\\n",
+ " libeigen3-dev \\\n",
+ " libflann-dev \\\n",
+ " libfreeimage-dev \\\n",
+ " libmetis-dev \\\n",
+ " libgoogle-glog-dev \\\n",
+ " libgflags-dev \\\n",
+ " libsqlite3-dev \\\n",
+ " libglew-dev \\\n",
+ " qtbase5-dev \\\n",
+ " libqt5opengl5-dev \\\n",
+ " libcgal-dev \\\n",
+ " libceres-dev\n",
+ "!unzip local.zip -d /usr/\n",
+ "!chmod +x /usr/local/bin/colmap\n",
"\n",
"# Install nerfstudio\n",
"%cd /content/\n",
- "# !pip install nerfstudio\n",
- "!pip install git+https://github.com/nerfstudio-project/nerfstudio.git\n",
- "!conda remove --force qt-main"
+ "!pip install git+https://github.com/nerfstudio-project/nerfstudio.git"
]
},
{
diff --git a/docs/quickstart/custom_dataset.md b/docs/quickstart/custom_dataset.md
index 009cb287c..4295bf4bc 100644
--- a/docs/quickstart/custom_dataset.md
+++ b/docs/quickstart/custom_dataset.md
@@ -15,6 +15,7 @@ We Currently support the following custom data types:
| ----------------------------- | -------------- | ----------------------------------------------- | ----------------------- |
| 📷 [Images](images_and_video) | Any | [COLMAP](https://colmap.github.io/install.html) | 🐢 |
| 📹 [Video](images_and_video) | Any | [COLMAP](https://colmap.github.io/install.html) | 🐢 |
+| 🌎 [360 Data](360_data) | 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/) | 🐇 |
@@ -291,11 +292,14 @@ ns-process-data metashape --data {data directory} --xml {xml file} --output-dir
ns-train nerfacto --data {output directory}
```
-## Equirectangular data
+(360_data)=
+
+## 360 Data (Equirectangular)
Equirectangular data is data that has been taken by a 360 camera such as Insta360. Both equirectangular image sets and videos can be processed by nerfstudio.
### Images
+
For a set of equirectangular images, process the data using the following command:
```bash
@@ -305,20 +309,21 @@ ns-process-data images --camera-type equirectangular --images-per-equirect {8, o
The images-per-equirect argument is the number of images that will be sampled from each equirectangular image. We have found that 8 images per equirectangular image is sufficient for most use cases so it defaults to that. However, if you find that there isn't enough detail in the nerf or that colmap is having trouble aligning the images, you can try increasing the number of images per equirectangular image to 14.
### Videos
-For videos we reccomend taking a video with the camera held on top of your head. This will result in any unwanted capturer to just be in the bottom of each frame image and therefore can be cropped out.
+
+For videos we recommend taking a video with the camera held on top of your head. This will result in any unwanted capturer to just be in the bottom of each frame image and therefore can be cropped out.
For a video, process the data using the following command:
```bash
ns-process-data video --camera-type equirectangular --images-per-equirect {8, or 14} --num-frames-target {num equirectangular frames to sample from} --crop-factor {top bottom left right} --data {data directory} --output-dir {output directory}
```
+
See the equirectangular images section above for a description of the `--images-per-equirect` argument.
-The `num-frames-target` argument is optional but it is recomended to set it to 3*(seconds of video) frames. For example, if you have a 30 second video, you would use `--num-frames-target 90` (3*30=90). This number was chosen from a bit of experimentation and seems to work well for most videos. It is by no means a hard rule and you can experiment with different values.
+The `num-frames-target` argument is optional but it is recommended to set it to 3*(seconds of video) frames. For example, if you have a 30 second video, you would use `--num-frames-target 90` (3*30=90). This number was chosen from a bit of experimentation and seems to work well for most videos. It is by no means a hard rule and you can experiment with different values.
The `crop-factor` argument is optional but often very helpful. This is because equirectangular videos taken by 360 cameras tend to have a portion of the bottom of the image that is the person who was holding the camera over their head.
This obscene mesh of human is obviously not helpful in training a nerf so we can remove it by cropping the bottom 20% of the image. This can be done by using the `--crop-factor 0 0.2 0 0` argument.
-
diff --git a/nerfstudio/cameras/camera_utils.py b/nerfstudio/cameras/camera_utils.py
index d2da627b7..79490ae67 100644
--- a/nerfstudio/cameras/camera_utils.py
+++ b/nerfstudio/cameras/camera_utils.py
@@ -471,7 +471,7 @@ def auto_orient_and_center_poses(
poses: TensorType["num_poses":..., 4, 4],
method: Literal["pca", "up", "vertical", "none"] = "up",
center_method: Literal["poses", "focus", "none"] = "poses",
-) -> Tuple[TensorType["num_poses":..., 3, 4], TensorType[4, 4]]:
+) -> Tuple[TensorType["num_poses":..., 3, 4], TensorType[3, 4]]:
"""Orients and centers the poses. We provide two methods for orientation: pca and up.
pca: Orient the poses so that the principal directions of the camera centers are aligned
diff --git a/nerfstudio/cameras/rays.py b/nerfstudio/cameras/rays.py
index fd4e8129e..0184b4c80 100644
--- a/nerfstudio/cameras/rays.py
+++ b/nerfstudio/cameras/rays.py
@@ -17,7 +17,7 @@ Some ray datastructures.
"""
import random
from dataclasses import dataclass
-from typing import Callable, Dict, Optional, Union
+from typing import Callable, Dict, Optional, Tuple, Union
import torch
from torchtyping import TensorType
@@ -56,6 +56,14 @@ class Frustums(TensorDataclass):
pos = pos + self.offsets
return pos
+ def get_start_positions(self) -> TensorType[..., 3]:
+ """Calulates "start" position of frustum.
+
+ Returns:
+ xyz positions.
+ """
+ return self.origins + self.directions * self.starts
+
def set_offsets(self, offsets):
"""Sets offsets for this frustum for computing positions"""
self.offsets = offsets
@@ -140,6 +148,27 @@ class RaySamples(TensorDataclass):
return weights
+ @staticmethod
+ def get_weights_and_transmittance_from_alphas(
+ alphas: TensorType[..., "num_samples", 1], weights_only: bool = False
+ ) -> Tuple[TensorType[..., "num_samples", 1], TensorType[..., "num_samples", 1]]:
+ """Return weights based on predicted alphas
+ Args:
+ alphas: Predicted alphas (maybe from sdf) for samples along ray
+ weights_only: If function should return only weights
+ Returns:
+ Tuple of weights and transmittance for each sample
+ """
+
+ transmittance = torch.cumprod(
+ torch.cat([torch.ones((*alphas.shape[:1], 1, 1), device=alphas.device), 1.0 - alphas + 1e-7], 1), 1
+ )
+
+ weights = alphas * transmittance[:, :-1, :]
+ if weights_only:
+ return weights
+ return weights, transmittance
+
@dataclass
class RayBundle(TensorDataclass):
diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py
index 312b99585..679f23331 100644
--- a/nerfstudio/configs/method_configs.py
+++ b/nerfstudio/configs/method_configs.py
@@ -30,6 +30,7 @@ from nerfstudio.data.datamanagers.depth_datamanager import DepthDataManagerConfi
from nerfstudio.data.datamanagers.dreamfusion_datamanager import (
DreamFusionDataManagerConfig,
)
+from nerfstudio.data.datamanagers.sdf_datamanager import SDFDataManagerConfig
from nerfstudio.data.datamanagers.semantic_datamanager import SemanticDataManagerConfig
from nerfstudio.data.datamanagers.variable_res_datamanager import (
VariableResDataManagerConfig,
@@ -44,9 +45,13 @@ from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataPars
from nerfstudio.data.dataparsers.phototourism_dataparser import (
PhototourismDataParserConfig,
)
+from nerfstudio.data.dataparsers.sdfstudio_dataparser import SDFStudioDataParserConfig
from nerfstudio.data.dataparsers.sitcoms3d_dataparser import Sitcoms3DDataParserConfig
from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig
-from nerfstudio.engine.schedulers import ExponentialDecaySchedulerConfig
+from nerfstudio.engine.schedulers import (
+ CosineDecaySchedulerConfig,
+ ExponentialDecaySchedulerConfig,
+)
from nerfstudio.engine.trainer import TrainerConfig
from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind
from nerfstudio.models.depth_nerfacto import DepthNerfactoModelConfig
@@ -57,6 +62,7 @@ from nerfstudio.models.mipnerf import MipNerfModel
from nerfstudio.models.nerfacto import NerfactoModelConfig
from nerfstudio.models.nerfplayer_nerfacto import NerfplayerNerfactoModelConfig
from nerfstudio.models.nerfplayer_ngp import NerfplayerNGPModelConfig
+from nerfstudio.models.neus import NeuSModelConfig
from nerfstudio.models.semantic_nerfw import SemanticNerfWModelConfig
from nerfstudio.models.tensorf import TensoRFModelConfig
from nerfstudio.models.vanilla_nerf import NeRFModel, VanillaModelConfig
@@ -80,6 +86,7 @@ descriptions = {
"dreamfusion": "Generative Text to NeRF model",
"nerfplayer-nerfacto": "NeRFPlayer with nerfacto backbone.",
"nerfplayer-ngp": "NeRFPlayer with InstantNGP backbone.",
+ "neus": "Implementation of NeuS. (slow)",
}
method_configs["nerfacto"] = TrainerConfig(
@@ -529,6 +536,39 @@ method_configs["nerfplayer-ngp"] = TrainerConfig(
vis="viewer",
)
+method_configs["neus"] = TrainerConfig(
+ method_name="neus",
+ steps_per_eval_image=500,
+ steps_per_eval_batch=5000,
+ steps_per_save=20000,
+ steps_per_eval_all_images=1000000, # set to a very large number so we don't eval with all images
+ max_num_iterations=100000,
+ mixed_precision=False,
+ pipeline=VanillaPipelineConfig(
+ datamanager=SDFDataManagerConfig(
+ dataparser=SDFStudioDataParserConfig(),
+ train_num_rays_per_batch=1024,
+ eval_num_rays_per_batch=1024,
+ camera_optimizer=CameraOptimizerConfig(
+ mode="off", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
+ ),
+ ),
+ model=NeuSModelConfig(eval_num_rays_per_chunk=1024),
+ ),
+ optimizers={
+ "fields": {
+ "optimizer": AdamOptimizerConfig(lr=5e-4, eps=1e-15),
+ "scheduler": CosineDecaySchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=300000),
+ },
+ "field_background": {
+ "optimizer": AdamOptimizerConfig(lr=5e-4, eps=1e-15),
+ "scheduler": CosineDecaySchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=300000),
+ },
+ },
+ viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
+ vis="viewer",
+)
+
external_methods, external_descriptions = discover_methods()
method_configs.update(external_methods)
descriptions.update(external_descriptions)
diff --git a/nerfstudio/data/datamanagers/sdf_datamanager.py b/nerfstudio/data/datamanagers/sdf_datamanager.py
new file mode 100644
index 000000000..e1aca46df
--- /dev/null
+++ b/nerfstudio/data/datamanagers/sdf_datamanager.py
@@ -0,0 +1,48 @@
+# 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.
+
+"""
+SDF datamanager.
+"""
+
+from dataclasses import dataclass, field
+from typing import Type
+
+from nerfstudio.data.datamanagers import base_datamanager
+from nerfstudio.data.datasets.sdf_dataset import SDFDataset
+
+
+@dataclass
+class SDFDataManagerConfig(base_datamanager.VanillaDataManagerConfig):
+ """A sdf datamanager - required to use with .setup()"""
+
+ _target: Type = field(default_factory=lambda: SDFDataManager)
+
+
+class SDFDataManager(base_datamanager.VanillaDataManager): # pylint: disable=abstract-method
+ """Data manager implementation for data that also requires processing sdf data.
+ Args:
+ config: the DataManagerConfig used to instantiate class
+ """
+
+ def create_train_dataset(self) -> SDFDataset:
+ self.train_dataparser_outputs = self.dataparser.get_dataparser_outputs(split="train")
+ return SDFDataset(
+ dataparser_outputs=self.train_dataparser_outputs,
+ )
+
+ def create_eval_dataset(self) -> SDFDataset:
+ return SDFDataset(
+ dataparser_outputs=self.dataparser.get_dataparser_outputs(split=self.test_split),
+ )
diff --git a/nerfstudio/data/dataparsers/minimal_dataparser.py b/nerfstudio/data/dataparsers/minimal_dataparser.py
index 2fa2218a4..6d8de4ec5 100644
--- a/nerfstudio/data/dataparsers/minimal_dataparser.py
+++ b/nerfstudio/data/dataparsers/minimal_dataparser.py
@@ -52,15 +52,17 @@ class MinimalDataParser(DataParser):
filepath = self.config.data / f"{split}.npz"
data = np.load(filepath, allow_pickle=True)
- image_filenames = data["image_filenames"].tolist()
- mask_filenames = data["mask_filenames"].tolist() if "mask_filenames" in data.keys() else None
+ image_filenames = [filepath.parent / path for path in data["image_filenames"].tolist()]
+ mask_filenames = None
+ if "mask_filenames" in data.keys():
+ mask_filenames = [filepath.parent / path for path in data["mask_filenames"].tolist()]
metadata = None
if "semantics" in data.keys():
semantics = data["semantics"].item()
metadata = {
"semantics": Semantics(
- filenames=semantics["filenames"].tolist(),
+ filenames=[filepath.parent / path for path in semantics["filenames"].tolist()],
classes=semantics["classes"].tolist(),
colors=torch.from_numpy(semantics["colors"]),
mask_classes=semantics["mask_classes"].tolist(),
@@ -87,11 +89,21 @@ class MinimalDataParser(DataParser):
times=torch.from_numpy(camera_np["times"]) if "times" in camera_np.keys() else None,
)
+ applied_scale = 1.0
+ applied_transform = np.eye(4, dtype=np.float32)[:3, :]
+ if "applied_scale" in data.keys():
+ applied_scale = float(data["applied_scale"])
+ if "applied_transform" in data.keys():
+ applied_transform = data["applied_transform"].astype(np.float32)
+ assert applied_transform.shape == (3, 4)
+
dataparser_outputs = DataparserOutputs(
image_filenames=image_filenames,
cameras=cameras,
scene_box=scene_box,
mask_filenames=mask_filenames,
+ dataparser_transform=applied_transform,
+ dataparser_scale=applied_scale,
metadata=metadata,
)
return dataparser_outputs
diff --git a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py
index adb09eb48..4a74cd961 100644
--- a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py
+++ b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py
@@ -182,22 +182,36 @@ class Nerfstudio(DataParser):
You should check that depth_file_path is specified for every frame (or zero frames) in transforms.json.
"""
- # filter image_filenames and poses based on train/eval split percentage
- num_images = len(image_filenames)
- num_train_images = math.ceil(num_images * self.config.train_split_fraction)
- num_eval_images = num_images - num_train_images
- i_all = np.arange(num_images)
- i_train = np.linspace(
- 0, num_images - 1, num_train_images, dtype=int
- ) # equally spaced training images starting and ending at 0 and num_images-1
- i_eval = np.setdiff1d(i_all, i_train) # eval images are the remaining images
- assert len(i_eval) == num_eval_images
- if split == "train":
- indices = i_train
- elif split in ["val", "test"]:
- indices = i_eval
+ has_split_files_spec = any(f"{split}_filenames" in meta for split in ("train", "val", "test"))
+ if f"{split}_filenames" in meta:
+ # Validate split first
+ split_filenames = set(self._get_fname(PurePath(x), data_dir) for x in meta[f"{split}_filenames"])
+ unmatched_filenames = split_filenames.difference(image_filenames)
+ if unmatched_filenames:
+ raise RuntimeError(f"Some filenames for split {split} were not found: {unmatched_filenames}.")
+
+ indices = [i for i, path in enumerate(image_filenames) if path in split_filenames]
+ CONSOLE.log(f"[yellow] Dataset is overriding {split}_indices to {indices}")
+ indices = np.array(indices, dtype=np.int32)
+ elif has_split_files_spec:
+ raise RuntimeError(f"The dataset's list of filenames for split {split} is missing.")
else:
- raise ValueError(f"Unknown dataparser split {split}")
+ # filter image_filenames and poses based on train/eval split percentage
+ num_images = len(image_filenames)
+ num_train_images = math.ceil(num_images * self.config.train_split_fraction)
+ num_eval_images = num_images - num_train_images
+ i_all = np.arange(num_images)
+ i_train = np.linspace(
+ 0, num_images - 1, num_train_images, dtype=int
+ ) # equally spaced training images starting and ending at 0 and num_images-1
+ i_eval = np.setdiff1d(i_all, i_train) # eval images are the remaining images
+ assert len(i_eval) == num_eval_images
+ if split == "train":
+ indices = i_train
+ elif split in ["val", "test"]:
+ indices = i_eval
+ else:
+ raise ValueError(f"Unknown dataparser split {split}")
if "orientation_override" in meta:
orientation_method = meta["orientation_override"]
@@ -274,6 +288,15 @@ class Nerfstudio(DataParser):
assert self.downscale_factor is not None
cameras.rescale_output_resolution(scaling_factor=1.0 / self.downscale_factor)
+ if "applied_transform" in meta:
+ applied_transform = torch.tensor(meta["applied_transform"], dtype=transform_matrix.dtype)
+ transform_matrix = transform_matrix @ torch.cat(
+ [applied_transform, torch.tensor([[0, 0, 0, 1]], dtype=transform_matrix.dtype)], 0
+ )
+ if "applied_scale" in meta:
+ applied_scale = float(meta["applied_scale"])
+ scale_factor *= applied_scale
+
dataparser_outputs = DataparserOutputs(
image_filenames=image_filenames,
cameras=cameras,
diff --git a/nerfstudio/data/dataparsers/sdfstudio_dataparser.py b/nerfstudio/data/dataparsers/sdfstudio_dataparser.py
index 43781a00b..9b692c65b 100644
--- a/nerfstudio/data/dataparsers/sdfstudio_dataparser.py
+++ b/nerfstudio/data/dataparsers/sdfstudio_dataparser.py
@@ -107,6 +107,7 @@ class SDFStudio(DataParser):
fy = torch.stack(fy)
cx = torch.stack(cx)
cy = torch.stack(cy)
+ c2w_colmap = torch.stack(camera_to_worlds)
camera_to_worlds = torch.stack(camera_to_worlds)
# Convert from COLMAP's/OPENCV's camera coordinate system to nerfstudio
@@ -151,7 +152,8 @@ class SDFStudio(DataParser):
"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,
+ # required for normal maps, these are in colmap format so they require c2w before conversion
+ "camera_to_worlds": c2w_colmap if len(c2w_colmap) > 0 else None,
"include_mono_prior": self.config.include_mono_prior,
},
)
diff --git a/nerfstudio/data/datasets/sdf_dataset.py b/nerfstudio/data/datasets/sdf_dataset.py
index 48e0cff6d..db0bb9796 100644
--- a/nerfstudio/data/datasets/sdf_dataset.py
+++ b/nerfstudio/data/datasets/sdf_dataset.py
@@ -42,7 +42,7 @@ class SDFDataset(InputDataset):
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.transform = self.metadata["transform"]
self.include_mono_prior = self.metadata["include_mono_prior"]
def get_metadata(self, data: Dict) -> Dict:
@@ -57,8 +57,8 @@ class SDFDataset(InputDataset):
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
+ metadata["depth"] = depth_image
+ metadata["normal"] = normal_image
return metadata
diff --git a/nerfstudio/data/utils/colmap_parsing_utils.py b/nerfstudio/data/utils/colmap_parsing_utils.py
index 652e2795e..3ca853a27 100644
--- a/nerfstudio/data/utils/colmap_parsing_utils.py
+++ b/nerfstudio/data/utils/colmap_parsing_utils.py
@@ -4,9 +4,12 @@ This file copied with small modifications from:
TODO(1480) Delete this file when moving to pycolmap.
-"""
-# Copyright (c) 2018, ETH Zurich and UNC Chapel Hill.
+"""
+# pylint: disable-all
+# Disabling all pylinting since this file is copy-pasted over from COLMAP
+
+# Copyright (c) 2023, ETH Zurich and UNC Chapel Hill.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
@@ -35,11 +38,9 @@ TODO(1480) Delete this file when moving to pycolmap.
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
#
-# Author: Johannes L. Schoenberger (jsch at inf.ethz.ch)
-
-# pylint: disable-all
-# Disabling all pylinting since this file is copy-pasted over from COLMAP
+# Author: Johannes L. Schoenberger (jsch-at-demuc-dot-de)
+import argparse
import collections
import os
import struct
@@ -71,6 +72,7 @@ CAMERA_MODELS = {
CameraModel(model_id=10, model_name="THIN_PRISM_FISHEYE", num_params=12),
}
CAMERA_MODEL_IDS = dict([(camera_model.model_id, camera_model) for camera_model in CAMERA_MODELS])
+CAMERA_MODEL_NAMES = dict([(camera_model.model_name, camera_model) for camera_model in CAMERA_MODELS])
def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"):
@@ -85,6 +87,22 @@ def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"):
return struct.unpack(endian_character + format_char_sequence, data)
+def write_next_bytes(fid, data, format_char_sequence, endian_character="<"):
+ """pack and write to a binary file.
+ :param fid:
+ :param data: data to send, if multiple elements are sent at the same time,
+ they should be encapsuled either in a list or a tuple
+ :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}.
+ should be the same length as the data list or tuple
+ :param endian_character: Any of {@, =, <, >, !}
+ """
+ if isinstance(data, (list, tuple)):
+ bytes = struct.pack(endian_character + format_char_sequence, *data)
+ else:
+ bytes = struct.pack(endian_character + format_char_sequence, data)
+ fid.write(bytes)
+
+
def read_cameras_text(path):
"""
see: src/base/reconstruction.cc
@@ -118,7 +136,7 @@ def read_cameras_binary(path_to_model_file):
cameras = {}
with open(path_to_model_file, "rb") as fid:
num_cameras = read_next_bytes(fid, 8, "Q")[0]
- for camera_line_index in range(num_cameras):
+ for _ in range(num_cameras):
camera_properties = read_next_bytes(fid, num_bytes=24, format_char_sequence="iiQQ")
camera_id = camera_properties[0]
model_id = camera_properties[1]
@@ -134,6 +152,42 @@ def read_cameras_binary(path_to_model_file):
return cameras
+def write_cameras_text(cameras, path):
+ """
+ see: src/base/reconstruction.cc
+ void Reconstruction::WriteCamerasText(const std::string& path)
+ void Reconstruction::ReadCamerasText(const std::string& path)
+ """
+ HEADER = (
+ "# Camera list with one line of data per camera:\n"
+ + "# CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[]\n"
+ + "# Number of cameras: {}\n".format(len(cameras))
+ )
+ with open(path, "w") as fid:
+ fid.write(HEADER)
+ for _, cam in cameras.items():
+ to_write = [cam.id, cam.model, cam.width, cam.height, *cam.params]
+ line = " ".join([str(elem) for elem in to_write])
+ fid.write(line + "\n")
+
+
+def write_cameras_binary(cameras, path_to_model_file):
+ """
+ see: src/base/reconstruction.cc
+ void Reconstruction::WriteCamerasBinary(const std::string& path)
+ void Reconstruction::ReadCamerasBinary(const std::string& path)
+ """
+ with open(path_to_model_file, "wb") as fid:
+ write_next_bytes(fid, len(cameras), "Q")
+ for _, cam in cameras.items():
+ model_id = CAMERA_MODEL_NAMES[cam.model].model_id
+ camera_properties = [cam.id, model_id, cam.width, cam.height]
+ write_next_bytes(fid, camera_properties, "iiQQ")
+ for p in cam.params:
+ write_next_bytes(fid, float(p), "d")
+ return cameras
+
+
def read_images_text(path):
"""
see: src/base/reconstruction.cc
@@ -178,7 +232,7 @@ def read_images_binary(path_to_model_file):
images = {}
with open(path_to_model_file, "rb") as fid:
num_reg_images = read_next_bytes(fid, 8, "Q")[0]
- for image_index in range(num_reg_images):
+ for _ in range(num_reg_images):
binary_image_properties = read_next_bytes(fid, num_bytes=64, format_char_sequence="idddddddi")
image_id = binary_image_properties[0]
qvec = np.array(binary_image_properties[1:5])
@@ -205,6 +259,57 @@ def read_images_binary(path_to_model_file):
return images
+def write_images_text(images, path):
+ """
+ see: src/base/reconstruction.cc
+ void Reconstruction::ReadImagesText(const std::string& path)
+ void Reconstruction::WriteImagesText(const std::string& path)
+ """
+ if len(images) == 0:
+ mean_observations = 0
+ else:
+ mean_observations = sum((len(img.point3D_ids) for _, img in images.items())) / len(images)
+ HEADER = (
+ "# Image list with two lines of data per image:\n"
+ + "# IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME\n"
+ + "# POINTS2D[] as (X, Y, POINT3D_ID)\n"
+ + "# Number of images: {}, mean observations per image: {}\n".format(len(images), mean_observations)
+ )
+
+ with open(path, "w") as fid:
+ fid.write(HEADER)
+ for _, img in images.items():
+ image_header = [img.id, *img.qvec, *img.tvec, img.camera_id, img.name]
+ first_line = " ".join(map(str, image_header))
+ fid.write(first_line + "\n")
+
+ points_strings = []
+ for xy, point3D_id in zip(img.xys, img.point3D_ids):
+ points_strings.append(" ".join(map(str, [*xy, point3D_id])))
+ fid.write(" ".join(points_strings) + "\n")
+
+
+def write_images_binary(images, path_to_model_file):
+ """
+ see: src/base/reconstruction.cc
+ void Reconstruction::ReadImagesBinary(const std::string& path)
+ void Reconstruction::WriteImagesBinary(const std::string& path)
+ """
+ with open(path_to_model_file, "wb") as fid:
+ write_next_bytes(fid, len(images), "Q")
+ for _, img in images.items():
+ write_next_bytes(fid, img.id, "i")
+ write_next_bytes(fid, img.qvec.tolist(), "dddd")
+ write_next_bytes(fid, img.tvec.tolist(), "ddd")
+ write_next_bytes(fid, img.camera_id, "i")
+ for char in img.name:
+ write_next_bytes(fid, char.encode("utf-8"), "c")
+ write_next_bytes(fid, b"\x00", "c")
+ write_next_bytes(fid, len(img.point3D_ids), "Q")
+ for xy, p3d_id in zip(img.xys, img.point3D_ids):
+ write_next_bytes(fid, [*xy, p3d_id], "ddq")
+
+
def read_points3D_text(path):
"""
see: src/base/reconstruction.cc
@@ -232,7 +337,7 @@ def read_points3D_text(path):
return points3D
-def read_points3d_binary(path_to_model_file):
+def read_points3D_binary(path_to_model_file):
"""
see: src/base/reconstruction.cc
void Reconstruction::ReadPoints3DBinary(const std::string& path)
@@ -241,7 +346,7 @@ def read_points3d_binary(path_to_model_file):
points3D = {}
with open(path_to_model_file, "rb") as fid:
num_points = read_next_bytes(fid, 8, "Q")[0]
- for point_line_index in range(num_points):
+ for _ in range(num_points):
binary_point_line_properties = read_next_bytes(fid, num_bytes=43, format_char_sequence="QdddBBBd")
point3D_id = binary_point_line_properties[0]
xyz = np.array(binary_point_line_properties[1:4])
@@ -257,7 +362,75 @@ def read_points3d_binary(path_to_model_file):
return points3D
-def read_model(path, ext):
+def write_points3D_text(points3D, path):
+ """
+ see: src/base/reconstruction.cc
+ void Reconstruction::ReadPoints3DText(const std::string& path)
+ void Reconstruction::WritePoints3DText(const std::string& path)
+ """
+ if len(points3D) == 0:
+ mean_track_length = 0
+ else:
+ mean_track_length = sum((len(pt.image_ids) for _, pt in points3D.items())) / len(points3D)
+ HEADER = (
+ "# 3D point list with one line of data per point:\n"
+ + "# POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[] as (IMAGE_ID, POINT2D_IDX)\n"
+ + "# Number of points: {}, mean track length: {}\n".format(len(points3D), mean_track_length)
+ )
+
+ with open(path, "w") as fid:
+ fid.write(HEADER)
+ for _, pt in points3D.items():
+ point_header = [pt.id, *pt.xyz, *pt.rgb, pt.error]
+ fid.write(" ".join(map(str, point_header)) + " ")
+ track_strings = []
+ for image_id, point2D in zip(pt.image_ids, pt.point2D_idxs):
+ track_strings.append(" ".join(map(str, [image_id, point2D])))
+ fid.write(" ".join(track_strings) + "\n")
+
+
+def write_points3D_binary(points3D, path_to_model_file):
+ """
+ see: src/base/reconstruction.cc
+ void Reconstruction::ReadPoints3DBinary(const std::string& path)
+ void Reconstruction::WritePoints3DBinary(const std::string& path)
+ """
+ with open(path_to_model_file, "wb") as fid:
+ write_next_bytes(fid, len(points3D), "Q")
+ for _, pt in points3D.items():
+ write_next_bytes(fid, pt.id, "Q")
+ write_next_bytes(fid, pt.xyz.tolist(), "ddd")
+ write_next_bytes(fid, pt.rgb.tolist(), "BBB")
+ write_next_bytes(fid, pt.error, "d")
+ track_length = pt.image_ids.shape[0]
+ write_next_bytes(fid, track_length, "Q")
+ for image_id, point2D_id in zip(pt.image_ids, pt.point2D_idxs):
+ write_next_bytes(fid, [image_id, point2D_id], "ii")
+
+
+def detect_model_format(path, ext):
+ if (
+ os.path.isfile(os.path.join(path, "cameras" + ext))
+ and os.path.isfile(os.path.join(path, "images" + ext))
+ and os.path.isfile(os.path.join(path, "points3D" + ext))
+ ):
+ print("Detected model format: '" + ext + "'")
+ return True
+
+ return False
+
+
+def read_model(path, ext=""):
+ # try to detect the extension automatically
+ if ext == "":
+ if detect_model_format(path, ".bin"):
+ ext = ".bin"
+ elif detect_model_format(path, ".txt"):
+ ext = ".txt"
+ else:
+ print("Provide model format: '.bin' or '.txt'")
+ return
+
if ext == ".txt":
cameras = read_cameras_text(os.path.join(path, "cameras" + ext))
images = read_images_text(os.path.join(path, "images" + ext))
@@ -265,7 +438,19 @@ def read_model(path, ext):
else:
cameras = read_cameras_binary(os.path.join(path, "cameras" + ext))
images = read_images_binary(os.path.join(path, "images" + ext))
- points3D = read_points3d_binary(os.path.join(path, "points3D") + ext)
+ points3D = read_points3D_binary(os.path.join(path, "points3D") + ext)
+ return cameras, images, points3D
+
+
+def write_model(cameras, images, points3D, path, ext=".bin"):
+ if ext == ".txt":
+ write_cameras_text(cameras, os.path.join(path, "cameras" + ext))
+ write_images_text(images, os.path.join(path, "images" + ext))
+ write_points3D_text(points3D, os.path.join(path, "points3D") + ext)
+ else:
+ write_cameras_binary(cameras, os.path.join(path, "cameras" + ext))
+ write_images_binary(images, os.path.join(path, "images" + ext))
+ write_points3D_binary(points3D, os.path.join(path, "points3D") + ext)
return cameras, images, points3D
diff --git a/nerfstudio/engine/schedulers.py b/nerfstudio/engine/schedulers.py
index 0b6586ffc..48a656650 100644
--- a/nerfstudio/engine/schedulers.py
+++ b/nerfstudio/engine/schedulers.py
@@ -107,3 +107,36 @@ class ExponentialDecayScheduler(Scheduler):
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=func)
return scheduler
+
+
+@dataclass
+class CosineDecaySchedulerConfig(SchedulerConfig):
+ """Basic scheduler config with self-defined exponential decay schedule"""
+
+ _target: Type = field(default_factory=lambda: CosineDecayScheduler)
+ """target class to instantiate"""
+ warm_up_end: int = 5000
+ """Iteration number where warmp ends"""
+ learning_rate_alpha: float = 0.05
+ """Learning rate alpha value"""
+ max_steps: int = 300000
+ """The maximum number of steps."""
+
+
+class CosineDecayScheduler(Scheduler):
+ """Starts with a flat lr schedule until it reaches N epochs then applies a given scheduler"""
+
+ config: CosineDecaySchedulerConfig
+
+ def get_scheduler(self, optimizer: Optimizer, lr_init: float) -> lr_scheduler._LRScheduler:
+ def func(step):
+ if step < self.config.warm_up_end:
+ learning_factor = step / self.config.warm_up_end
+ else:
+ alpha = self.config.learning_rate_alpha
+ progress = (step - self.config.warm_up_end) / (self.config.max_steps - self.config.warm_up_end)
+ learning_factor = (np.cos(np.pi * progress) + 1.0) * 0.5 * (1 - alpha) + alpha
+ return learning_factor
+
+ scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=func)
+ return scheduler
diff --git a/nerfstudio/field_components/field_heads.py b/nerfstudio/field_components/field_heads.py
index 953b5c555..45c420894 100644
--- a/nerfstudio/field_components/field_heads.py
+++ b/nerfstudio/field_components/field_heads.py
@@ -39,6 +39,9 @@ class FieldHeadNames(Enum):
TRANSIENT_DENSITY = "transient_density"
SEMANTICS = "semantics"
FEATURE = "feature"
+ SDF = "sdf"
+ ALPHA = "alpha"
+ GRADIENT = "gradient"
class FieldHead(FieldComponent):
diff --git a/nerfstudio/fields/base_field.py b/nerfstudio/fields/base_field.py
index e20774d88..d8e4713d1 100644
--- a/nerfstudio/fields/base_field.py
+++ b/nerfstudio/fields/base_field.py
@@ -17,16 +17,26 @@ Base class for the graphs.
"""
from abc import abstractmethod
-from typing import Dict, Optional, Tuple
+from dataclasses import dataclass, field
+from typing import Dict, Optional, Tuple, Type
import torch
from torch import nn
from torchtyping import TensorType
from nerfstudio.cameras.rays import Frustums, RaySamples
+from nerfstudio.configs.base_config import InstantiateConfig
from nerfstudio.field_components.field_heads import FieldHeadNames
+@dataclass
+class FieldConfig(InstantiateConfig):
+ """Configuration for field instantiation"""
+
+ _target: Type = field(default_factory=lambda: Field)
+ """target class to instantiate"""
+
+
class Field(nn.Module):
"""Base class for fields."""
@@ -116,3 +126,12 @@ class Field(nn.Module):
normals = self.get_normals()
field_outputs[FieldHeadNames.NORMALS] = normals # type: ignore
return field_outputs
+
+
+def shift_directions_for_tcnn(directions: TensorType["bs":..., 3]) -> TensorType["bs":..., 3]:
+ """Shift directions from [-1, 1] to [0, 1]
+
+ Args:
+ directions: batch of directions
+ """
+ return (directions + 1.0) / 2.0
diff --git a/nerfstudio/fields/instant_ngp_field.py b/nerfstudio/fields/instant_ngp_field.py
index b0f4e97e4..d1b3e2fa7 100644
--- a/nerfstudio/fields/instant_ngp_field.py
+++ b/nerfstudio/fields/instant_ngp_field.py
@@ -30,7 +30,7 @@ from nerfstudio.data.scene_box import SceneBox
from nerfstudio.field_components.activations import trunc_exp
from nerfstudio.field_components.embedding import Embedding
from nerfstudio.field_components.field_heads import FieldHeadNames
-from nerfstudio.fields.base_field import Field
+from nerfstudio.fields.base_field import Field, shift_directions_for_tcnn
try:
import tinycudann as tcnn
@@ -39,15 +39,6 @@ except ImportError:
pass
-def get_normalized_directions(directions: TensorType["bs":..., 3]) -> TensorType["bs":..., 3]:
- """SH encoding must be in the range [0, 1]
-
- Args:
- directions: batch of directions
- """
- return (directions + 1.0) / 2.0
-
-
class TCNNInstantNGPField(Field):
"""TCNN implementation of the Instant-NGP field.
@@ -159,7 +150,7 @@ class TCNNInstantNGPField(Field):
def get_outputs(
self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None
) -> Dict[FieldHeadNames, TensorType]:
- directions = get_normalized_directions(ray_samples.frustums.directions)
+ directions = shift_directions_for_tcnn(ray_samples.frustums.directions)
directions_flat = directions.view(-1, 3)
d = self.direction_encoding(directions_flat)
diff --git a/nerfstudio/fields/nerfacto_field.py b/nerfstudio/fields/nerfacto_field.py
index 830962edf..de0fd7c7d 100644
--- a/nerfstudio/fields/nerfacto_field.py
+++ b/nerfstudio/fields/nerfacto_field.py
@@ -46,7 +46,7 @@ from nerfstudio.field_components.spatial_distortions import (
SceneContraction,
SpatialDistortion,
)
-from nerfstudio.fields.base_field import Field
+from nerfstudio.fields.base_field import Field, shift_directions_for_tcnn
try:
import tinycudann as tcnn
@@ -55,15 +55,6 @@ except ImportError:
pass
-def get_normalized_directions(directions: TensorType["bs":..., 3]) -> TensorType["bs":..., 3]:
- """SH encoding must be in the range [0, 1]
-
- Args:
- directions: batch of directions
- """
- return (directions + 1.0) / 2.0
-
-
class TCNNNerfactoField(Field):
"""Compound Field that uses TCNN
@@ -267,7 +258,7 @@ class TCNNNerfactoField(Field):
if ray_samples.camera_indices is None:
raise AttributeError("Camera indices are not provided.")
camera_indices = ray_samples.camera_indices.squeeze()
- directions = get_normalized_directions(ray_samples.frustums.directions)
+ directions = shift_directions_for_tcnn(ray_samples.frustums.directions)
directions_flat = directions.view(-1, 3)
d = self.direction_encoding(directions_flat)
diff --git a/nerfstudio/fields/nerfplayer_nerfacto_field.py b/nerfstudio/fields/nerfplayer_nerfacto_field.py
index 0b186464f..cb21013e6 100644
--- a/nerfstudio/fields/nerfplayer_nerfacto_field.py
+++ b/nerfstudio/fields/nerfplayer_nerfacto_field.py
@@ -38,7 +38,7 @@ from nerfstudio.field_components.field_heads import (
)
from nerfstudio.field_components.spatial_distortions import SpatialDistortion
from nerfstudio.field_components.temporal_grid import TemporalGridEncoder
-from nerfstudio.fields.base_field import Field
+from nerfstudio.fields.base_field import Field, shift_directions_for_tcnn
try:
import tinycudann as tcnn
@@ -47,15 +47,6 @@ except ImportError:
pass
-def get_normalized_directions(directions: TensorType["bs":..., 3]) -> TensorType["bs":..., 3]:
- """SH encoding must be in the range [0, 1]
-
- Args:
- directions: batch of directions
- """
- return (directions + 1.0) / 2.0
-
-
class TemporalHashMLPDensityField(Field):
"""A lightweight temporal density field module.
@@ -347,7 +338,7 @@ class NerfplayerNerfactoField(Field):
if ray_samples.camera_indices is None:
raise AttributeError("Camera indices are not provided.")
camera_indices = ray_samples.camera_indices.squeeze()
- directions = get_normalized_directions(ray_samples.frustums.directions)
+ directions = shift_directions_for_tcnn(ray_samples.frustums.directions)
directions_flat = directions.view(-1, 3)
d = self.direction_encoding(directions_flat)
diff --git a/nerfstudio/fields/nerfplayer_ngp_field.py b/nerfstudio/fields/nerfplayer_ngp_field.py
index 829ff0d33..2952a4369 100644
--- a/nerfstudio/fields/nerfplayer_ngp_field.py
+++ b/nerfstudio/fields/nerfplayer_ngp_field.py
@@ -30,8 +30,7 @@ from nerfstudio.field_components.activations import trunc_exp
from nerfstudio.field_components.embedding import Embedding
from nerfstudio.field_components.field_heads import FieldHeadNames
from nerfstudio.field_components.temporal_grid import TemporalGridEncoder
-from nerfstudio.fields.base_field import Field
-from nerfstudio.fields.instant_ngp_field import get_normalized_directions
+from nerfstudio.fields.base_field import Field, shift_directions_for_tcnn
try:
import tinycudann as tcnn
@@ -164,7 +163,7 @@ class NerfplayerNGPField(Field):
def get_outputs(
self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None
) -> Dict[FieldHeadNames, TensorType]:
- directions = get_normalized_directions(ray_samples.frustums.directions)
+ directions = shift_directions_for_tcnn(ray_samples.frustums.directions)
directions_flat = directions.view(-1, 3)
if self.direction_encoding is not None:
diff --git a/nerfstudio/fields/sdf_field.py b/nerfstudio/fields/sdf_field.py
new file mode 100644
index 000000000..e9fd10cb8
--- /dev/null
+++ b/nerfstudio/fields/sdf_field.py
@@ -0,0 +1,458 @@
+# 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.
+
+"""
+Field for SDF based model, rather then estimating density to generate a surface,
+a signed distance function (SDF) for surface representation is used to help with extracting high fidelity surfaces
+"""
+
+from dataclasses import dataclass, field
+from typing import Dict, Optional, Type
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn.parameter import Parameter
+from torchtyping import TensorType
+from typing_extensions import Literal
+
+from nerfstudio.cameras.rays import RaySamples
+from nerfstudio.field_components.embedding import Embedding
+from nerfstudio.field_components.encodings import NeRFEncoding
+from nerfstudio.field_components.field_heads import FieldHeadNames
+from nerfstudio.field_components.spatial_distortions import SpatialDistortion
+from nerfstudio.fields.base_field import Field, FieldConfig
+
+try:
+ import tinycudann as tcnn
+except ImportError:
+ # tinycudann module doesn't exist
+ pass
+
+
+class LearnedVariance(nn.Module):
+ """Variance network in NeuS
+
+ Args:
+ init_val: initial value in NeuS variance network
+ """
+
+ def __init__(self, init_val):
+ super().__init__()
+ self.register_parameter("variance", nn.Parameter(init_val * torch.ones(1), requires_grad=True))
+
+ def forward(self, x: TensorType[1]) -> TensorType[1]:
+ """Returns current variance value"""
+ return torch.ones([len(x), 1], device=x.device) * torch.exp(self.variance * 10.0)
+
+ def get_variance(self) -> TensorType[1]:
+ """return current variance value"""
+ return torch.exp(self.variance * 10.0).clip(1e-6, 1e6)
+
+
+@dataclass
+class SDFFieldConfig(FieldConfig):
+ """SDF Field Config"""
+
+ _target: Type = field(default_factory=lambda: SDFField)
+ num_layers: int = 8
+ """Number of layers for geometric network"""
+ hidden_dim: int = 256
+ """Number of hidden dimension of geometric network"""
+ geo_feat_dim: int = 256
+ """Dimension of geometric feature"""
+ num_layers_color: int = 4
+ """Number of layers for color network"""
+ hidden_dim_color: int = 256
+ """Number of hidden dimension of color network"""
+ appearance_embedding_dim: int = 32
+ """Dimension of appearance embedding"""
+ use_appearance_embedding: bool = False
+ """Dimension of appearance embedding"""
+ bias: float = 0.8
+ """sphere size of geometric initializaion"""
+ geometric_init: bool = True
+ """Whether to use geometric initialization"""
+ inside_outside: bool = True
+ """whether to revert signed distance value, set to True for indoor scene"""
+ weight_norm: bool = True
+ """Whether to use weight norm for linear layer"""
+ use_grid_feature: bool = False
+ """Whether to use multi-resolution feature grids"""
+ divide_factor: float = 2.0
+ """Normalization factor for multi-resolution grids"""
+ beta_init: float = 0.1
+ """Init learnable beta value for transformation of sdf to density"""
+ encoding_type: Literal["hash", "periodic", "tensorf_vm"] = "hash"
+ num_levels = 16
+ """Number of encoding levels"""
+ max_res = 2048
+ """Maximum resolution of the encoding"""
+ base_res = 16
+ """Base resolution of the encoding"""
+ log2_hashmap_size = 19
+ """Size of the hash map"""
+ features_per_level = 2
+ """Number of features per encoding level"""
+ use_hash = True
+ """Whether to use hash encoding"""
+ smoothstep = True
+ """Whether to use the smoothstep function"""
+
+
+class SDFField(Field):
+ """
+ A field for Signed Distance Functions (SDF).
+
+ Args:
+ config: The configuration for the SDF field.
+ aabb: An axis-aligned bounding box for the SDF field.
+ num_images: The number of images for embedding appearance.
+ use_average_appearance_embedding: Whether to use average appearance embedding. Defaults to False.
+ spatial_distortion: The spatial distortion. Defaults to None.
+ """
+
+ config: SDFFieldConfig
+
+ def __init__(
+ self,
+ config: SDFFieldConfig,
+ aabb: TensorType[2, 3],
+ num_images: int,
+ use_average_appearance_embedding: bool = False,
+ spatial_distortion: Optional[SpatialDistortion] = None,
+ ) -> None:
+ super().__init__()
+ self.config = config
+
+ self.aabb = Parameter(aabb, requires_grad=False)
+
+ self.spatial_distortion = spatial_distortion
+ self.num_images = num_images
+
+ self.embedding_appearance = Embedding(self.num_images, self.config.appearance_embedding_dim)
+ self.use_average_appearance_embedding = use_average_appearance_embedding
+ self.use_grid_feature = self.config.use_grid_feature
+ self.divide_factor = self.config.divide_factor
+
+ growth_factor = np.exp((np.log(config.max_res) - np.log(config.base_res)) / (config.num_levels - 1))
+
+ if self.config.encoding_type == "hash":
+ # feature encoding
+ self.encoding = tcnn.Encoding(
+ n_input_dims=3,
+ encoding_config={
+ "otype": "HashGrid" if config.use_hash else "DenseGrid",
+ "n_levels": config.num_levels,
+ "n_features_per_level": config.features_per_level,
+ "log2_hashmap_size": config.log2_hashmap_size,
+ "base_resolution": config.base_res,
+ "per_level_scale": growth_factor,
+ "interpolation": "Smoothstep" if config.smoothstep else "Linear",
+ },
+ )
+
+ # we concat inputs position ourselves
+ self.position_encoding = NeRFEncoding(
+ in_dim=3, num_frequencies=6, min_freq_exp=0.0, max_freq_exp=5.0, include_input=False
+ )
+
+ self.direction_encoding = NeRFEncoding(
+ in_dim=3, num_frequencies=4, min_freq_exp=0.0, max_freq_exp=3.0, include_input=True
+ )
+
+ # initialize geometric network
+ self.initialize_geo_layers()
+
+ # deviation_network to compute alpha from sdf from NeuS
+ self.deviation_network = LearnedVariance(init_val=self.config.beta_init)
+
+ # color network
+ dims = [self.config.hidden_dim_color for _ in range(self.config.num_layers_color)]
+ # point, view_direction, normal, feature, embedding
+ in_dim = (
+ 3
+ + self.direction_encoding.get_out_dim()
+ + 3
+ + self.config.geo_feat_dim
+ + self.embedding_appearance.get_out_dim()
+ )
+ dims = [in_dim] + dims + [3]
+ self.num_layers_color = len(dims)
+
+ for l in range(0, self.num_layers_color - 1):
+ out_dim = dims[l + 1]
+ lin = nn.Linear(dims[l], out_dim)
+
+ if self.config.weight_norm:
+ lin = nn.utils.weight_norm(lin)
+ setattr(self, "clin" + str(l), lin)
+
+ self.softplus = nn.Softplus(beta=100)
+ self.relu = nn.ReLU()
+ self.sigmoid = torch.nn.Sigmoid()
+
+ self._cos_anneal_ratio = 1.0
+
+ def initialize_geo_layers(self) -> None:
+ """
+ Initialize layers for geometric network (sdf)
+ """
+ # MLP with geometric initialization
+ dims = [self.config.hidden_dim for _ in range(self.config.num_layers)]
+ in_dim = 3 + self.position_encoding.get_out_dim() + self.encoding.n_output_dims
+ dims = [in_dim] + dims + [1 + self.config.geo_feat_dim]
+ self.num_layers = len(dims)
+ self.skip_in = [4]
+
+ for l in range(0, self.num_layers - 1):
+ if l + 1 in self.skip_in:
+ out_dim = dims[l + 1] - dims[0]
+ else:
+ out_dim = dims[l + 1]
+
+ lin = nn.Linear(dims[l], out_dim)
+
+ if self.config.geometric_init:
+ if l == self.num_layers - 2:
+ if not self.config.inside_outside:
+ torch.nn.init.normal_(lin.weight, mean=np.sqrt(np.pi) / np.sqrt(dims[l]), std=0.0001)
+ torch.nn.init.constant_(lin.bias, -self.config.bias)
+ else:
+ torch.nn.init.normal_(lin.weight, mean=-np.sqrt(np.pi) / np.sqrt(dims[l]), std=0.0001)
+ torch.nn.init.constant_(lin.bias, self.config.bias)
+ elif l == 0:
+ torch.nn.init.constant_(lin.bias, 0.0)
+ torch.nn.init.constant_(lin.weight[:, 3:], 0.0)
+ torch.nn.init.normal_(lin.weight[:, :3], 0.0, np.sqrt(2) / np.sqrt(out_dim))
+ elif l in self.skip_in:
+ torch.nn.init.constant_(lin.bias, 0.0)
+ torch.nn.init.normal_(lin.weight, 0.0, np.sqrt(2) / np.sqrt(out_dim))
+ torch.nn.init.constant_(lin.weight[:, -(dims[0] - 3) :], 0.0)
+ else:
+ torch.nn.init.constant_(lin.bias, 0.0)
+ torch.nn.init.normal_(lin.weight, 0.0, np.sqrt(2) / np.sqrt(out_dim))
+
+ if self.config.weight_norm:
+ lin = nn.utils.weight_norm(lin)
+ setattr(self, "glin" + str(l), lin)
+
+ def set_cos_anneal_ratio(self, anneal: float) -> None:
+ """Set the anneal value for the proposal network."""
+ self._cos_anneal_ratio = anneal
+
+ def forward_geonetwork(self, inputs: TensorType[..., 3]) -> TensorType[..., "geo-features+1"]:
+ """forward the geonetwork"""
+ if self.use_grid_feature:
+ positions = self.spatial_distortion(inputs)
+
+ positions = (positions + 1.0) / 2.0
+ feature = self.encoding(positions)
+ else:
+ feature = torch.zeros_like(inputs[:, :1].repeat(1, self.encoding.n_output_dims))
+
+ pe = self.position_encoding(inputs)
+
+ inputs = torch.cat((inputs, pe, feature), dim=-1)
+
+ # Pass through layers
+ outputs = inputs
+
+ for l in range(0, self.num_layers - 1):
+ lin = getattr(self, "glin" + str(l))
+
+ if l in self.skip_in:
+ outputs = torch.cat([outputs, inputs], 1) / np.sqrt(2)
+
+ outputs = lin(outputs)
+
+ if l < self.num_layers - 2:
+ outputs = self.softplus(outputs)
+ return outputs
+
+ def get_sdf(self, ray_samples: RaySamples) -> TensorType["num_samples", -1, 1]:
+ """predict the sdf value for ray samples"""
+ positions = ray_samples.frustums.get_start_positions()
+ positions_flat = positions.view(-1, 3)
+ hidden_output = self.forward_geonetwork(positions_flat).view(*ray_samples.frustums.shape, -1)
+ sdf, _ = torch.split(hidden_output, [1, self.config.geo_feat_dim], dim=-1)
+ return sdf
+
+ def get_alpha(
+ self,
+ ray_samples: RaySamples,
+ sdf: Optional[TensorType["num_samples", -1, 1]] = None,
+ gradients: Optional[TensorType["num_samples", -1, 1]] = None,
+ ) -> TensorType["num_samples", -1, 1]:
+ """compute alpha from sdf as in NeuS"""
+ if sdf is None or gradients is None:
+ inputs = ray_samples.frustums.get_start_positions()
+ inputs.requires_grad_(True)
+ with torch.enable_grad():
+ hidden_output = self.forward_geonetwork(inputs)
+ sdf, _ = torch.split(hidden_output, [1, self.config.geo_feat_dim], dim=-1)
+ d_output = torch.ones_like(sdf, requires_grad=False, device=sdf.device)
+ gradients = torch.autograd.grad(
+ outputs=sdf,
+ inputs=inputs,
+ grad_outputs=d_output,
+ create_graph=True,
+ retain_graph=True,
+ only_inputs=True,
+ )[0]
+
+ inv_s = self.deviation_network.get_variance() # Single parameter
+
+ true_cos = (ray_samples.frustums.directions * gradients).sum(-1, keepdim=True)
+
+ # anneal as NeuS
+ cos_anneal_ratio = self._cos_anneal_ratio
+
+ # "cos_anneal_ratio" grows from 0 to 1 in the beginning training iterations. The anneal strategy below makes
+ # the cos value "not dead" at the beginning training iterations, for better convergence.
+ iter_cos = -(
+ F.relu(-true_cos * 0.5 + 0.5) * (1.0 - cos_anneal_ratio) + F.relu(-true_cos) * cos_anneal_ratio
+ ) # always non-positive
+
+ # Estimate signed distances at section points
+ estimated_next_sdf = sdf + iter_cos * ray_samples.deltas * 0.5
+ estimated_prev_sdf = sdf - iter_cos * ray_samples.deltas * 0.5
+
+ prev_cdf = torch.sigmoid(estimated_prev_sdf * inv_s)
+ next_cdf = torch.sigmoid(estimated_next_sdf * inv_s)
+
+ p = prev_cdf - next_cdf
+ c = prev_cdf
+
+ alpha = ((p + 1e-5) / (c + 1e-5)).clip(0.0, 1.0)
+
+ return alpha
+
+ def get_density(self, ray_samples: RaySamples):
+ raise NotImplementedError
+
+ def get_colors(
+ self,
+ points: TensorType[..., 3],
+ directions: TensorType[..., 3],
+ normals: TensorType[..., 3],
+ geo_features: TensorType[..., "geo-feat-dim"],
+ camera_indices: TensorType,
+ ) -> TensorType[..., 3]:
+ """compute colors"""
+ d = self.direction_encoding(directions)
+
+ # appearance
+ if self.training:
+ embedded_appearance = self.embedding_appearance(camera_indices)
+ # set it to zero if don't use it
+ if not self.config.use_appearance_embedding:
+ embedded_appearance = torch.zeros_like(embedded_appearance)
+ else:
+ if self.use_average_appearance_embedding:
+ embedded_appearance = torch.ones(
+ (*directions.shape[:-1], self.config.appearance_embedding_dim), device=directions.device
+ ) * self.embedding_appearance.mean(dim=0)
+ else:
+ embedded_appearance = torch.zeros(
+ (*directions.shape[:-1], self.config.appearance_embedding_dim), device=directions.device
+ )
+
+ hidden_input = torch.cat(
+ [
+ points,
+ d,
+ normals,
+ geo_features.view(-1, self.config.geo_feat_dim),
+ embedded_appearance.view(-1, self.config.appearance_embedding_dim),
+ ],
+ dim=-1,
+ )
+
+ for layer in range(0, self.num_layers_color - 1):
+ lin = getattr(self, "clin" + str(layer))
+
+ hidden_input = lin(hidden_input)
+
+ if layer < self.num_layers_color - 2:
+ hidden_input = self.relu(hidden_input)
+
+ rgb = self.sigmoid(hidden_input)
+
+ return rgb
+
+ def get_outputs(
+ self,
+ ray_samples: RaySamples,
+ density_embedding: Optional[TensorType] = None,
+ return_alphas: bool = False,
+ ) -> Dict[FieldHeadNames, TensorType]:
+ """compute output of ray samples"""
+ if ray_samples.camera_indices is None:
+ raise AttributeError("Camera indices are not provided.")
+
+ outputs = {}
+
+ camera_indices = ray_samples.camera_indices.squeeze()
+
+ inputs = ray_samples.frustums.get_start_positions()
+ inputs = inputs.view(-1, 3)
+
+ directions = ray_samples.frustums.directions
+ directions_flat = directions.reshape(-1, 3)
+
+ inputs.requires_grad_(True)
+ with torch.enable_grad():
+ hidden_output = self.forward_geonetwork(inputs)
+ sdf, geo_feature = torch.split(hidden_output, [1, self.config.geo_feat_dim], dim=-1)
+ d_output = torch.ones_like(sdf, requires_grad=False, device=sdf.device)
+ gradients = torch.autograd.grad(
+ outputs=sdf, inputs=inputs, grad_outputs=d_output, create_graph=True, retain_graph=True, only_inputs=True
+ )[0]
+
+ rgb = self.get_colors(inputs, directions_flat, gradients, geo_feature, camera_indices)
+
+ rgb = rgb.view(*ray_samples.frustums.directions.shape[:-1], -1)
+ sdf = sdf.view(*ray_samples.frustums.directions.shape[:-1], -1)
+ gradients = gradients.view(*ray_samples.frustums.directions.shape[:-1], -1)
+ normals = torch.nn.functional.normalize(gradients, p=2, dim=-1)
+
+ outputs.update(
+ {
+ FieldHeadNames.RGB: rgb,
+ FieldHeadNames.SDF: sdf,
+ FieldHeadNames.NORMALS: normals,
+ FieldHeadNames.GRADIENT: gradients,
+ }
+ )
+
+ if return_alphas:
+ alphas = self.get_alpha(ray_samples, sdf, gradients)
+ outputs.update({FieldHeadNames.ALPHA: alphas})
+
+ return outputs
+
+ def forward(
+ self, ray_samples: RaySamples, compute_normals: bool = False, return_alphas: bool = False
+ ) -> Dict[FieldHeadNames, TensorType]:
+ """Evaluates the field at points along the ray.
+
+ Args:
+ ray_samples: Samples to evaluate field on.
+ compute normals: not currently used in this implementation.
+ return_alphas: Whether to return alpha values
+ """
+ field_outputs = self.get_outputs(ray_samples, return_alphas=return_alphas)
+ return field_outputs
diff --git a/nerfstudio/model_components/losses.py b/nerfstudio/model_components/losses.py
index 419a8a1f4..174c131cc 100644
--- a/nerfstudio/model_components/losses.py
+++ b/nerfstudio/model_components/losses.py
@@ -23,7 +23,7 @@ from torchtyping import TensorType
from typing_extensions import Literal
from nerfstudio.cameras.rays import RaySamples
-from nerfstudio.utils.math import masked_reduction
+from nerfstudio.utils.math import masked_reduction, normalized_depth_scale_and_shift
L1Loss = nn.L1Loss
MSELoss = nn.MSELoss
@@ -355,7 +355,8 @@ class MiDaSMSELoss(nn.Module):
"""
summed_mask = torch.sum(mask, (1, 2))
image_loss = torch.sum(self.mse_loss(prediction, target) * mask, (1, 2))
- image_loss = masked_reduction(image_loss, summed_mask, self.reduction_type)
+ # multiply by 2 magic number?
+ image_loss = masked_reduction(image_loss, 2 * summed_mask, self.reduction_type)
return image_loss
@@ -456,47 +457,6 @@ class ScaleAndShiftInvariantLoss(nn.Module):
self.__prediction_ssi = None
- @classmethod
- def compute_scale_and_shift(
- cls, prediction: TensorType[1, 32, "mult"], target: TensorType[1, 32, "mult"], mask: TensorType[1, 32, "mult"]
- ):
- """
- More info here: https://arxiv.org/pdf/2206.00665.pdf supplementary section A2 Depth Consistency Loss
- This function computes scale/shift required to normalizes predicted depth map,
- to allow for using normalized depth maps as input from monocular depth estimation networks.
- These networks are trained such that they predict normalized depth maps.
-
- Solves for scale/shift using a least squares approach with a closed form solution:
- Based on:
- https://github.com/autonomousvision/monosdf/blob/d9619e948bf3d85c6adec1a643f679e2e8e84d4b/code/model/loss.py#L7
- Args:
- prediction: predicted depth map
- target: ground truth depth map
- mask: mask of valid pixels
- Returns:
- scale and shift for depth prediction
- """
- # system matrix: A = [[a_00, a_01], [a_10, a_11]]
- a_00 = torch.sum(mask * prediction * prediction, (1, 2))
- a_01 = torch.sum(mask * prediction, (1, 2))
- a_11 = torch.sum(mask, (1, 2))
-
- # right hand side: b = [b_0, b_1]
- b_0 = torch.sum(mask * prediction * target, (1, 2))
- b_1 = torch.sum(mask * target, (1, 2))
-
- # solution: x = A^-1 . b = [[a_11, -a_01], [-a_10, a_00]] / (a_00 * a_11 - a_01 * a_10) . b
- scale = torch.zeros_like(b_0)
- shift = torch.zeros_like(b_1)
-
- det = a_00 * a_11 - a_01 * a_01
- valid = det.nonzero()
-
- scale[valid] = (a_11[valid] * b_0[valid] - a_01[valid] * b_1[valid]) / det[valid]
- shift[valid] = (-a_01[valid] * b_0[valid] + a_00[valid] * b_1[valid]) / det[valid]
-
- return scale, shift
-
def forward(
self, prediction: TensorType[1, 32, "mult"], target: TensorType[1, 32, "mult"], mask: TensorType[1, 32, "mult"]
) -> TensorType[0]:
@@ -508,7 +468,7 @@ class ScaleAndShiftInvariantLoss(nn.Module):
Returns:
scale and shift invariant loss
"""
- scale, shift = ScaleAndShiftInvariantLoss.compute_scale_and_shift(prediction, target, mask)
+ scale, shift = normalized_depth_scale_and_shift(prediction, target, mask)
self.__prediction_ssi = scale.view(-1, 1, 1) * prediction + shift.view(-1, 1, 1)
total = self.__data_loss(self.__prediction_ssi, target, mask)
diff --git a/nerfstudio/model_components/ray_samplers.py b/nerfstudio/model_components/ray_samplers.py
index 63a01f4fa..d6cb080c6 100644
--- a/nerfstudio/model_components/ray_samplers.py
+++ b/nerfstudio/model_components/ray_samplers.py
@@ -17,7 +17,7 @@ Collection of sampling strategies
"""
from abc import abstractmethod
-from typing import Callable, List, Optional, Tuple
+from typing import Callable, List, Optional, Tuple, Union
import nerfacc
import torch
@@ -602,3 +602,166 @@ class ProposalNetworkSampler(Sampler):
assert ray_samples is not None
return ray_samples, weights_list, ray_samples_list
+
+
+class NeuSSampler(Sampler):
+ """NeuS sampler that uses a sdf network to generate samples with fixed variance value in each iterations."""
+
+ def __init__(
+ self,
+ num_samples: int = 64,
+ num_samples_importance: int = 64,
+ num_samples_outside: int = 32,
+ num_upsample_steps: int = 4,
+ base_variance: float = 64,
+ single_jitter: bool = True,
+ ) -> None:
+ super().__init__()
+ self.num_samples = num_samples
+ self.num_samples_importance = num_samples_importance
+ self.num_samples_outside = num_samples_outside
+ self.num_upsample_steps = num_upsample_steps
+ self.base_variance = base_variance
+ self.single_jitter = single_jitter
+
+ # samplers
+ self.uniform_sampler = UniformSampler(single_jitter=single_jitter)
+ self.pdf_sampler = PDFSampler(
+ include_original=False,
+ single_jitter=single_jitter,
+ histogram_padding=1e-5,
+ )
+ self.outside_sampler = LinearDisparitySampler()
+
+ def generate_ray_samples(
+ self,
+ ray_bundle: Optional[RayBundle] = None,
+ sdf_fn: Optional[Callable] = None,
+ ray_samples: Optional[RaySamples] = None,
+ ) -> Union[Tuple[RaySamples, torch.Tensor], RaySamples]:
+ assert ray_bundle is not None
+ assert sdf_fn is not None
+
+ # Start with uniform sampling
+ if ray_samples is None:
+ ray_samples = self.uniform_sampler(ray_bundle, num_samples=self.num_samples)
+
+ total_iters = 0
+ sorted_index = None
+ new_samples = ray_samples
+
+ base_variance = self.base_variance
+
+ while total_iters < self.num_upsample_steps:
+
+ with torch.no_grad():
+ new_sdf = sdf_fn(new_samples)
+
+ # merge sdf predictions
+ if sorted_index is not None:
+ sdf_merge = torch.cat([sdf.squeeze(-1), new_sdf.squeeze(-1)], -1)
+ sdf = torch.gather(sdf_merge, 1, sorted_index).unsqueeze(-1)
+ else:
+ sdf = new_sdf
+
+ # compute with fix variances
+ alphas = self.rendering_sdf_with_fixed_inv_s(
+ ray_samples, sdf.reshape(ray_samples.shape), inv_s=base_variance * 2**total_iters
+ )
+
+ weights = ray_samples.get_weights_and_transmittance_from_alphas(alphas[..., None], weights_only=True)
+ weights = torch.cat((weights, torch.zeros_like(weights[:, :1])), dim=1)
+
+ new_samples = self.pdf_sampler(
+ ray_bundle,
+ ray_samples,
+ weights,
+ num_samples=self.num_samples_importance // self.num_upsample_steps,
+ )
+
+ ray_samples, sorted_index = self.merge_ray_samples(ray_bundle, ray_samples, new_samples)
+
+ total_iters += 1
+
+ return ray_samples
+
+ @staticmethod
+ def rendering_sdf_with_fixed_inv_s(
+ ray_samples: RaySamples, sdf: TensorType["num_samples", -1], inv_s: int
+ ) -> TensorType["num_samples", -1]:
+ """
+ rendering given a fixed inv_s as NeuS
+
+ Args:
+ ray_samples: samples along ray
+ sdf: sdf values along ray
+ inv_s: fixed variance value
+ Returns:
+ alpha value
+ """
+ batch_size = ray_samples.shape[0]
+ prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
+ deltas = ray_samples.deltas[:, :-1, 0]
+ mid_sdf = (prev_sdf + next_sdf) * 0.5
+ cos_val = (next_sdf - prev_sdf) / (deltas + 1e-5)
+
+ # ----------------------------------------------------------------------------------------------------------
+ # Use min value of [ cos, prev_cos ]
+ # Though it makes the sampling (not rendering) a little bit biased, this strategy can make the sampling more
+ # robust when meeting situations like below:
+ #
+ # SDF
+ # ^
+ # |\ -----x----...
+ # | \ /
+ # | x x
+ # |---\----/-------------> 0 level
+ # | \ /
+ # | \/
+ # |
+ # ----------------------------------------------------------------------------------------------------------
+ prev_cos_val = torch.cat([torch.zeros([batch_size, 1], device=sdf.device), cos_val[:, :-1]], dim=-1)
+ cos_val = torch.stack([prev_cos_val, cos_val], dim=-1)
+ cos_val, _ = torch.min(cos_val, dim=-1, keepdim=False)
+ cos_val = cos_val.clip(-1e3, 0.0)
+
+ dist = deltas
+ prev_esti_sdf = mid_sdf - cos_val * dist * 0.5
+ next_esti_sdf = mid_sdf + cos_val * dist * 0.5
+ prev_cdf = torch.sigmoid(prev_esti_sdf * inv_s)
+ next_cdf = torch.sigmoid(next_esti_sdf * inv_s)
+ alpha = (prev_cdf - next_cdf + 1e-5) / (prev_cdf + 1e-5)
+
+ return alpha
+
+ @staticmethod
+ def merge_ray_samples(ray_bundle: RayBundle, ray_samples_1: RaySamples, ray_samples_2: RaySamples):
+ """Merge two set of ray samples and return sorted index which can be used to merge sdf values
+ Args:
+ ray_samples_1 : ray_samples to merge
+ ray_samples_2 : ray_samples to merge
+ """
+
+ starts_1 = ray_samples_1.spacing_starts[..., 0]
+ starts_2 = ray_samples_2.spacing_starts[..., 0]
+
+ ends = torch.maximum(ray_samples_1.spacing_ends[..., -1:, 0], ray_samples_2.spacing_ends[..., -1:, 0])
+
+ bins, sorted_index = torch.sort(torch.cat([starts_1, starts_2], -1), -1)
+
+ bins = torch.cat([bins, ends], dim=-1)
+
+ # Stop gradients
+ bins = bins.detach()
+
+ euclidean_bins = ray_samples_1.spacing_to_euclidean_fn(bins)
+
+ ray_samples = ray_bundle.get_ray_samples(
+ bin_starts=euclidean_bins[..., :-1, None],
+ bin_ends=euclidean_bins[..., 1:, None],
+ spacing_starts=bins[..., :-1, None],
+ spacing_ends=bins[..., 1:, None],
+ spacing_to_euclidean_fn=ray_samples_1.spacing_to_euclidean_fn,
+ )
+
+ return ray_samples, sorted_index
diff --git a/nerfstudio/models/base_surface_model.py b/nerfstudio/models/base_surface_model.py
new file mode 100644
index 000000000..c83a07032
--- /dev/null
+++ b/nerfstudio/models/base_surface_model.py
@@ -0,0 +1,397 @@
+# 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.
+
+"""
+Implementation of Base surface model.
+"""
+
+from __future__ import annotations
+
+from abc import abstractmethod
+from dataclasses import dataclass, field
+from typing import Dict, List, Tuple, Type
+
+import torch
+import torch.nn.functional as F
+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.field_components.encodings import NeRFEncoding
+from nerfstudio.field_components.field_heads import FieldHeadNames
+from nerfstudio.field_components.spatial_distortions import SceneContraction
+from nerfstudio.fields.nerfacto_field import TCNNNerfactoField
+from nerfstudio.fields.sdf_field import SDFFieldConfig
+from nerfstudio.fields.vanilla_nerf_field import NeRFField
+from nerfstudio.model_components.losses import (
+ L1Loss,
+ MSELoss,
+ ScaleAndShiftInvariantLoss,
+ monosdf_normal_loss,
+)
+from nerfstudio.model_components.ray_samplers import LinearDisparitySampler
+from nerfstudio.model_components.renderers import (
+ AccumulationRenderer,
+ DepthRenderer,
+ RGBRenderer,
+ SemanticRenderer,
+)
+from nerfstudio.model_components.scene_colliders import AABBBoxCollider, NearFarCollider
+from nerfstudio.models.base_model import Model, ModelConfig
+from nerfstudio.utils import colormaps
+from nerfstudio.utils.colors import get_color
+from nerfstudio.utils.math import normalized_depth_scale_and_shift
+
+
+@dataclass
+class SurfaceModelConfig(ModelConfig):
+ """Surface Model Config"""
+
+ _target: Type = field(default_factory=lambda: SurfaceModel)
+ near_plane: float = 0.05
+ """How far along the ray to start sampling."""
+ far_plane: float = 4.0
+ """How far along the ray to stop sampling."""
+ far_plane_bg: float = 1000.0
+ """How far along the ray to stop sampling of the background model."""
+ background_color: Literal["random", "last_sample", "white", "black"] = "black"
+ """Whether to randomize the background color."""
+ use_average_appearance_embedding: bool = False
+ """Whether to use average appearance embedding or zeros for inference."""
+ eikonal_loss_mult: float = 0.1
+ """Monocular normal consistency loss multiplier."""
+ fg_mask_loss_mult: float = 0.01
+ """Foreground mask loss multiplier."""
+ mono_normal_loss_mult: float = 0.0
+ """Monocular normal consistency loss multiplier."""
+ mono_depth_loss_mult: float = 0.0
+ """Monocular depth consistency loss multiplier."""
+ sdf_field: SDFFieldConfig = SDFFieldConfig()
+ """Config for SDF Field"""
+ background_model: Literal["grid", "mlp", "none"] = "mlp"
+ """background models"""
+ num_samples_outside: int = 32
+ """Number of samples outside the bounding sphere for backgound"""
+ periodic_tvl_mult: float = 0.0
+ """Total variational loss mutliplier"""
+ overwrite_near_far_plane: bool = False
+ """whether to use near and far collider from command line"""
+
+
+class SurfaceModel(Model):
+ """Base surface model
+
+ Args:
+ config: Base surface model configuration to instantiate model
+ """
+
+ config: SurfaceModelConfig
+
+ def populate_modules(self):
+ """Set the fields and modules."""
+ super().populate_modules()
+
+ self.scene_contraction = SceneContraction(order=float("inf"))
+
+ # Can we also use contraction for sdf?
+ # Fields
+ self.field = self.config.sdf_field.setup(
+ aabb=self.scene_box.aabb,
+ spatial_distortion=self.scene_contraction,
+ num_images=self.num_train_data,
+ use_average_appearance_embedding=self.config.use_average_appearance_embedding,
+ )
+
+ # Collider
+ self.collider = AABBBoxCollider(self.scene_box, near_plane=0.05)
+
+ # command line near and far has highest priority
+ if self.config.overwrite_near_far_plane:
+ self.collider = NearFarCollider(near_plane=self.config.near_plane, far_plane=self.config.far_plane)
+
+ # background model
+ if self.config.background_model == "grid":
+ self.field_background = TCNNNerfactoField(
+ self.scene_box.aabb,
+ spatial_distortion=self.scene_contraction,
+ num_images=self.num_train_data,
+ use_average_appearance_embedding=self.config.use_average_appearance_embedding,
+ )
+ elif self.config.background_model == "mlp":
+ position_encoding = NeRFEncoding(
+ in_dim=3, num_frequencies=10, min_freq_exp=0.0, max_freq_exp=9.0, include_input=True
+ )
+ direction_encoding = NeRFEncoding(
+ in_dim=3, num_frequencies=4, min_freq_exp=0.0, max_freq_exp=3.0, include_input=True
+ )
+
+ self.field_background = NeRFField(
+ position_encoding=position_encoding,
+ direction_encoding=direction_encoding,
+ spatial_distortion=self.scene_contraction,
+ )
+ else:
+ # dummy background model
+ self.field_background = Parameter(torch.ones(1), requires_grad=False)
+
+ self.sampler_bg = LinearDisparitySampler(num_samples=self.config.num_samples_outside)
+
+ # renderers
+ background_color = (
+ get_color(self.config.background_color)
+ if self.config.background_color in set(["white", "black"])
+ else self.config.background_color
+ )
+ self.renderer_rgb = RGBRenderer(background_color=background_color)
+ self.renderer_accumulation = AccumulationRenderer()
+ self.renderer_depth = DepthRenderer(method="expected")
+ self.renderer_normal = SemanticRenderer()
+
+ # losses
+ self.rgb_loss = L1Loss()
+ self.eikonal_loss = MSELoss()
+ self.depth_loss = ScaleAndShiftInvariantLoss(alpha=0.5, scales=1)
+
+ # metrics
+ self.psnr = PeakSignalNoiseRatio(data_range=1.0)
+ self.ssim = structural_similarity_index_measure
+ self.lpips = LearnedPerceptualImagePatchSimilarity()
+
+ def get_param_groups(self) -> Dict[str, List[Parameter]]:
+ param_groups = {}
+ param_groups["fields"] = list(self.field.parameters())
+ if self.config.background_model != "none":
+ param_groups["field_background"] = list(self.field_background.parameters())
+ else:
+ param_groups["field_background"] = list(self.field_background)
+ return param_groups
+
+ @abstractmethod
+ def sample_and_forward_field(self, ray_bundle: RayBundle) -> Dict[str, torch.Tensor]:
+ """Takes in a Ray Bundle and returns a dictionary of samples and field output.
+
+ Args:
+ ray_bundle: Input bundle of rays. This raybundle should have all the
+ needed information to compute the outputs.
+
+ Returns:
+ Outputs of model. (ie. rendered colors)
+ """
+
+ def get_outputs(self, ray_bundle: RayBundle) -> Dict[str, torch.Tensor]:
+ """Takes in a Ray Bundle and returns a dictionary of outputs.
+
+ Args:
+ ray_bundle: Input bundle of rays. This raybundle should have all the
+ needed information to compute the outputs.
+
+ Returns:
+ Outputs of model. (ie. rendered colors)
+ """
+ samples_and_field_outputs = self.sample_and_forward_field(ray_bundle=ray_bundle)
+
+ # shortcuts
+ field_outputs = samples_and_field_outputs["field_outputs"]
+ ray_samples = samples_and_field_outputs["ray_samples"]
+ weights = samples_and_field_outputs["weights"]
+ bg_transmittance = samples_and_field_outputs["bg_transmittance"]
+
+ rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)
+ depth = self.renderer_depth(weights=weights, ray_samples=ray_samples)
+ # the rendered depth is point-to-point distance and we should convert to depth
+ depth = depth / ray_bundle.metadata["directions_norm"]
+
+ normal = self.renderer_normal(semantics=field_outputs[FieldHeadNames.NORMALS], weights=weights)
+ accumulation = self.renderer_accumulation(weights=weights)
+
+ # background model
+ if self.config.background_model != "none":
+ # sample inversely from far to 1000 and points and forward the bg model
+ ray_bundle.nears = ray_bundle.fars
+ ray_bundle.fars = torch.ones_like(ray_bundle.fars) * self.config.far_plane_bg
+
+ ray_samples_bg = self.sampler_bg(ray_bundle)
+ # use the same background model for both density field and occupancy field
+ field_outputs_bg = self.field_background(ray_samples_bg)
+ weights_bg = ray_samples_bg.get_weights(field_outputs_bg[FieldHeadNames.DENSITY])
+
+ rgb_bg = self.renderer_rgb(rgb=field_outputs_bg[FieldHeadNames.RGB], weights=weights_bg)
+ depth_bg = self.renderer_depth(weights=weights_bg, ray_samples=ray_samples_bg)
+ accumulation_bg = self.renderer_accumulation(weights=weights_bg)
+
+ # merge background color to forgound color
+ rgb = rgb + bg_transmittance * rgb_bg
+
+ bg_outputs = {
+ "bg_rgb": rgb_bg,
+ "bg_accumulation": accumulation_bg,
+ "bg_depth": depth_bg,
+ "bg_weights": weights_bg,
+ }
+ else:
+ bg_outputs = {}
+
+ outputs = {
+ "rgb": rgb,
+ "accumulation": accumulation,
+ "depth": depth,
+ "normal": normal,
+ "weights": weights,
+ # used to scale z_vals for free space and sdf loss
+ "directions_norm": ray_bundle.metadata["directions_norm"],
+ }
+ outputs.update(bg_outputs)
+
+ if self.training:
+ grad_points = field_outputs[FieldHeadNames.GRADIENT]
+ outputs.update({"eik_grad": grad_points})
+ outputs.update(samples_and_field_outputs)
+
+ if "weights_list" in samples_and_field_outputs:
+ weights_list = samples_and_field_outputs["weights_list"]
+ ray_samples_list = samples_and_field_outputs["ray_samples_list"]
+
+ for i in range(len(weights_list) - 1):
+ outputs[f"prop_depth_{i}"] = self.renderer_depth(
+ weights=weights_list[i], ray_samples=ray_samples_list[i]
+ )
+ # this is used only in viewer
+ outputs["normal_vis"] = (outputs["normal"] + 1.0) / 2.0
+ return outputs
+
+ def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]:
+ """Computes and returns the losses dict.
+
+ Args:
+ outputs: the output to compute loss dict to
+ batch: ground truth batch corresponding to outputs
+ metrics_dict: dictionary of metrics, some of which we can use for loss
+ """
+ loss_dict = {}
+ image = batch["image"].to(self.device)
+ loss_dict["rgb_loss"] = self.rgb_loss(image, outputs["rgb"])
+ if self.training:
+ # eikonal loss
+ grad_theta = outputs["eik_grad"]
+ loss_dict["eikonal_loss"] = ((grad_theta.norm(2, dim=-1) - 1) ** 2).mean() * self.config.eikonal_loss_mult
+
+ # foreground mask loss
+ if "fg_mask" in batch and self.config.fg_mask_loss_mult > 0.0:
+ fg_label = batch["fg_mask"].float().to(self.device)
+ weights_sum = outputs["weights"].sum(dim=1).clip(1e-3, 1.0 - 1e-3)
+ loss_dict["fg_mask_loss"] = (
+ F.binary_cross_entropy(weights_sum, fg_label) * self.config.fg_mask_loss_mult
+ )
+
+ # monocular normal loss
+ if "normal" in batch and self.config.mono_normal_loss_mult > 0.0:
+ normal_gt = batch["normal"].to(self.device)
+ normal_pred = outputs["normal"]
+ loss_dict["normal_loss"] = (
+ monosdf_normal_loss(normal_pred, normal_gt) * self.config.mono_normal_loss_mult
+ )
+
+ # monocular depth loss
+ if "depth" in batch and self.config.mono_depth_loss_mult > 0.0:
+ depth_gt = batch["depth"].to(self.device)[..., None]
+ depth_pred = outputs["depth"]
+
+ mask = torch.ones_like(depth_gt).reshape(1, 32, -1).bool()
+ loss_dict["depth_loss"] = (
+ self.depth_loss(depth_pred.reshape(1, 32, -1), (depth_gt * 50 + 0.5).reshape(1, 32, -1), mask)
+ * self.config.mono_depth_loss_mult
+ )
+
+ return loss_dict
+
+ def get_metrics_dict(self, outputs, batch) -> Dict[str, torch.Tensor]:
+ """Compute and returns metrics.
+
+ Args:
+ outputs: the output to compute loss dict to
+ batch: ground truth batch corresponding to outputs
+ """
+ metrics_dict = {}
+ image = batch["image"].to(self.device)
+ metrics_dict["psnr"] = self.psnr(outputs["rgb"], image)
+ return metrics_dict
+
+ def get_image_metrics_and_images(
+ self, outputs: Dict[str, torch.Tensor], batch: Dict[str, torch.Tensor]
+ ) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:
+ """Writes the test image outputs.
+ Args:
+ outputs: Outputs of the model.
+ batch: Batch of data.
+
+ Returns:
+ A dictionary of metrics.
+ """
+ image = batch["image"].to(self.device)
+ rgb = outputs["rgb"]
+ acc = colormaps.apply_colormap(outputs["accumulation"])
+
+ normal = outputs["normal"]
+ normal = (normal + 1.0) / 2.0
+
+ combined_rgb = torch.cat([image, rgb], dim=1)
+ combined_acc = torch.cat([acc], dim=1)
+ if "depth" in batch:
+ depth_gt = batch["depth"].to(self.device)
+ depth_pred = outputs["depth"]
+
+ # align to predicted depth and normalize
+ scale, shift = normalized_depth_scale_and_shift(
+ depth_pred[None, ..., 0], depth_gt[None, ...], depth_gt[None, ...] > 0.0
+ )
+ depth_pred = depth_pred * scale + shift
+
+ combined_depth = torch.cat([depth_gt[..., None], depth_pred], dim=1)
+ combined_depth = colormaps.apply_depth_colormap(combined_depth)
+ else:
+ depth = colormaps.apply_depth_colormap(
+ outputs["depth"],
+ accumulation=outputs["accumulation"],
+ )
+ combined_depth = torch.cat([depth], dim=1)
+
+ if "normal" in batch:
+ normal_gt = (batch["normal"].to(self.device) + 1.0) / 2.0
+ combined_normal = torch.cat([normal_gt, normal], dim=1)
+ else:
+ combined_normal = torch.cat([normal], dim=1)
+
+ images_dict = {
+ "img": combined_rgb,
+ "accumulation": combined_acc,
+ "depth": combined_depth,
+ "normal": combined_normal,
+ }
+
+ # Switch images from [H, W, C] to [1, C, H, W] for metrics computations
+ image = torch.moveaxis(image, -1, 0)[None, ...]
+ rgb = torch.moveaxis(rgb, -1, 0)[None, ...]
+
+ psnr = self.psnr(image, rgb)
+ ssim = self.ssim(image, rgb)
+ lpips = self.lpips(image, rgb)
+
+ # all of these metrics will be logged as scalars
+ metrics_dict = {"psnr": float(psnr.item()), "ssim": float(ssim)} # type: ignore
+ metrics_dict["lpips"] = float(lpips)
+
+ return metrics_dict, images_dict
diff --git a/nerfstudio/models/neus.py b/nerfstudio/models/neus.py
new file mode 100644
index 000000000..0e88b6843
--- /dev/null
+++ b/nerfstudio/models/neus.py
@@ -0,0 +1,119 @@
+# 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.
+
+"""
+Implementation of NeuS.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Dict, List, Type
+
+from nerfstudio.cameras.rays import RayBundle
+from nerfstudio.engine.callbacks import (
+ TrainingCallback,
+ TrainingCallbackAttributes,
+ TrainingCallbackLocation,
+)
+from nerfstudio.field_components.field_heads import FieldHeadNames
+from nerfstudio.model_components.ray_samplers import NeuSSampler
+from nerfstudio.models.base_surface_model import SurfaceModel, SurfaceModelConfig
+
+
+@dataclass
+class NeuSModelConfig(SurfaceModelConfig):
+ """NeuS Model Config"""
+
+ _target: Type = field(default_factory=lambda: NeuSModel)
+ num_samples: int = 64
+ """Number of uniform samples"""
+ num_samples_importance: int = 64
+ """Number of importance samples"""
+ num_up_sample_steps: int = 4
+ """number of up sample step, 1 for simple coarse-to-fine sampling"""
+ base_variance: float = 64
+ """fixed base variance in NeuS sampler, the inv_s will be base * 2 ** iter during upsample"""
+ perturb: bool = True
+ """use to use perturb for the sampled points"""
+
+
+class NeuSModel(SurfaceModel):
+ """NeuS model
+
+ Args:
+ config: NeuS configuration to instantiate model
+ """
+
+ config: NeuSModelConfig
+
+ def populate_modules(self):
+ """Set the fields and modules."""
+ super().populate_modules()
+
+ self.sampler = NeuSSampler(
+ num_samples=self.config.num_samples,
+ num_samples_importance=self.config.num_samples_importance,
+ num_samples_outside=self.config.num_samples_outside,
+ num_upsample_steps=self.config.num_up_sample_steps,
+ base_variance=self.config.base_variance,
+ )
+
+ self.anneal_end = 50000
+
+ def get_training_callbacks(
+ self, training_callback_attributes: TrainingCallbackAttributes
+ ) -> List[TrainingCallback]:
+ callbacks = []
+ # anneal for cos in NeuS
+ if self.anneal_end > 0:
+
+ def set_anneal(step):
+ anneal = min([1.0, step / self.anneal_end])
+ self.field.set_cos_anneal_ratio(anneal)
+
+ callbacks.append(
+ TrainingCallback(
+ where_to_run=[TrainingCallbackLocation.BEFORE_TRAIN_ITERATION],
+ update_every_num_iters=1,
+ func=set_anneal,
+ )
+ )
+
+ return callbacks
+
+ def sample_and_forward_field(self, ray_bundle: RayBundle) -> Dict:
+ ray_samples = self.sampler(ray_bundle, sdf_fn=self.field.get_sdf)
+ field_outputs = self.field(ray_samples, return_alphas=True)
+ weights, transmittance = ray_samples.get_weights_and_transmittance_from_alphas(
+ field_outputs[FieldHeadNames.ALPHA]
+ )
+ bg_transmittance = transmittance[:, -1, :]
+
+ samples_and_field_outputs = {
+ "ray_samples": ray_samples,
+ "field_outputs": field_outputs,
+ "weights": weights,
+ "bg_transmittance": bg_transmittance,
+ }
+ return samples_and_field_outputs
+
+ def get_metrics_dict(self, outputs, batch) -> Dict:
+ metrics_dict = super().get_metrics_dict(outputs, batch)
+ if self.training:
+ # training statics
+ metrics_dict["s_val"] = self.field.deviation_network.get_variance().item()
+ metrics_dict["inv_s"] = 1.0 / self.field.deviation_network.get_variance().item()
+
+ return metrics_dict
diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py
index 6db227971..22405f60c 100644
--- a/nerfstudio/pipelines/base_pipeline.py
+++ b/nerfstudio/pipelines/base_pipeline.py
@@ -109,9 +109,7 @@ class Pipeline(nn.Module):
return self.model.device
def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True):
- model_state = {
- key.replace("_model.", ""): value for key, value in state_dict.items() if key.startswith("_model.")
- }
+ model_state = {key[len("_model.") :]: value for key, value in state_dict.items() if key.startswith("_model.")}
pipeline_state = {key: value for key, value in state_dict.items() if not key.startswith("_model.")}
self._model.load_state_dict(model_state, strict=strict)
super().load_state_dict(pipeline_state, strict=False)
@@ -378,7 +376,9 @@ class VanillaPipeline(Pipeline):
loaded_state: pre-trained model state dict
step: training step of the loaded checkpoint
"""
- state = {key.replace("module.", ""): value for key, value in loaded_state.items()}
+ state = {
+ (key[len("module.") :] if key.startswith("module.") else key): value for key, value in loaded_state.items()
+ }
self._model.update_to_step(step)
self.load_state_dict(state, strict=True)
diff --git a/nerfstudio/process_data/colmap_utils.py b/nerfstudio/process_data/colmap_utils.py
index f8254f0a8..859222441 100644
--- a/nerfstudio/process_data/colmap_utils.py
+++ b/nerfstudio/process_data/colmap_utils.py
@@ -35,7 +35,7 @@ from nerfstudio.data.utils.colmap_parsing_utils import (
qvec2rotmat,
read_cameras_binary,
read_images_binary,
- read_points3d_binary,
+ read_points3D_binary,
)
from nerfstudio.process_data.process_data_utils import CameraModel
from nerfstudio.utils import colormaps
@@ -455,6 +455,11 @@ def colmap_to_json(
out = parse_colmap_camera_params(cam_id_to_camera[1])
out["frames"] = frames
+ applied_transform = np.eye(4)[:3, :]
+ applied_transform = applied_transform[np.array([1, 0, 2]), :]
+ applied_transform[2, :] *= -1
+ out["applied_transform"] = applied_transform.tolist()
+
with open(output_dir / "transforms.json", "w", encoding="utf-8") as f:
json.dump(out, f, indent=4)
@@ -510,7 +515,7 @@ def create_sfm_depth(
# ptid_to_info = recon.points3D
# cam_id_to_camera = recon.cameras
# im_id_to_image = recon.images
- ptid_to_info = read_points3d_binary(recon_dir / "points3D.bin")
+ ptid_to_info = read_points3D_binary(recon_dir / "points3D.bin")
cam_id_to_camera = read_cameras_binary(recon_dir / "cameras.bin")
im_id_to_image = read_images_binary(recon_dir / "images.bin")
diff --git a/nerfstudio/process_data/hloc_utils.py b/nerfstudio/process_data/hloc_utils.py
index 0f46dd12c..0aec2376e 100644
--- a/nerfstudio/process_data/hloc_utils.py
+++ b/nerfstudio/process_data/hloc_utils.py
@@ -42,6 +42,13 @@ except ImportError:
else:
_HAS_HLOC = True
+try:
+ from pixsfm.refine_hloc import PixSfM
+except ImportError:
+ _HAS_PIXSFM = False
+else:
+ _HAS_PIXSFM = True
+
CONSOLE = Console(width=120)
@@ -58,6 +65,7 @@ def run_hloc(
"superglue", "superglue-fast", "NN-superpoint", "NN-ratio", "NN-mutual", "adalam"
] = "superglue",
num_matched: int = 50,
+ refine_pixsfm: bool = False,
) -> None:
"""Runs hloc on the images.
@@ -75,6 +83,10 @@ def run_hloc(
)
sys.exit(1)
+ if refine_pixsfm and not _HAS_PIXSFM:
+ CONSOLE.print("[bold red]Error: use refine_pixsfm, you must install pixel-perfect-sfm toolbox!!")
+ sys.exit(1)
+
outputs = colmap_dir
sfm_pairs = outputs / "pairs-netvlad.txt"
sfm_dir = outputs / "sparse" / "0"
@@ -99,13 +111,29 @@ def run_hloc(
image_options = pycolmap.ImageReaderOptions( # pylint: disable=c-extension-no-member
camera_model=camera_model.value
)
- reconstruction.main(
- sfm_dir,
- image_dir,
- sfm_pairs,
- features,
- matches,
- camera_mode=pycolmap.CameraMode.SINGLE, # pylint: disable=c-extension-no-member
- image_options=image_options,
- verbose=verbose,
- )
+ if refine_pixsfm:
+ sfm = PixSfM({"dense_features": {"max_edge": 1024}})
+ refined, _ = sfm.reconstruction(
+ sfm_dir,
+ image_dir,
+ sfm_pairs,
+ features,
+ matches,
+ image_list=references,
+ camera_mode=pycolmap.CameraMode.SINGLE, # pylint: disable=c-extension-no-member
+ image_options=image_options,
+ verbose=verbose,
+ )
+ print("Refined", refined.summary())
+
+ else:
+ reconstruction.main(
+ sfm_dir,
+ image_dir,
+ sfm_pairs,
+ features,
+ matches,
+ camera_mode=pycolmap.CameraMode.SINGLE, # pylint: disable=c-extension-no-member
+ image_options=image_options,
+ verbose=verbose,
+ )
diff --git a/nerfstudio/process_data/process_data_utils.py b/nerfstudio/process_data/process_data_utils.py
index 14ae33662..a73127fcb 100644
--- a/nerfstudio/process_data/process_data_utils.py
+++ b/nerfstudio/process_data/process_data_utils.py
@@ -159,7 +159,7 @@ def convert_video_to_images(
ffmpeg_cmd += f" -vf thumbnail={spacing},setpts=N/TB{crop_cmd} -r 1"
else:
CONSOLE.print("[bold red]Can't satisfy requested number of frames. Extracting all frames.")
- ffmpeg_cmd += " -pix_fmt bgr8"
+ ffmpeg_cmd += f" -pix_fmt bgr8 -vf {crop_cmd[1:]}"
ffmpeg_cmd += f" {out_filename}"
run_command(ffmpeg_cmd, verbose=verbose)
@@ -174,13 +174,18 @@ def convert_video_to_images(
def copy_images_list(
- image_paths: List[Path], image_dir: Path, crop_border_pixels: Optional[int] = None, verbose: bool = False
+ image_paths: List[Path],
+ image_dir: Path,
+ crop_border_pixels: Optional[int] = None,
+ crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0),
+ verbose: bool = False,
) -> List[Path]:
"""Copy all images in a list of Paths. Useful for filtering from a directory.
Args:
image_paths: List of Paths of images to copy to a new directory.
image_dir: Path to the output directory.
crop_border_pixels: If not None, crops each edge by the specified number of pixels.
+ crop_factor: Portion of the image to crop. Should be in [0,1] (top, bottom, left, right)
verbose: If True, print extra logging.
Returns:
A list of the copied image Paths.
@@ -207,6 +212,18 @@ def copy_images_list(
crop = f"crop=iw-{crop_border_pixels*2}:ih-{crop_border_pixels*2}"
ffmpeg_cmd = f'ffmpeg -y -noautorotate -i "{image_dir / filename}" -q:v 2 -vf {crop} "{image_dir / filename}"'
run_command(ffmpeg_cmd, verbose=verbose)
+ elif crop_factor != (0.0, 0.0, 0.0, 0.0):
+ file_type = image_paths[0].suffix
+ filename = f"frame_%05d{file_type}"
+ height = 1 - crop_factor[0] - crop_factor[1]
+ width = 1 - crop_factor[2] - crop_factor[3]
+ start_x = crop_factor[2]
+ start_y = crop_factor[0]
+ crop_cmd = f',"crop=w=iw*{width}:h=ih*{height}:x=iw*{start_x}:y=ih*{start_y}"'
+ ffmpeg_cmd = (
+ f'ffmpeg -y -noautorotate -i "{image_dir / filename}" -q:v 2 -vf {crop_cmd[1:]} "{image_dir / filename}"'
+ )
+ run_command(ffmpeg_cmd, verbose=verbose)
num_frames = len(image_paths)
@@ -266,13 +283,16 @@ def copy_and_upscale_polycam_depth_maps_list(
return copied_depth_map_paths
-def copy_images(data: Path, image_dir: Path, verbose) -> OrderedDict[Path, Path]:
+def copy_images(
+ data: Path, image_dir: Path, verbose, crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0)
+) -> OrderedDict[Path, Path]:
"""Copy images from a directory to a new directory.
Args:
data: Path to the directory of images.
image_dir: Path to the output directory.
verbose: If True, print extra logging.
+ crop_factor: Portion of the image to crop. Should be in [0,1] (top, bottom, left, right)
Returns:
The mapping from the original filenames to the new ones.
"""
@@ -283,7 +303,9 @@ def copy_images(data: Path, image_dir: Path, verbose) -> OrderedDict[Path, Path]
CONSOLE.log("[bold red]:skull: No usable images in the data folder.")
sys.exit(1)
- copied_images = copy_images_list(image_paths=image_paths, image_dir=image_dir, verbose=verbose)
+ copied_images = copy_images_list(
+ image_paths=image_paths, image_dir=image_dir, crop_factor=crop_factor, verbose=verbose
+ )
return OrderedDict((original_path, new_path) for original_path, new_path in zip(image_paths, copied_images))
diff --git a/nerfstudio/utils/math.py b/nerfstudio/utils/math.py
index b0a0aaec4..f54b6509a 100644
--- a/nerfstudio/utils/math.py
+++ b/nerfstudio/utils/math.py
@@ -311,3 +311,44 @@ def masked_reduction(
input_tensor[valid] = input_tensor[valid] / mask[valid]
input_tensor = torch.mean(input_tensor)
return input_tensor
+
+
+def normalized_depth_scale_and_shift(
+ prediction: TensorType[1, 32, "mult"], target: TensorType[1, 32, "mult"], mask: TensorType[1, 32, "mult"]
+):
+ """
+ More info here: https://arxiv.org/pdf/2206.00665.pdf supplementary section A2 Depth Consistency Loss
+ This function computes scale/shift required to normalizes predicted depth map,
+ to allow for using normalized depth maps as input from monocular depth estimation networks.
+ These networks are trained such that they predict normalized depth maps.
+
+ Solves for scale/shift using a least squares approach with a closed form solution:
+ Based on:
+ https://github.com/autonomousvision/monosdf/blob/d9619e948bf3d85c6adec1a643f679e2e8e84d4b/code/model/loss.py#L7
+ Args:
+ prediction: predicted depth map
+ target: ground truth depth map
+ mask: mask of valid pixels
+ Returns:
+ scale and shift for depth prediction
+ """
+ # system matrix: A = [[a_00, a_01], [a_10, a_11]]
+ a_00 = torch.sum(mask * prediction * prediction, (1, 2))
+ a_01 = torch.sum(mask * prediction, (1, 2))
+ a_11 = torch.sum(mask, (1, 2))
+
+ # right hand side: b = [b_0, b_1]
+ b_0 = torch.sum(mask * prediction * target, (1, 2))
+ b_1 = torch.sum(mask * target, (1, 2))
+
+ # solution: x = A^-1 . b = [[a_11, -a_01], [-a_10, a_00]] / (a_00 * a_11 - a_01 * a_10) . b
+ scale = torch.zeros_like(b_0)
+ shift = torch.zeros_like(b_1)
+
+ det = a_00 * a_11 - a_01 * a_01
+ valid = det.nonzero()
+
+ scale[valid] = (a_11[valid] * b_0[valid] - a_01[valid] * b_1[valid]) / det[valid]
+ shift[valid] = (-a_01[valid] * b_0[valid] + a_00[valid] * b_1[valid]) / det[valid]
+
+ return scale, shift
diff --git a/pyproject.toml b/pyproject.toml
index 8ac4f04f2..e8dbf3e96 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "nerfstudio"
-version = "0.1.18"
+version = "0.1.19"
description = "All-in-one repository for state-of-the-art NeRFs"
readme = "README.md"
license = { text="Apache 2.0"}
diff --git a/scripts/process_data.py b/scripts/process_data.py
index 570fda59e..2b0647434 100644
--- a/scripts/process_data.py
+++ b/scripts/process_data.py
@@ -55,6 +55,8 @@ class ProcessImages:
sfm_tool: Literal["any", "colmap", "hloc"] = "any"
"""Structure from motion tool to use. Colmap will use sift features, hloc can use many modern methods
such as superpoint features and superglue matcher"""
+ refine_pixsfm: bool = False
+ """If True, runs refinement using Pixel Perfect SFM. Only works with hloc sfm_tool"""
feature_type: Literal[
"any",
"sift",
@@ -89,6 +91,8 @@ class ProcessImages:
"""Number of samples per image to take from each equirectangular image.
Used only when camera-type is equirectangular.
"""
+ crop_factor: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0)
+ """Portion of the image to crop. All values should be in [0,1]. (top, bottom, left, right)"""
gpu: bool = True
"""If True, use GPU."""
use_sfm_depth: bool = False
@@ -127,7 +131,7 @@ class ProcessImages:
pers_size = equirect_utils.compute_resolution_from_equirect(self.data, self.images_per_equirect)
CONSOLE.log(f"Generating {self.images_per_equirect} {pers_size} sized images per equirectangular image")
self.data = equirect_utils.generate_planar_projections_from_equirectangular(
- self.data, pers_size, self.images_per_equirect
+ self.data, pers_size, self.images_per_equirect, crop_factor=self.crop_factor
)
summary_log = []
@@ -136,7 +140,7 @@ class ProcessImages:
if not self.skip_image_processing:
# Copy images to output directory
image_rename_map_paths = process_data_utils.copy_images(
- self.data, image_dir=image_dir, verbose=self.verbose
+ self.data, image_dir=image_dir, crop_factor=self.crop_factor, verbose=self.verbose
)
image_rename_map = dict((a.name, b.name) for a, b in image_rename_map_paths.items())
num_frames = len(image_rename_map)
@@ -211,6 +215,9 @@ class ProcessImages:
(sfm_tool, feature_type, matcher_type) = process_data_utils.find_tool_feature_matcher_combination(
self.sfm_tool, self.feature_type, self.matcher_type
)
+ # check that sfm_tool is hloc if using refine_pixsfm
+ if self.refine_pixsfm:
+ assert sfm_tool == "hloc", "refine_pixsfm only works with sfm_tool hloc"
if sfm_tool == "colmap":
colmap_utils.run_colmap(
diff --git a/tests/dataparsers/test_nerfstudio_dataparser.py b/tests/dataparsers/test_nerfstudio_dataparser.py
new file mode 100644
index 000000000..c9dce5284
--- /dev/null
+++ b/tests/dataparsers/test_nerfstudio_dataparser.py
@@ -0,0 +1,99 @@
+"""
+Nerfstudio dataparser
+"""
+# pylint: disable=all
+import json
+from pathlib import Path
+
+import numpy as np
+import pytest
+from PIL import Image
+from pytest import fixture
+
+
+@fixture
+def mocked_dataset(tmp_path: Path):
+ """Mocked dataset with transforms"""
+ (tmp_path / "images_4").mkdir()
+ frames = []
+ for i in range(10):
+ Image.new("RGB", (100, 150)).save(tmp_path / "images_4" / f"img_{i}.png")
+ frames.append(
+ {
+ "file_path": f"img_{i}.png",
+ "transform_matrix": np.eye(4).tolist(),
+ }
+ )
+ with (tmp_path / "transforms.json").open("w+", encoding="utf8") as f:
+ json.dump({"fl_x": 2, "fl_y": 3, "cx": 4, "cy": 5, "h": 150, "w": 100, "frames": frames}, f)
+ return tmp_path
+
+
+@pytest.mark.parametrize("orientation_method", ["up", "none", "pca"])
+def test_nerfstudio_dataparser_no_filelist(mocked_dataset, orientation_method):
+ """Tests basic load"""
+ assert (mocked_dataset / "images_4").exists()
+ from nerfstudio.data.dataparsers.nerfstudio_dataparser import (
+ DataparserOutputs,
+ Nerfstudio,
+ NerfstudioDataParserConfig,
+ )
+
+ parser: Nerfstudio = NerfstudioDataParserConfig(
+ data=mocked_dataset,
+ downscale_factor=4,
+ orientation_method=orientation_method,
+ center_method="none",
+ auto_scale_poses=False,
+ ).setup()
+
+ paths = set()
+ for split in ("train", "test", "val"):
+ out = parser.get_dataparser_outputs(split)
+ assert isinstance(out, DataparserOutputs)
+ assert len(out.image_filenames) > 0
+ paths.update(out.image_filenames)
+ train_files = set(parser.get_dataparser_outputs("train").image_filenames)
+ assert len(train_files.intersection(parser.get_dataparser_outputs("val").image_filenames)) == 0
+ assert len(train_files.intersection(parser.get_dataparser_outputs("test").image_filenames)) == 0
+ assert len(paths) == 10
+
+
+def test_nerfstudio_dataparser_split_filelist(mocked_dataset):
+ """Tests basic load"""
+ assert (mocked_dataset / "images_4").exists()
+ with open(mocked_dataset / "transforms.json", "r+") as f:
+ data = json.load(f)
+ data["train_filenames"] = ["img_0.png", "img_1.png"]
+ data["val_filenames"] = ["img_2.png", "img_3.png"]
+ data["test_filenames"] = ["img_4.png", "img_5.png"]
+ f.seek(0)
+ f.truncate(0)
+ json.dump(data, f)
+
+ from nerfstudio.data.dataparsers.nerfstudio_dataparser import (
+ DataparserOutputs,
+ Nerfstudio,
+ NerfstudioDataParserConfig,
+ )
+
+ parser: Nerfstudio = NerfstudioDataParserConfig(
+ data=mocked_dataset,
+ downscale_factor=4,
+ orientation_method="none",
+ center_method="none",
+ auto_scale_poses=False,
+ ).setup()
+
+ assert parser.get_dataparser_outputs("train").image_filenames == [
+ mocked_dataset / "images_4/img_0.png",
+ mocked_dataset / "images_4/img_1.png",
+ ]
+ assert parser.get_dataparser_outputs("val").image_filenames == [
+ mocked_dataset / "images_4/img_2.png",
+ mocked_dataset / "images_4/img_3.png",
+ ]
+ assert parser.get_dataparser_outputs("test").image_filenames == [
+ mocked_dataset / "images_4/img_4.png",
+ mocked_dataset / "images_4/img_5.png",
+ ]
diff --git a/tests/process_data/test_process_images.py b/tests/process_data/test_process_images.py
new file mode 100644
index 000000000..3e4fa82a9
--- /dev/null
+++ b/tests/process_data/test_process_images.py
@@ -0,0 +1,80 @@
+"""
+Process images test
+"""
+import os
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image
+
+from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig
+from nerfstudio.data.utils.colmap_parsing_utils import Camera
+from nerfstudio.data.utils.colmap_parsing_utils import Image as ColmapImage
+from nerfstudio.data.utils.colmap_parsing_utils import (
+ write_cameras_binary,
+ write_images_binary,
+)
+from scripts.process_data import ProcessImages
+
+
+def test_process_images_skip_colmap(tmp_path: Path):
+ """
+ Test ns-process-data images
+ """
+ # Mock a colmap sparse model
+ width = 100
+ height = 150
+ sparse_path = tmp_path / "sparse" / "0"
+ sparse_path.mkdir(exist_ok=True, parents=True)
+ (tmp_path / "images").mkdir(exist_ok=True, parents=True)
+ write_cameras_binary(
+ {1: Camera(1, "OPENCV", width, height, [110, 110, 50, 75, 0, 0, 0, 0, 0, 0])},
+ sparse_path / "cameras.bin",
+ )
+ frames = {}
+ num_frames = 10
+ qvecs = np.random.uniform(size=(num_frames, 4))
+ tvecs = np.random.uniform(size=(num_frames, 3))
+ # original_poses = np.concatenate(
+ # (
+ # np.concatenate(
+ # (
+ # np.stack(list(map(qvec2rotmat, qvecs))),
+ # tvecs[:, :, None],
+ # ),
+ # -1,
+ # ),
+ # np.array([[[0, 0, 0, 1]]], dtype=qvecs.dtype).repeat(num_frames, 0),
+ # ),
+ # -2,
+ # )
+ for i in range(num_frames):
+ frames[i + 1] = ColmapImage(i + 1, qvecs[i], tvecs[i], 1, f"image_{i}.png", [], [])
+ Image.new("RGB", (width, height)).save(tmp_path / "images" / f"image_{i}.png")
+ write_images_binary(frames, sparse_path / "images.bin")
+
+ # Mock missing COLMAP and ffmpeg in the dev env
+ old_path = os.environ.get("PATH", "")
+ os.environ["PATH"] = str(tmp_path / "mocked_bin") + f":{old_path}"
+ (tmp_path / "mocked_bin").mkdir()
+ (tmp_path / "mocked_bin" / "colmap").touch(mode=0o777)
+ (tmp_path / "mocked_bin" / "ffmpeg").touch(mode=0o777)
+
+ # Run ProcessImages
+ cmd = ProcessImages(tmp_path / "images", tmp_path / "nerfstudio", colmap_model_path=sparse_path, skip_colmap=True)
+ cmd.main()
+ os.environ["PATH"] = old_path
+
+ assert (tmp_path / "nerfstudio" / "transforms.json").exists()
+ parser = NerfstudioDataParserConfig(
+ data=tmp_path / "nerfstudio",
+ downscale_factor=None,
+ orientation_method="none", # orientation_method,
+ center_method="none",
+ auto_scale_poses=False,
+ ).setup()
+ outputs = parser.get_dataparser_outputs("train")
+ assert len(outputs.image_filenames) == 9
+ assert torch.is_tensor(outputs.dataparser_transform)
+ # TODO @jkulhanek: Add tests if the loaded poses are the same
diff --git a/tests/test_train.py b/tests/test_train.py
index 1e5c55cd5..3f1d191d3 100644
--- a/tests/test_train.py
+++ b/tests/test_train.py
@@ -27,6 +27,7 @@ BLACKLIST = [
"nerfplayer-ngp",
"nerfplayer-nerfacto",
"dreamfusion",
+ "neus",
]