From 7ae2fce52b34c395a107a9ae43330191e4fb517d Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Fri, 10 Mar 2023 18:10:33 -0800 Subject: [PATCH 01/10] v0.1.19 (#1585) --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index a89f966d8..e75ddbf1c 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"} From 44ad94073d759fd6e6e755a6bb6343409e452377 Mon Sep 17 00:00:00 2001 From: Terrance Wang Date: Fri, 10 Mar 2023 23:09:46 -0800 Subject: [PATCH 02/10] bug fix for model loading from checkpoint (#1567) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * bug fix for model loading from checkpoint * removing prefix instead * fix in load_pipeline if ddp as well * fix * formatting --------- Co-authored-by: JonΓ‘Ε‘ KulhΓ‘nek --- nerfstudio/pipelines/base_pipeline.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index fd5d6fcc7..bb4a25f95 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -107,9 +107,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) @@ -370,7 +368,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) From e25c39e1c0f97f4752b9842a04d4f1290ddb8c01 Mon Sep 17 00:00:00 2001 From: Alexander Kristoffersen Date: Sat, 11 Mar 2023 07:53:59 -0800 Subject: [PATCH 03/10] Updated Colab to Python3.9 and added colmap install speedup. (#1587) --- colab/demo.ipynb | 70 ++++++++++++++++++++---------------------------- 1 file changed, 29 insertions(+), 41 deletions(-) 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" + "\"Open" ] }, { @@ -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" ] }, { From d9a391f79a7301a091148003b57d1ab86b5a808a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jon=C3=A1=C5=A1=20Kulh=C3=A1nek?= Date: Sat, 11 Mar 2023 22:14:45 +0100 Subject: [PATCH 04/10] Fixed: MinimalDataParser add options, allow Nerfstudio data format to specify test indices (#1586) * MinimalDataParser add options, allow Nerfstudio data format to specify test indices * Switch from indices to filenames and add tests * Fix #1581 * Fix changed interface --- nerfstudio/cameras/camera_utils.py | 2 +- .../data/dataparsers/minimal_dataparser.py | 18 +++- .../data/dataparsers/nerfstudio_dataparser.py | 53 +++++++--- nerfstudio/process_data/colmap_utils.py | 5 + .../dataparsers/test_nerfstudio_dataparser.py | 99 +++++++++++++++++++ 5 files changed, 158 insertions(+), 19 deletions(-) create mode 100644 tests/dataparsers/test_nerfstudio_dataparser.py 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/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 ef25653fd..a712d5005 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 = np.array(meta["applied_transform"], dtype=np.float32) + transform_matrix = transform_matrix @ np.concatenate( + (applied_transform, np.array([0, 0, 0, 1], dtype=np.float32)), 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/process_data/colmap_utils.py b/nerfstudio/process_data/colmap_utils.py index f8254f0a8..24517a331 100644 --- a/nerfstudio/process_data/colmap_utils.py +++ b/nerfstudio/process_data/colmap_utils.py @@ -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 + with open(output_dir / "transforms.json", "w", encoding="utf-8") as f: json.dump(out, f, indent=4) 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", + ] From 2b8d5a9e5c2dba4332a62434355631c13f6d9c96 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jon=C3=A1=C5=A1=20Kulh=C3=A1nek?= Date: Sun, 12 Mar 2023 00:12:14 +0100 Subject: [PATCH 05/10] Adds ns-process-data images tests and fixes #1586 (#1591) * Adds ns-process-data images tests and fixes #1586 * Fix missing colmap and ffmpeg in github actions --- .../data/dataparsers/nerfstudio_dataparser.py | 6 +- nerfstudio/data/utils/colmap_parsing_utils.py | 209 +++++++++++++++++- nerfstudio/process_data/colmap_utils.py | 6 +- tests/process_data/test_process_images.py | 80 +++++++ 4 files changed, 283 insertions(+), 18 deletions(-) create mode 100644 tests/process_data/test_process_images.py diff --git a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py index a712d5005..cffc09be1 100644 --- a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py +++ b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py @@ -289,9 +289,9 @@ class Nerfstudio(DataParser): cameras.rescale_output_resolution(scaling_factor=1.0 / self.downscale_factor) if "applied_transform" in meta: - applied_transform = np.array(meta["applied_transform"], dtype=np.float32) - transform_matrix = transform_matrix @ np.concatenate( - (applied_transform, np.array([0, 0, 0, 1], dtype=np.float32)), 0 + 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"]) 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/process_data/colmap_utils.py b/nerfstudio/process_data/colmap_utils.py index 24517a331..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 @@ -458,7 +458,7 @@ def colmap_to_json( applied_transform = np.eye(4)[:3, :] applied_transform = applied_transform[np.array([1, 0, 2]), :] applied_transform[2, :] *= -1 - out["applied_transform"] = applied_transform + out["applied_transform"] = applied_transform.tolist() with open(output_dir / "transforms.json", "w", encoding="utf-8") as f: json.dump(out, f, indent=4) @@ -515,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/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 From 1b18b44d814d51db061fbb91116fe8d30283e5df Mon Sep 17 00:00:00 2001 From: Matthew Tancik Date: Sun, 12 Mar 2023 11:14:58 -0700 Subject: [PATCH 06/10] Update 360 docs (#1593) * Update 360 docs * Update custom_dataset.md * Update custom_dataset.md --- README.md | 1 + docs/quickstart/custom_dataset.md | 13 +++++++++---- 2 files changed, 10 insertions(+), 4 deletions(-) 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/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. - From ddfa463f4d778f0ebe0db304628cf38981c48c10 Mon Sep 17 00:00:00 2001 From: Pablo Vela Date: Sun, 12 Mar 2023 16:01:18 -0400 Subject: [PATCH 07/10] NeuS Implementation (#1444) * add initial method configs and data manager * add neus scheduler * add initial surface model and neus * add Neus ray sampler * add sdf field * move scale/shift calc to math utils * add debug and fix sdf_dataset typo * update vscode neus debug * fix dataset and get initial working version * add magic number from original repo * black fixes * clean up get weights/transmittance from alphas * clean up typos in sdf datamanager * type annotations on neus scheduler and remove unused params * clean up and annotate sdf field along w/ other fixes * more clean up * remove unnessary directions norm * resolve comments * update neus scheduler to follow refactor * fix linter errors * fix some typos * remove get_weights_from_alpha * more comment fixes * more linter error fixes * clean up and add comments * more little fixes for types/comments * more comments * rename NeuS Scheduler to DelayedScheduler * rename DelayedScheduler to CosineDecayScheduler --- .vscode/launch.json | 21 +- nerfstudio/cameras/rays.py | 31 +- nerfstudio/configs/method_configs.py | 42 +- .../data/datamanagers/sdf_datamanager.py | 48 ++ .../data/dataparsers/sdfstudio_dataparser.py | 4 +- nerfstudio/data/datasets/sdf_dataset.py | 6 +- nerfstudio/engine/schedulers.py | 33 ++ nerfstudio/field_components/field_heads.py | 3 + nerfstudio/fields/base_field.py | 12 +- nerfstudio/fields/sdf_field.py | 458 ++++++++++++++++++ nerfstudio/model_components/losses.py | 48 +- nerfstudio/model_components/ray_samplers.py | 165 ++++++- nerfstudio/models/base_surface_model.py | 397 +++++++++++++++ nerfstudio/models/neus.py | 119 +++++ nerfstudio/utils/math.py | 41 ++ tests/test_train.py | 1 + 16 files changed, 1376 insertions(+), 53 deletions(-) create mode 100644 nerfstudio/data/datamanagers/sdf_datamanager.py create mode 100644 nerfstudio/fields/sdf_field.py create mode 100644 nerfstudio/models/base_surface_model.py create mode 100644 nerfstudio/models/neus.py 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/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 e888747b4..5d799fb01 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -27,6 +27,7 @@ from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig from nerfstudio.configs.base_config import ViewerConfig from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManagerConfig from nerfstudio.data.datamanagers.depth_datamanager import DepthDataManagerConfig +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, @@ -41,9 +42,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 @@ -52,6 +57,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 @@ -74,6 +80,7 @@ descriptions = { "phototourism": "Uses the Phototourism data.", "nerfplayer-nerfacto": "NeRFPlayer with nerfacto backbone.", "nerfplayer-ngp": "NeRFPlayer with InstantNGP backbone.", + "neus": "Implementation of NeuS. (slow)", } method_configs["nerfacto"] = TrainerConfig( @@ -430,6 +437,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/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/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 cba0ee7ea..6fffb5111 100644 --- a/nerfstudio/field_components/field_heads.py +++ b/nerfstudio/field_components/field_heads.py @@ -37,6 +37,9 @@ class FieldHeadNames(Enum): TRANSIENT_RGB = "transient_rgb" TRANSIENT_DENSITY = "transient_density" SEMANTICS = "semantics" + 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..3bcc51d56 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.""" 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 866218ba1..d75bfec94 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 @@ -598,3 +598,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/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/tests/test_train.py b/tests/test_train.py index 89aa1c29a..b38e84205 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -26,6 +26,7 @@ BLACKLIST = [ "depth-nerfacto", "nerfplayer-ngp", "nerfplayer-nerfacto", + "neus", ] From 630babacd6aea32999f40c181b7e8857d9a483b4 Mon Sep 17 00:00:00 2001 From: Pablo Vela Date: Sun, 12 Mar 2023 16:32:33 -0400 Subject: [PATCH 08/10] Add option to use Pixel Perfect SFM (#1554) add pixsfm option when using hloc --- .devcontainer/docker-compose.yml | 8 ++--- Dockerfile | 14 ++++++++ nerfstudio/process_data/hloc_utils.py | 48 +++++++++++++++++++++------ scripts/process_data.py | 5 +++ 4 files changed, 61 insertions(+), 14 deletions(-) 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/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/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/scripts/process_data.py b/scripts/process_data.py index 570fda59e..2e08bf9a9 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", @@ -211,6 +213,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( From de6af6c4e519a93d781fdf58dc97bedfb234a7ff Mon Sep 17 00:00:00 2001 From: Saurav Maheshkar Date: Tue, 14 Mar 2023 17:25:48 +0000 Subject: [PATCH 09/10] feat(fields): create common `get_normalized_directions` fn (#1599) * feat(fields): create common get_normalized_directions fn * style: run black * refactor: rename func and update docstring --- nerfstudio/fields/base_field.py | 9 +++++++++ nerfstudio/fields/instant_ngp_field.py | 13 ++----------- nerfstudio/fields/nerfacto_field.py | 13 ++----------- nerfstudio/fields/nerfplayer_nerfacto_field.py | 13 ++----------- nerfstudio/fields/nerfplayer_ngp_field.py | 5 ++--- 5 files changed, 17 insertions(+), 36 deletions(-) diff --git a/nerfstudio/fields/base_field.py b/nerfstudio/fields/base_field.py index 3bcc51d56..d8e4713d1 100644 --- a/nerfstudio/fields/base_field.py +++ b/nerfstudio/fields/base_field.py @@ -126,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: From 2838cfc8a57629778abbda9670798c7f1b08ecfe Mon Sep 17 00:00:00 2001 From: Rohan Mathur Date: Tue, 14 Mar 2023 12:04:28 -0700 Subject: [PATCH 10/10] Adding the crop code made for videos to images (#1598) * added equirect support * added cropping for planar image input * random comment * bug fix and running command * added crop-factor to docstrings --- nerfstudio/process_data/process_data_utils.py | 30 ++++++++++++++++--- scripts/process_data.py | 6 ++-- 2 files changed, 30 insertions(+), 6 deletions(-) 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/scripts/process_data.py b/scripts/process_data.py index 2e08bf9a9..2b0647434 100644 --- a/scripts/process_data.py +++ b/scripts/process_data.py @@ -91,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 @@ -129,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 = [] @@ -138,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)