diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..21f9e9035e7b055956d3fd84e0ccdec278374a3d 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +neuralangelo-main/assets/teaser.gif filter=lfs diff=lfs merge=lfs -text diff --git a/neuralangelo-main/.DS_Store b/neuralangelo-main/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..8fc4d3f4af301bbef4ca77b2380e53e17cdbe5c3 Binary files /dev/null and b/neuralangelo-main/.DS_Store differ diff --git a/neuralangelo-main/.gitignore b/neuralangelo-main/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..4e44ba4278be17d092ae3544a4c2061acb38b066 --- /dev/null +++ b/neuralangelo-main/.gitignore @@ -0,0 +1,199 @@ +checkpoints + +# Other uncheckable file types +*.zip +*.exe +*.dll +*.swp +*.vscode +*.ipynb +*.DS_Store +*.pyc + +# Credential information that should never be checked in +*.secret + +_datasets/ +scripts/wandb/_staff/ + +# Data types +*.png +*.hdr +*.jpg +*.jpeg +*.pgm +*.tiff +*.tif +*.mp4 +*.tar +*.tar.gz +*.pkl +*.pt +*.bin + +# log folder +logs/ + +# ------------------------ BELOW IS AUTO-GENERATED FOR PYTHON REPOS ------------------------ + +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/#use-with-ide +.pdm.toml + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ +CLIP diff --git a/neuralangelo-main/.gitmodules b/neuralangelo-main/.gitmodules new file mode 100644 index 0000000000000000000000000000000000000000..7cb07c0226faa7ff3af183e2a4647a7c7f370e82 --- /dev/null +++ b/neuralangelo-main/.gitmodules @@ -0,0 +1,3 @@ +[submodule "third_party/colmap"] + path = third_party/colmap + url = https://github.com/colmap/colmap.git diff --git a/neuralangelo-main/.pre-commit-config.yaml b/neuralangelo-main/.pre-commit-config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ef6f05697bcbcb2c418b9bf68637a26135cca3e5 --- /dev/null +++ b/neuralangelo-main/.pre-commit-config.yaml @@ -0,0 +1,7 @@ +repos: +- repo: https://github.com/pycqa/flake8 + rev: 4.0.0 + hooks: + - id: flake8 + args: [--max-line-length=120] + exclude: third_party diff --git a/neuralangelo-main/DATA_PROCESSING.md b/neuralangelo-main/DATA_PROCESSING.md new file mode 100644 index 0000000000000000000000000000000000000000..1b1004448adf40e7689d8d06416fa25a12b96198 --- /dev/null +++ b/neuralangelo-main/DATA_PROCESSING.md @@ -0,0 +1,126 @@ +# Data Preparation + +The following sections provide a step-by-step guide on how to convert a video to a json file that Neuralangelo parses. + +## Prerequisites +Initialize the COLMAP submodule: +```bash +git submodule update --init --recursive +``` + +## Self-captured video sequence +To capture your own data, we recommend using a high-shutter speed to avoid motion blur (which is very common when using a phone camera). To follow the instructions below, you can download a toy example from the link: https://drive.google.com/file/d/1VJeWYNJEBK0MFIzHjI8xZzwf3_I3eLwH/view?usp=drive_link + +### Preprocessing +You can run the following command to preprocess your data: + +```bash +EXPERIMENT_NAME=toy_example +PATH_TO_VIDEO=toy_example.MOV +SKIP_FRAME_RATE=24 +SCENE_TYPE=object # {outdoor,indoor,object} +bash projects/neuralangelo/scripts/preprocess.sh ${EXPERIMENT_NAME} ${PATH_TO_VIDEO} ${SKIP_FRAME_RATE} ${SCENE_TYPE} +``` + +Alternatively, you can follow the steps below if you want more fine-grained control. + +1. Convert video to images + + ```bash + PATH_TO_VIDEO=toy_example.MOV + SKIP_FRAME_RATE=30 + bash projects/neuralangelo/scripts/run_ffmpeg.sh ${PATH_TO_VIDEO} ${SKIP_FRAME_RATE} + ``` + `PATH_TO_VIDEO`: path to video + `SKIP_FRAME_RATE`: downsampling rate (recommended 10 for 24 fps captured videos) + +2. Run COLMAP + + ```bash + PATH_TO_IMAGES=toy_example_skip30 + bash projects/neuralangelo/scripts/run_colmap.sh ${PATH_TO_IMAGES} + ``` + `PATH_TO_IMAGES`: path to extracted images + + After COLMAP finishes, the folder structure will look like following: + ```bash + PATH_TO_IMAGES + |__ database.db (COLMAP databse) + |__ raw_images (raw input images) + |__ dense + |____ images (undistorted images) + |____ sparse (COLMAP correspondences, intrinsics and sparse point cloud) + |____ stereo (COLMAP files for MVS) + ``` + `dense/images` will be the input for surface reconstruction. + +3. Generate json file for data loading + + In this step, we define the bounding region for reconstruction and convert the COLMAP data to json format following Instant NGP. We strongly recommend you go to step 5 to validate the quality of the automatic bounding region extraction for improved performance. + + ```bash + PATH_TO_IMAGES=toy_example_skip30 + SCENE_TYPE=object # {outdoor,indoor,object} + python3 projects/neuralangelo/scripts/convert_data_to_json.py --data_dir ${PATH_TO_IMAGES}/dense --scene_type ${SCENE_TYPE} + ``` + `PATH_TO_IMAGES`: path to extracted images + +4. Config files + + Use the following to configure and generate your config files + ```bash + EXPERIMENT_NAME=toy_example + SCENE_TYPE=object # {outdoor,indoor,object} + python3 projects/neuralangelo/scripts/generate_config.py --experiment_name ${EXPERIMENT_NAME} --data_dir ${PATH_TO_IMAGES}/dense --scene_type ${SCENE_TYPE} --auto_exposure_wb + ``` + The config file will be generated as `projects/neuralangelo/configs/custom/{EXPERIMENT_NAME}.yaml`. + + To find more arguments and how they work: + ```bash + python3 projects/neuralangelo/scripts/generate_config.py -h + ``` + You can also manually adjust the parameters in the yaml file directly. + +5. Inspect results in Blender (optional but recommended) + + For certain cases, the camera poses estimated by COLMAP could be wrong, and the bounding region estimation could be off. + We offer some tools to to inspect the pre-processing results. Below are some options: + + - Blender: Download [Blender](https://www.blender.org/download/) and follow the instructions in our [add-on repo](https://github.com/mli0603/BlenderNeuralangelo). + - This [Jupyter notebook](projects/neuralangelo/scripts/visualize_colmap.ipynb) (using K3D) can be helpful for visualizing the COLMAP results. + +## DTU dataset +- Please use respecting the license terms of the dataset. + +You can run the following command to download [the DTU dataset](https://roboimagedata.compute.dtu.dk/?page_id=36) that is preprocessed by NeuS authors and generate json files: +```bash +PATH_TO_DTU=datasets/dtu # Modify this to be the DTU dataset root directory. +bash projects/neuralangelo/scripts/preprocess_dtu.sh ${PATH_TO_DTU} +``` + +## Tanks and Temples dataset +- Please use respecting the license terms of the dataset. + +Download the data from [Tanks and Temples](https://tanksandtemples.org/download/) website. +You will also need to download additional [COLMAP/camera/alignment](https://drive.google.com/file/d/1jAr3IDvhVmmYeDWi0D_JfgiHcl70rzVE/view?resourcekey=) and the images of each scene. +The file structure should look like (you may need to move around the downloaded images): +``` +tanks_and_temples +├─ Barn +│ ├─ Barn_COLMAP_SfM.log (camera poses) +│ ├─ Barn.json (cropfiles) +│ ├─ Barn.ply (ground-truth point cloud) +│ ├─ Barn_trans.txt (colmap-to-ground-truth transformation) +│ └─ images (folder of images) +│ ├─ 000001.png +│ ├─ 000002.png +│ ... +├─ Caterpillar +│ ├─ ... +... +``` +Run the following command to generate json files: +```bash +PATH_TO_TNT=datasets/tanks_and_temples # Modify this to be the Tanks and Temples root directory. +bash projects/neuralangelo/scripts/preprocess_tnt.sh ${PATH_TO_TNT} +``` diff --git a/neuralangelo-main/LICENSE.md b/neuralangelo-main/LICENSE.md new file mode 100644 index 0000000000000000000000000000000000000000..fb2dee8255000e3e9738ae146db21d5e41dcbeca --- /dev/null +++ b/neuralangelo-main/LICENSE.md @@ -0,0 +1,55 @@ +# NVIDIA Source Code License for Neuralangelo + +## 1. Definitions + +- “Licensor” means any person or entity that distributes its Work. + +- “Software” means the original work of authorship made available under this License. + +- “Work” means the Software and any additions to or derivative works of the Software that are made available under this License. + +- “NVIDIA Processors” means any central processing unit (CPU), graphics processing unit (GPU), field-programmable gate array (FPGA), application-specific integrated circuit (ASIC) or any combination thereof designed, made, sold, or provided by NVIDIA or its affiliates. + +- The terms “reproduce,” “reproduction,” “derivative works,” and “distribution” have the meaning as provided under U.S. copyright law; provided, however, that for the purposes of this License, derivative works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work. + +- Works, including the Software, are “made available” under this License by including in or with the Work either (a) a copyright notice referencing the applicability of this License to the Work, or (b) a copy of this License. + +## 2. License Grant + +### 2.1 Copyright Grant. + +Subject to the terms and conditions of this License, each Licensor grants to you a perpetual, worldwide, non-exclusive, royalty-free, copyright license to reproduce, prepare derivative works of, publicly display, publicly perform, sublicense and distribute its Work and any resulting derivative works in any form. + +## 3. Limitations + +### 3.1 Redistribution. + +You may reproduce or distribute the Work only if (a) you do so under this License, (b) you include a complete copy of this License with your distribution, and (c) you retain without modification any copyright, patent, trademark, or attribution notices that are present in the Work. + +### 3.2 Derivative Works. + +You may specify that additional or different terms apply to the use, reproduction, and distribution of your derivative works of the Work (“Your Terms”) only if (a) Your Terms provide that the use limitation in Section 3.3 applies to your derivative works, and (b) you identify the specific derivative works that are subject to Your Terms. Notwithstanding Your Terms, this License (including the redistribution requirements in Section 3.1) will continue to apply to the Work itself. + +### 3.3 Use Limitation. + +The Work and any derivative works thereof only may be used or intended for use non-commercially and with NVIDIA Processors. Notwithstanding the foregoing, NVIDIA and its affiliates may use the Work and any derivative works commercially. As used herein, “non-commercially” means for research or evaluation purposes only. + +### 3.4 Patent Claims. + +If you bring or threaten to bring a patent claim against any Licensor (including any claim, cross-claim or counterclaim in a lawsuit) to enforce any patents that you allege are infringed by any Work, then your rights under this License from such Licensor (including the grant in Section 2.1) will terminate immediately. + +### 3.5 Trademarks. + +This License does not grant any rights to use any Licensor’s or its affiliates’ names, logos, or trademarks, except as necessary to reproduce the notices described in this License. + +### 3.6 Termination. + +If you violate any term of this License, then your rights under this License (including the grant in Section 2.1) will terminate immediately. + +## 4. Disclaimer of Warranty. + +THE WORK IS PROVIDED “AS IS” WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF M ERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR NON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER THIS LICENSE. + +## 5. Limitation of Liability. + +EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE SHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF OR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK (INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION, LOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER COMM ERCIAL DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES. diff --git a/neuralangelo-main/README.md b/neuralangelo-main/README.md new file mode 100644 index 0000000000000000000000000000000000000000..aebd823169ba69ddf025f7cb1d73e1726429daac --- /dev/null +++ b/neuralangelo-main/README.md @@ -0,0 +1,79 @@ +# Neuralangelo + +## [Project Page](https://research.nvidia.com/labs/dir/neuralangelo/) | [Paper](https://arxiv.org/abs/2306.03092/) +This is the official repo for the implementation of **Neuralangelo: High-Fidelity Neural Surface Reconstruction**. +The code is built upon the Imaginaire library from the Deep Imagination Research Group at NVIDIA. + + + +For business inquiries, please submit the [NVIDIA research licensing form](https://www.nvidia.com/en-us/research/inquiries/). + +## Installation +We offer two ways to setup the environment: +1. We provide prebuilt Docker images, where + - `docker.io/chenhsuanlin/colmap:3.9` is for running COLMAP and the data preprocessing scripts. This includes the prebuilt COLMAP library (CUDA-supported). + - `docker.io/chenhsuanlin/neuralangelo:23.04-py3` is for running the main Neuralangelo pipeline. + + The corresponding Dockerfiles can be found in the `docker` directory. +2. The conda environment for Neuralangelo. Install the dependencies and activate the environment `neuralangelo` with + ```bash + conda env create --file neuralangelo.yaml + conda activate neuralangelo + ``` +For COLMAP, alternative installation options are also available on the [COLMAP website](https://colmap.github.io/). + +## Data preparation +Please refer to [Data Preparation](DATA_PROCESSING.md) for step-by-step instructions. +We assume known camera poses for each extracted frame from the video. +The code uses the same json format as [Instant NGP](https://github.com/NVlabs/instant-ngp). + +## Run Neuralangelo! +```bash +EXPERIMENT=toy_example +GROUP=example_group +NAME=example_name +CONFIG=projects/neuralangelo/configs/custom/${EXPERIMENT}.yaml +GPUS=1 # use >1 for multi-GPU training! +torchrun --nproc_per_node=${GPUS} train.py \ + --logdir=logs/${GROUP}/${NAME} \ + --config=${CONFIG} \ + --show_pbar +``` +Some useful notes: +- This codebase supports logging with [Weights & Biases](https://wandb.ai/site). You should have a W&B account for this. + - Add `--wandb` to the command line argument to enable W&B logging. + - Add `--wandb_name` to specify the W&B project name. + - More detailed control can be found in the `init_wandb()` function in `imaginaire/trainers/base.py`. +- Configs can be overridden through the command line (e.g. `--optim.params.lr=1e-2`). +- Set `--checkpoint={CHECKPOINT_PATH}` to initialize with a certain checkpoint; set `--resume=True` to resume training. +- If appearance embeddings are enabled, make sure `data.num_images` is set to the number of training images. + +## Isosurface extraction +Use the following command to run isosurface mesh extraction: +```bash +CHECKPOINT=logs/${GROUP}/${NAME}/xxx.pt +OUTPUT_MESH=xxx.ply +CONFIG=projects/neuralangelo/configs/custom/${EXPERIMENT}.yaml +RESOLUTION=2048 +BLOCK_RES=128 +GPUS=1 # use >1 for multi-GPU mesh extraction +torchrun --nproc_per_node=${GPUS} projects/neuralangelo/scripts/extract_mesh.py \ + --logdir=logs/${GROUP}/${NAME} \ + --config=${CONFIG} \ + --checkpoint=${CHECKPOINT} \ + --output_file=${OUTPUT_MESH} \ + --resolution=${RESOLUTION} \ + --block_res=${BLOCK_RES} +``` + +-------------------------------------- + +If you find our code useful for your research, please cite +``` +@inproceedings{li2023neuralangelo, + title={Neuralangelo: High-Fidelity Neural Surface Reconstruction}, + author={Li, Zhaoshuo and M\"uller, Thomas and Evans, Alex and Taylor, Russell H and Unberath, Mathias and Liu, Ming-Yu and Lin, Chen-Hsuan}, + booktitle={IEEE Conference on Computer Vision and Pattern Recognition ({CVPR})}, + year={2023} +} +``` diff --git a/neuralangelo-main/assets/teaser.gif b/neuralangelo-main/assets/teaser.gif new file mode 100644 index 0000000000000000000000000000000000000000..652415902db392b42280fc5c6de9827c8390a99e --- /dev/null +++ b/neuralangelo-main/assets/teaser.gif @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:34ac48ce705c8301950957e5f6089d0ba8adadfa41301d2a37837c045b5224f6 +size 15522649 diff --git a/neuralangelo-main/docker/Dockerfile-colmap b/neuralangelo-main/docker/Dockerfile-colmap new file mode 100644 index 0000000000000000000000000000000000000000..67c9c1aed1f6bafd0daa6c468940522662afc3c2 --- /dev/null +++ b/neuralangelo-main/docker/Dockerfile-colmap @@ -0,0 +1,47 @@ +# docker build -f docker/Dockerfile-colmap -t chenhsuanlin/colmap:3.9 . +# docker push chenhsuanlin/colmap:3.9 + +FROM nvcr.io/nvidia/cuda:11.8.0-devel-ubuntu20.04 +ARG DEBIAN_FRONTEND=noninteractive + +# colmap dependencies +RUN apt-get update && apt-get install -y \ + git \ + cmake \ + ninja-build \ + build-essential \ + libboost-program-options-dev \ + libboost-filesystem-dev \ + libboost-graph-dev \ + libboost-system-dev \ + libboost-test-dev \ + libeigen3-dev \ + libflann-dev \ + libfreeimage-dev \ + libmetis-dev \ + libgoogle-glog-dev \ + libgflags-dev \ + libsqlite3-dev \ + libglew-dev \ + qtbase5-dev \ + libqt5opengl5-dev \ + libcgal-dev \ + libceres-dev +# headless servers +RUN apt-get update && apt-get install -y \ + xvfb +# Colmap +RUN git clone https://github.com/colmap/colmap.git && cd colmap +RUN cd colmap && mkdir build && cd build && cmake .. -DCUDA_ENABLED=ON -DCMAKE_CUDA_ARCHITECTURES="70;72;75;80;86" -GNinja +RUN cd colmap/build && ninja && ninja install + +# additional python packages +RUN apt-get update && apt-get install -y \ + pip \ + ffmpeg +RUN pip install \ + addict \ + opencv-python-headless \ + pillow \ + pyyaml \ + trimesh diff --git a/neuralangelo-main/docker/Dockerfile-neuralangelo b/neuralangelo-main/docker/Dockerfile-neuralangelo new file mode 100644 index 0000000000000000000000000000000000000000..8e205eb7c6b33857d3acee7b83537b961bd6c366 --- /dev/null +++ b/neuralangelo-main/docker/Dockerfile-neuralangelo @@ -0,0 +1,33 @@ +# docker build -f docker/Dockerfile-neuralangelo -t chenhsuanlin/neuralangelo:23.04-py3 . +# docker push chenhsuanlin/neuralangelo:23.04-py3 + +FROM nvcr.io/nvidia/pytorch:23.04-py3 +ARG DEBIAN_FRONTEND=noninteractive + +# Install basics +RUN apt-get update && apt-get install -y --no-install-recommends \ + build-essential \ + bzip2 \ + ca-certificates \ + cmake \ + curl \ + ffmpeg \ + g++ \ + git \ + libx264-dev \ + tmux \ + wget + +# Update pip +RUN pip install --upgrade pip + +# Code formatting +RUN pip install --upgrade \ + flake8 \ + pre-commit + +# Install base Python libraries for Imaginaire +COPY requirements.txt requirements.txt +ARG FORCE_CUDA=1 +ARG TCNN_CUDA_ARCHITECTURES=70,72,75,80,86 +RUN pip install --upgrade -r requirements.txt diff --git a/neuralangelo-main/imaginaire/config.py b/neuralangelo-main/imaginaire/config.py new file mode 100644 index 0000000000000000000000000000000000000000..ffa0bb9057dd35348c1e42f326653c6d8a8674f8 --- /dev/null +++ b/neuralangelo-main/imaginaire/config.py @@ -0,0 +1,223 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import collections +import functools +import os +import re + +import yaml +from imaginaire.utils.distributed import master_only_print as print +from imaginaire.utils.termcolor import cyan, green, yellow + +DEBUG = False +USE_JIT = False + + +class AttrDict(dict): + """Dict as attribute trick.""" + + def __init__(self, *args, **kwargs): + super(AttrDict, self).__init__(*args, **kwargs) + self.__dict__ = self + for key, value in self.__dict__.items(): + if isinstance(value, dict): + self.__dict__[key] = AttrDict(value) + elif isinstance(value, (list, tuple)): + if value and isinstance(value[0], dict): + self.__dict__[key] = [AttrDict(item) for item in value] + else: + self.__dict__[key] = value + + def yaml(self): + """Convert object to yaml dict and return.""" + yaml_dict = {} + for key, value in self.__dict__.items(): + if isinstance(value, AttrDict): + yaml_dict[key] = value.yaml() + elif isinstance(value, list): + if value and isinstance(value[0], AttrDict): + new_l = [] + for item in value: + new_l.append(item.yaml()) + yaml_dict[key] = new_l + else: + yaml_dict[key] = value + else: + yaml_dict[key] = value + return yaml_dict + + def __repr__(self): + """Print all variables.""" + ret_str = [] + for key, value in self.__dict__.items(): + if isinstance(value, AttrDict): + ret_str.append('{}:'.format(key)) + child_ret_str = value.__repr__().split('\n') + for item in child_ret_str: + ret_str.append(' ' + item) + elif isinstance(value, list): + if value and isinstance(value[0], AttrDict): + ret_str.append('{}:'.format(key)) + for item in value: + # Treat as AttrDict above. + child_ret_str = item.__repr__().split('\n') + for item in child_ret_str: + ret_str.append(' ' + item) + else: + ret_str.append('{}: {}'.format(key, value)) + else: + ret_str.append('{}: {}'.format(key, value)) + return '\n'.join(ret_str) + + +class Config(AttrDict): + r"""Configuration class. This should include every human specifiable + hyperparameter values for your training.""" + + def __init__(self, filename=None, verbose=False): + super(Config, self).__init__() + self.source_filename = filename + + # Load the base configuration file. + base_filename = os.path.join( + os.path.dirname(__file__), '../imaginaire/config_base.yaml' + ) + cfg_base = self.load_config(base_filename) + recursive_update(self, cfg_base) + + # Update with given configurations. + cfg_dict = self.load_config(filename) + recursive_update(self, cfg_dict) + + if verbose: + print(' imaginaire config '.center(80, '-')) + print(self.__repr__()) + print(''.center(80, '-')) + + def load_config(self, filename): + # Update with given configurations. + assert os.path.exists(filename), f'File {filename} not exist.' + yaml_loader = yaml.SafeLoader + yaml_loader.add_implicit_resolver( + u'tag:yaml.org,2002:float', + re.compile(u'''^(?: + [-+]?(?:[0-9][0-9_]*)\\.[0-9_]*(?:[eE][-+]?[0-9]+)? + |[-+]?(?:[0-9][0-9_]*)(?:[eE][-+]?[0-9]+) + |\\.[0-9_]+(?:[eE][-+][0-9]+)? + |[-+]?[0-9][0-9_]*(?::[0-5]?[0-9])+\\.[0-9_]* + |[-+]?\\.(?:inf|Inf|INF) + |\\.(?:nan|NaN|NAN))$''', re.X), + list(u'-+0123456789.')) + try: + with open(filename) as file: + cfg_dict = yaml.load(file, Loader=yaml_loader) + cfg_dict = AttrDict(cfg_dict) + except EnvironmentError: + print(f'Please check the file with name of "{filename}"') + # Inherit configurations from parent + parent_key = "_parent_" + if parent_key in cfg_dict: + parent_filename = cfg_dict.pop(parent_key) + cfg_parent = self.load_config(parent_filename) + recursive_update(cfg_parent, cfg_dict) + cfg_dict = cfg_parent + return cfg_dict + + def print_config(self, level=0): + """Recursively print the configuration (with termcolor).""" + for key, value in sorted(self.items()): + if isinstance(value, (dict, Config)): + print(" " * level + cyan("* ") + green(key) + ":") + Config.print_config(value, level + 1) + else: + print(" " * level + cyan("* ") + green(key) + ":", yellow(value)) + + def save_config(self, logdir): + """Save the final configuration to a yaml file.""" + cfg_fname = f"{logdir}/config.yaml" + with open(cfg_fname, "w") as file: + yaml.safe_dump(self.yaml(), file, default_flow_style=False, indent=4) + + +def rsetattr(obj, attr, val): + """Recursively find object and set value""" + pre, _, post = attr.rpartition('.') + return setattr(rgetattr(obj, pre) if pre else obj, post, val) + + +def rgetattr(obj, attr, *args): + """Recursively find object and return value""" + + def _getattr(obj, attr): + r"""Get attribute.""" + return getattr(obj, attr, *args) + + return functools.reduce(_getattr, [obj] + attr.split('.')) + + +def recursive_update(d, u): + """Recursively update AttrDict d with AttrDict u""" + for key, value in u.items(): + if isinstance(value, collections.abc.Mapping): + d.__dict__[key] = recursive_update(d.get(key, AttrDict({})), value) + elif isinstance(value, (list, tuple)): + if value and isinstance(value[0], dict): + d.__dict__[key] = [AttrDict(item) for item in value] + else: + d.__dict__[key] = value + else: + d.__dict__[key] = value + return d + + +def recursive_update_strict(d, u, stack=[]): + """Recursively update AttrDict d with AttrDict u with strict matching""" + for key, value in u.items(): + if key not in d: + key_full = ".".join(stack + [key]) + raise KeyError(f"The input key '{key_full}; does not exist in the config files.") + if isinstance(value, collections.abc.Mapping): + d.__dict__[key] = recursive_update_strict(d.get(key, AttrDict({})), value, stack + [key]) + elif isinstance(value, (list, tuple)): + if value and isinstance(value[0], dict): + d.__dict__[key] = [AttrDict(item) for item in value] + else: + d.__dict__[key] = value + else: + d.__dict__[key] = value + return d + + +def parse_cmdline_arguments(args): + """ + Parse arguments from command line. + Syntax: --key1.key2.key3=value --> value + --key1.key2.key3= --> None + --key1.key2.key3 --> True + --key1.key2.key3! --> False + """ + cfg_cmd = {} + for arg in args: + assert arg.startswith("--") + if "=" not in arg[2:]: + key_str, value = (arg[2:-1], "false") if arg[-1] == "!" else (arg[2:], "true") + else: + key_str, value = arg[2:].split("=") + keys_sub = key_str.split(".") + cfg_sub = cfg_cmd + for k in keys_sub[:-1]: + cfg_sub.setdefault(k, {}) + cfg_sub = cfg_sub[k] + assert keys_sub[-1] not in cfg_sub, keys_sub[-1] + cfg_sub[keys_sub[-1]] = yaml.safe_load(value) + return cfg_cmd diff --git a/neuralangelo-main/imaginaire/config_base.yaml b/neuralangelo-main/imaginaire/config_base.yaml new file mode 100644 index 0000000000000000000000000000000000000000..47a3517272d07a19ecbc5aa5c7060952ade1724a --- /dev/null +++ b/neuralangelo-main/imaginaire/config_base.yaml @@ -0,0 +1,115 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. +# ----------------------------------------------------------------------------- + +# This is the base configuration file. + +# We often dump images to understand what's going on in the training. +# image_save_iter specifies how often we dump images. +image_save_iter: 9999999999 +# metrics_iter and metrics_epoch specify how often we compute the performance metrics +# If these two numbers are not set, they are copied from checkpoint.save_iter and checkpoint.save_epoch respectively. +metrics_iter: +metrics_epoch: +# max_epoch and max_iter specify what is the maximum epoch and iteration that we will train our model. +# min( max_epoch * dataset_size / batch_size, max_iter) will be the total number of iterations that the model will be trained. +max_epoch: 9999999999 +max_iter: 9999999999 +# logging_iter controls how often we log the training stats. +logging_iter: 100 +# If speed_benchmark is True, we will print out time required for forward, backward, and gradient update. +speed_benchmark: False +# Kill the process if `timeout_period` seconds have passed since the last iteration. This usually means the process gets stuck. +timeout_period: 9999999 + +# Default local rank +local_rank: 0 +# Toggle NVTX profiler +nvtx_profile: False + +# Checkpointer +checkpoint: + # If save_iter is set to M, then we save the checkpoint every M iteration. + # If save_latest_iter is set to M, then we save the checkpoint every M iteration using the name + # 'latest_checkpoint.pt', so that the new checkpoint will overwrite previous ones. + # If save_epoch is set to N, then we save the checkpoint every N epoch. + # Both can be set at the same time. + save_iter: 9999999999 + save_latest_iter: 9999999999 + save_epoch: 9999999999 + save_period: 9999999999 + # If True, load state_dict to the models in strict mode + strict_resume: True + +# Trainer +trainer: + ema_config: + enabled: False + beta: 0.9999 + start_iteration: 0 + + image_to_tensorboard: False + ddp_config: + find_unused_parameters: False + static_graph: True + init: + type: none + gain: + amp_config: + init_scale: 65536.0 + growth_factor: 2.0 + backoff_factor: 0.5 + growth_interval: 2000 + enabled: False + grad_accum_iter: 1 + +# Networks +model: + type: dummy + +# Optimizers +optim: + type: Adam + params: + # This defines the parameters for the specified PyTorch optimizer class (e.g. betas, eps). + lr: 0.0001 + fused_opt: False + # Default learning rate policy is step with iteration_mode=False (epoch mode), step_size=10^10, and gamma=1. + # This means a constant learning rate + sched: + iteration_mode: False + type: step + step_size: 9999999999 + gamma: 1 + +# Data +data: + name: dummy + type: imaginaire.datasets.images + use_multi_epoch_loader: False + num_workers: 0 +test_data: + name: dummy + type: imaginaire.datasets.images + num_workers: 0 + test: + is_lmdb: False + roots: + batch_size: 1 + +# cuDNN +# set deterministic to True for better reproducibility of the results. When deterministic is True, it will only use CUDNN functions that are deterministic. +# If benchmark is set to True, cudnn will benchmark several algorithms and pick that which it found to be fastest at the first iteration. +cudnn: + deterministic: False + benchmark: True + +# Others +pretrained_weight: +inference_args: {} diff --git a/neuralangelo-main/imaginaire/datasets/base.py b/neuralangelo-main/imaginaire/datasets/base.py new file mode 100644 index 0000000000000000000000000000000000000000..5827f8a5571f349abb96a032002a73ac2d16b35c --- /dev/null +++ b/neuralangelo-main/imaginaire/datasets/base.py @@ -0,0 +1,603 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import importlib +import json +import os +import pickle +from collections import OrderedDict +from functools import partial +from inspect import signature + +import numpy as np +import torch +import torch.utils.data as data +import torchvision.transforms as transforms + +from imaginaire.datasets.folder import FolderDataset +from imaginaire.datasets.lmdb import \ + IMG_EXTENSIONS, HDR_IMG_EXTENSIONS, LMDBDataset +from imaginaire.datasets.object_store import ObjectStoreDataset +from imaginaire.datasets.utils.data import \ + (VIDEO_EXTENSIONS, Augmentor, + load_from_folder, load_from_lmdb, load_from_object_store) +from imaginaire.datasets.utils.lmdb import create_metadata + + +DATASET_TYPES = ['lmdb', 'folder', 'object_store'] + + +class BaseDataset(data.Dataset): + r"""Base class for image/video datasets. + + Args: + cfg (Config object): Input config. + is_inference (bool): Training if False, else validation. + is_test (bool): Final test set after training and validation. + """ + + def __init__(self, cfg, is_inference, is_test): + super(BaseDataset, self).__init__() + + self.cfg = cfg + self.is_inference = is_inference + self.is_test = is_test + if self.is_test: + self.cfgdata = self.cfg.test_data + data_info = self.cfgdata.test + else: + self.cfgdata = self.cfg.data + if self.is_inference: + data_info = self.cfgdata.val + else: + data_info = self.cfgdata.train + self.name = self.cfgdata.name + self.lmdb_roots = data_info.roots + self.dataset_type = getattr(data_info, 'dataset_type', None) + self.cache = getattr(self.cfgdata, 'cache', None) + self.interpolator = getattr(self.cfgdata, 'interpolator', "INTER_LINEAR") + + # Get AWS secret keys. + if self.dataset_type == 'object_store': + assert hasattr(cfg, 'aws_credentials_file') + self.aws_credentials_file = cfg.aws_credentials_file + + # Legacy lmdb/folder only support. + if self.dataset_type is None: + self.dataset_is_lmdb = getattr(data_info, 'is_lmdb', False) + if self.dataset_is_lmdb: + self.dataset_type = 'lmdb' + else: + self.dataset_type = 'folder' + # Legacy support ends. + + assert self.dataset_type in DATASET_TYPES + if self.dataset_type == 'lmdb': + # Add handle to function to load data from LMDB. + self.load_from_dataset = load_from_lmdb + elif self.dataset_type == 'folder': + # For some unpaired experiments, we would like the dataset to be presented in a paired way + + if hasattr(self.cfgdata, 'paired') is False: + self.cfgdata.paired = self.paired + # Add handle to function to load data from folder. + self.load_from_dataset = load_from_folder + # Create metadata for folders. + print('Creating metadata') + all_filenames, all_metadata = [], [] + if self.is_test: + cfg.data_backup = cfg.data + cfg.data = cfg.test_data + for root in self.lmdb_roots: + filenames, metadata = create_metadata( + data_root=root, cfg=cfg, paired=self.cfgdata['paired']) + all_filenames.append(filenames) + all_metadata.append(metadata) + if self.is_test: + cfg.data = cfg.data_backup + elif self.dataset_type == 'object_store': + # Add handle to function to load data from AWS S3. + self.load_from_dataset = load_from_object_store + + # Get the types of data stored in dataset, and their extensions. + self.data_types = [] # Names of data types. + self.dataset_data_types = [] # These data types are in the dataset. + self.image_data_types = [] # These types are images. + self.hdr_image_data_types = [] # These types are HDR images. + self.normalize = {} # Does this data type need normalization? + self.extensions = {} # What is this data type's file extension. + self.is_mask = {} # Whether this data type is discrete masks? + self.num_channels = {} # How many channels does this data type have? + self.pre_aug_ops = {} # Ops on data type before augmentation. + self.post_aug_ops = {} # Ops on data type after augmentation. + + # Extract info from data types. + for data_type in self.cfgdata.input_types: + name = list(data_type.keys()) + assert len(name) == 1 + name = name[0] + info = data_type[name] + + if 'ext' not in info: + info['ext'] = None + if 'normalize' not in info: + info['normalize'] = False + if 'is_mask' not in info: + info['is_mask'] = False + if 'pre_aug_ops' not in info: + info['pre_aug_ops'] = 'None' + if 'post_aug_ops' not in info: + info['post_aug_ops'] = 'None' + if 'computed_on_the_fly' not in info: + info['computed_on_the_fly'] = False + if 'num_channels' not in info: + info['num_channels'] = None + + self.data_types.append(name) + if not info['computed_on_the_fly']: + self.dataset_data_types.append(name) + + self.extensions[name] = info['ext'] + self.normalize[name] = info['normalize'] + self.num_channels[name] = info['num_channels'] + self.pre_aug_ops[name] = [op.strip() for op in + info['pre_aug_ops'].split(',')] + self.post_aug_ops[name] = [op.strip() for op in + info['post_aug_ops'].split(',')] + self.is_mask[name] = info['is_mask'] + if info['ext'] is not None and (info['ext'] in IMG_EXTENSIONS or info['ext'] in VIDEO_EXTENSIONS): + self.image_data_types.append(name) + if info['ext'] is not None and info['ext'] in HDR_IMG_EXTENSIONS: + self.hdr_image_data_types.append(name) + + # Add some info into cfgdata for legacy support. + self.cfgdata.data_types = self.data_types + self.cfgdata.num_channels = [self.num_channels[name] + for name in self.data_types] + + # Augmentations which need full dict. + self.full_data_post_aug_ops, self.full_data_ops = [], [] + if hasattr(self.cfgdata, 'full_data_ops'): + ops = self.cfgdata.full_data_ops + self.full_data_ops.extend([op.strip() for op in ops.split(',')]) + if hasattr(self.cfgdata, 'full_data_post_aug_ops'): + ops = self.cfgdata.full_data_post_aug_ops + self.full_data_post_aug_ops.extend( + [op.strip() for op in ops.split(',')]) + + # These are the labels which will be concatenated for generator input. + self.input_labels = [] + if hasattr(self.cfgdata, 'input_labels'): + self.input_labels = self.cfgdata.input_labels + + # These are the keypoints which also need to be augmented. + self.keypoint_data_types = [] + if hasattr(self.cfgdata, 'keypoint_data_types'): + self.keypoint_data_types = self.cfgdata.keypoint_data_types + + # Create augmentation operations. + aug_list = data_info.augmentations + individual_video_frame_aug_list = getattr(data_info, 'individual_video_frame_augmentations', dict()) + post_aug_list = getattr(data_info, 'post_augmentations', dict()) + self.augmentor = Augmentor( + aug_list, individual_video_frame_aug_list, post_aug_list, self.image_data_types, self.is_mask, + self.keypoint_data_types, self.interpolator) + self.augmentable_types = self.image_data_types + \ + self.keypoint_data_types + + # Create torch transformations. + self.transform = {} + for data_type in self.image_data_types: + normalize = self.normalize[data_type] + self.transform[data_type] = self._get_transform( + normalize, self.num_channels[data_type]) + + # Create torch transformations for HDR images. + for data_type in self.hdr_image_data_types: + normalize = self.normalize[data_type] + self.transform[data_type] = self._get_transform( + normalize, self.num_channels[data_type]) + + # Initialize handles. + self.sequence_lists = [] # List of sequences per dataset root. + self.lmdbs = {} # Dict for list of lmdb handles per data type. + for data_type in self.dataset_data_types: + self.lmdbs[data_type] = [] + self.dataset_probability = None + self.additional_lists = [] + + # Load each dataset. + for idx, root in enumerate(self.lmdb_roots): + if self.dataset_type == 'lmdb': + self._add_dataset(root) + elif self.dataset_type == 'folder': + self._add_dataset(root, filenames=all_filenames[idx], + metadata=all_metadata[idx]) + elif self.dataset_type == 'object_store': + self._add_dataset( + root, aws_credentials_file=self.aws_credentials_file) + + # Compute dataset statistics and create whatever self.variables required + # for the specific dataloader. + self._compute_dataset_stats() + + # Build index of data to sample. + self.mapping, self.epoch_length = self._create_mapping() + + def _create_mapping(self): + r"""Creates mapping from data sample idx to actual LMDB keys. + All children need to implement their own. + + Returns: + self.mapping (list): List of LMDB keys. + """ + raise NotImplementedError + + def _compute_dataset_stats(self): + r"""Computes required statistics about dataset. + All children need to implement their own. + """ + pass + + def __getitem__(self, index): + r"""Entry function for dataset.""" + raise NotImplementedError + + def _get_transform(self, normalize, num_channels): + r"""Convert numpy to torch tensor. + + Args: + normalize (bool): Normalize image i.e. (x - 0.5) * 2. + Goes from [0, 1] -> [-1, 1]. + Returns: + Composed list of torch transforms. + """ + transform_list = [transforms.ToTensor()] + if normalize: + transform_list.append( + transforms.Normalize((0.5, ) * num_channels, + (0.5, ) * num_channels, inplace=True)) + return transforms.Compose(transform_list) + + def _add_dataset(self, root, filenames=None, metadata=None, + aws_credentials_file=None): + r"""Adds an LMDB dataset to a list of datasets. + + Args: + root (str): Path to LMDB or folder dataset. + filenames: List of filenames for folder dataset. + metadata: Metadata for folder dataset. + aws_credentials_file: Path to file containing AWS credentials. + """ + if aws_credentials_file and self.dataset_type == 'object_store': + object_store_dataset = ObjectStoreDataset( + root, aws_credentials_file, cache=self.cache) + sequence_list = object_store_dataset.sequence_list + else: + # Get sequences associated with this dataset. + if filenames is None: + list_path = 'all_filenames.json' + with open(os.path.join(root, list_path)) as fin: + sequence_list = OrderedDict(json.load(fin)) + else: + sequence_list = filenames + + additional_path = 'all_indices.json' + if os.path.exists(os.path.join(root, additional_path)): + print('Using additional list for object indices.') + with open(os.path.join(root, additional_path)) as fin: + additional_list = OrderedDict(json.load(fin)) + self.additional_lists.append(additional_list) + self.sequence_lists.append(sequence_list) + + # Get LMDB dataset handles. + for data_type in self.dataset_data_types: + if self.dataset_type == 'lmdb': + self.lmdbs[data_type].append( + LMDBDataset(os.path.join(root, data_type))) + elif self.dataset_type == 'folder': + self.lmdbs[data_type].append( + FolderDataset(os.path.join(root, data_type), metadata)) + elif self.dataset_type == 'object_store': + # All data types use the same handle. + self.lmdbs[data_type].append(object_store_dataset) + + def perform_individual_video_frame(self, data, augment_ops): + r"""Perform data augmentation on images only. + + Args: + data (dict): Keys are from data types. Values can be numpy.ndarray + or list of numpy.ndarray (image or list of images). + augment_ops (list): The augmentation operations for individual frames. + Returns: + (tuple): + - data (dict): Augmented data, with same keys as input data. + - is_flipped (bool): Flag which tells if images have been + left-right flipped. + """ + if augment_ops: + all_data = dict() + for ix, key in enumerate(data.keys()): + if ix == 0: + num = len(data[key]) + for j in range(num): + all_data['%d' % j] = dict() + for j in range(num): + all_data['%d' % j][key] = data[key][j:(j+1)] + for j in range(num): + all_data['%d' % j], _ = self.perform_augmentation( + all_data['%d' % j], paired=True, augment_ops=augment_ops) + for key in data.keys(): + tmp = [] + for j in range(num): + tmp += all_data['%d' % j][key] + data[key] = tmp + return data + + def perform_augmentation(self, data, paired, augment_ops=None): + r"""Perform data augmentation on images only. + + Args: + data (dict): Keys are from data types. Values can be numpy.ndarray + or list of numpy.ndarray (image or list of images). + paired (bool): Apply same augmentation to all input keys? + augment_ops (list): The augmentation operations. + Returns: + (tuple): + - data (dict): Augmented data, with same keys as input data. + - is_flipped (bool): Flag which tells if images have been + left-right flipped. + """ + aug_inputs = {} + for data_type in self.augmentable_types: + aug_inputs[data_type] = data[data_type] + + augmented, is_flipped = self.augmentor.perform_augmentation( + aug_inputs, paired=paired, augment_ops=augment_ops) + + for data_type in self.augmentable_types: + data[data_type] = augmented[data_type] + + return data, is_flipped + + def flip_hdr(self, data, is_flipped=False): + r"""Flip hdr images. + + Args: + data (dict): Keys are from data types. Values can be numpy.ndarray + or list of numpy.ndarray (image or list of images). + is_flipped (bool): Applying left-right flip to the hdr images + Returns: + (tuple): + - data (dict): Augmented data, with same keys as input data. + """ + if is_flipped is False: + return data + + for data_type in self.hdr_image_data_types: + # print('Length of data: {}'.format(len(data[data_type]))) + data[data_type][0] = data[data_type][0][:, ::-1, :].copy() + return data + + def to_tensor(self, data): + r"""Convert all images to tensor. + + Args: + data (dict): Dict containing data_type as key, with each value + as a list of numpy.ndarrays. + Returns: + data (dict): Dict containing data_type as key, with each value + as a list of torch.Tensors. + """ + for data_type in self.image_data_types: + for idx in range(len(data[data_type])): + if data[data_type][idx].dtype == np.uint16: + data[data_type][idx] = data[data_type][idx].astype( + np.float32) + data[data_type][idx] = self.transform[data_type]( + data[data_type][idx]) + for data_type in self.hdr_image_data_types: + for idx in range(len(data[data_type])): + data[data_type][idx] = self.transform[data_type]( + data[data_type][idx]) + return data + + def apply_ops(self, data, op_dict, full_data=False): + r"""Apply any ops from op_dict to data types. + + Args: + data (dict): Dict containing data_type as key, with each value + as a list of numpy.ndarrays. + op_dict (dict): Dict containing data_type as key, with each value + containing string of operations to apply. + full_data (bool): Do these ops require access to the full data? + Returns: + data (dict): Dict containing data_type as key, with each value + modified by the op if any. + """ + if full_data: + # op needs entire data dict. + for op in op_dict: + if op == 'None': + continue + op, op_type = self.get_op(op) + assert op_type == 'full_data' + data = op(data) + else: + # op per data type. + if not op_dict: + return data + for data_type in data: + for op in op_dict[data_type]: + if op == 'None': + continue + op, op_type = self.get_op(op) + data[data_type] = op(data[data_type]) + + if op_type == 'vis': + # We have converted this data type to an image. Enter it + # in self.image_data_types and give it a torch + # transform. + if data_type not in self.image_data_types: + self.image_data_types.append(data_type) + normalize = self.normalize[data_type] + num_channels = self.num_channels[data_type] + self.transform[data_type] = \ + self._get_transform(normalize, num_channels) + elif op_type == 'convert': + continue + elif op_type is None: + continue + else: + raise NotImplementedError + return data + + def get_op(self, op): + r"""Get function to apply for specific op. + + Args: + op (str): Name of the op. + Returns: + function handle. + """ + def list_to_tensor(data): + r"""Convert list of numeric values to tensor.""" + assert isinstance(data, list) + return torch.from_numpy(np.array(data, dtype=np.float32)) + + def decode_json_list(data): + r"""Decode list of strings in json to objects.""" + assert isinstance(data, list) + return [json.loads(item) for item in data] + + def decode_pkl_list(data): + r"""Decode list of pickled strings to objects.""" + assert isinstance(data, list) + return [pickle.loads(item) for item in data] + + def list_to_numpy(data): + r"""Convert list of numeric values to numpy array.""" + assert isinstance(data, list) + return np.array(data) + + def l2_normalize(data): + r"""L2 normalization.""" + assert isinstance(data, torch.Tensor) + import torch.nn.functional as F + return F.normalize(data, dim=1) + + if op == 'to_tensor': + return list_to_tensor, None + elif op == 'decode_json': + return decode_json_list, None + elif op == 'decode_pkl': + return decode_pkl_list, None + elif op == 'to_numpy': + return list_to_numpy, None + elif op == 'l2_norm': + return l2_normalize, None + elif '::' in op: + parts = op.split('::') + if len(parts) == 2: + module, function = parts + module = importlib.import_module(module) + function = getattr(module, function) + sig = signature(function) + num_params = len(sig.parameters) + assert num_params in [3, 4], \ + 'Full data functions take in (cfgdata, is_inference, ' \ + 'full_data) or (cfgdata, is_inference, self, full_data) ' \ + 'as input.' + if num_params == 3: + function = partial( + function, self.cfgdata, self.is_inference) + elif num_params == 4: + function = partial( + function, self.cfgdata, self.is_inference, self) + function_type = 'full_data' + elif len(parts) == 3: + function_type, module, function = parts + module = importlib.import_module(module) + + # Get function inputs, if provided. + partial_fn = False + if '(' in function and ')' in function: + partial_fn = True + function, params = self._get_fn_params(function) + + function = getattr(module, function) + + # Create partial function. + if partial_fn: + function = partial(function, **params) + + # Get function signature. + sig = signature(function) + num_params = 0 + for param in sig.parameters.values(): + if param.kind == param.POSITIONAL_OR_KEYWORD: + num_params += 1 + + if function_type == 'vis': + if num_params != 9: + raise ValueError( + 'vis function type needs to take ' + + '(resize_h, resize_w, crop_h, crop_w, ' + + 'original_h, original_w, is_flipped, cfgdata, ' + + 'data) as input.') + function = partial(function, + self.augmentor.resize_h, + self.augmentor.resize_w, + self.augmentor.crop_h, + self.augmentor.crop_w, + self.augmentor.original_h, + self.augmentor.original_w, + self.augmentor.is_flipped, + self.cfgdata) + elif function_type == 'convert': + if num_params != 1: + raise ValueError( + 'convert function type needs to take ' + + '(data) as input.') + else: + raise ValueError('Unknown op: %s' % (op)) + else: + raise ValueError('Unknown op: %s' % (op)) + return function, function_type + else: + raise ValueError('Unknown op: %s' % (op)) + + def _get_fn_params(self, function_string): + r"""Find key-value inputs to function from string definition. + + Args: + function_string (str): String with function name and args. e.g. + my_function(a=10, b=20). + Returns: + function (str): Name of function. + params (dict): Key-value params for function. + """ + start = function_string.find('(') + end = function_string.find(')') + function = function_string[:start] + params_str = function_string[start+1:end] + params = {} + for item in params_str.split(':'): + key, value = item.split('=') + try: + params[key] = float(value) + except Exception: + params[key] = value + return function, params + + def __len__(self): + return self.epoch_length diff --git a/neuralangelo-main/imaginaire/datasets/utils/dataloader.py b/neuralangelo-main/imaginaire/datasets/utils/dataloader.py new file mode 100644 index 0000000000000000000000000000000000000000..f4411bc308e5003c5e2f9c75709d7cc563942d7f --- /dev/null +++ b/neuralangelo-main/imaginaire/datasets/utils/dataloader.py @@ -0,0 +1,49 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import torch + + +class MultiEpochsDataLoader(torch.utils.data.DataLoader): + """ + Relentlessly sample from the dataset. + This eliminates the overhead of prefetching data before each epoch. + https://github.com/rwightman/pytorch-image-models/blob/master/timm/data/loader.py + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._DataLoader__initialized = False + self.batch_sampler = _RepeatSampler(self.batch_sampler) + self._DataLoader__initialized = True + self.iterator = super().__iter__() + + def __len__(self): + return len(self.batch_sampler.sampler) + + def __iter__(self): + for i in range(len(self)): + yield next(self.iterator) + + +class _RepeatSampler(object): + """ Sampler that repeats forever. + Args: + sampler (Sampler) + """ + + def __init__(self, sampler): + self.sampler = sampler + + def __iter__(self): + while True: + yield from iter(self.sampler) diff --git a/neuralangelo-main/imaginaire/datasets/utils/get_dataloader.py b/neuralangelo-main/imaginaire/datasets/utils/get_dataloader.py new file mode 100644 index 0000000000000000000000000000000000000000..0b804e85ccc5d3b5d375ca291a22a498784f64c5 --- /dev/null +++ b/neuralangelo-main/imaginaire/datasets/utils/get_dataloader.py @@ -0,0 +1,181 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import importlib + +import torch +import torch.distributed as dist + +from imaginaire.utils.distributed import master_only_print as print + +from imaginaire.datasets.utils.sampler import DistributedSamplerPreemptable +from imaginaire.datasets.utils.dataloader import MultiEpochsDataLoader + + +def _get_train_dataset_objects(cfg, subset_indices=None): + r"""Return dataset objects for the training set. + Args: + cfg (obj): Global configuration file. + subset_indices (sequence): Indices of the subset to use. + + Returns: + train_dataset (obj): PyTorch training dataset object. + """ + dataset_module = importlib.import_module(cfg.data.type) + train_dataset = dataset_module.Dataset(cfg, is_inference=False) + if subset_indices is not None: + train_dataset = torch.utils.data.Subset(train_dataset, subset_indices) + print('Train dataset length:', len(train_dataset)) + return train_dataset + + +def _get_val_dataset_objects(cfg, subset_indices=None): + r"""Return dataset objects for the validation set. + Args: + cfg (obj): Global configuration file. + subset_indices (sequence): Indices of the subset to use. + Returns: + val_dataset (obj): PyTorch validation dataset object. + """ + dataset_module = importlib.import_module(cfg.data.type) + if hasattr(cfg.data.val, 'type'): + for key in ['type', 'input_types', 'input_image']: + setattr(cfg.data, key, getattr(cfg.data.val, key)) + dataset_module = importlib.import_module(cfg.data.type) + val_dataset = dataset_module.Dataset(cfg, is_inference=True) + + if subset_indices is not None: + val_dataset = torch.utils.data.Subset(val_dataset, subset_indices) + print('Val dataset length:', len(val_dataset)) + return val_dataset + + +def _get_test_dataset_object(cfg, subset_indices=None): + r"""Return dataset object for the test set + + Args: + cfg (obj): Global configuration file. + subset_indices (sequence): Indices of the subset to use. + Returns: + (obj): PyTorch dataset object. + """ + dataset_module = importlib.import_module(cfg.test_data.type) + test_dataset = dataset_module.Dataset(cfg, is_inference=True, is_test=True) + if subset_indices is not None: + test_dataset = torch.utils.data.Subset(test_dataset, subset_indices) + return test_dataset + + +def _get_data_loader(cfg, dataset, batch_size, not_distributed=False, + shuffle=True, drop_last=True, seed=0, use_multi_epoch_loader=False, + preemptable=False): + r"""Return data loader . + + Args: + cfg (obj): Global configuration file. + dataset (obj): PyTorch dataset object. + batch_size (int): Batch size. + not_distributed (bool): Do not use distributed samplers. + shuffle (bool): Whether to shuffle the data + drop_last (bool): Whether to drop the last batch is the number of samples is smaller than the batch size + seed (int): random seed. + preemptable (bool): Whether to handle preemptions. + Return: + (obj): Data loader. + """ + not_distributed = not_distributed or not dist.is_initialized() + if not_distributed: + sampler = None + else: + if preemptable: + sampler = DistributedSamplerPreemptable(dataset, shuffle=shuffle, seed=seed) + else: + sampler = torch.utils.data.distributed.DistributedSampler(dataset, shuffle=shuffle, seed=seed) + num_workers = getattr(cfg.data, 'num_workers', 8) + persistent_workers = getattr(cfg.data, 'persistent_workers', False) + data_loader = (MultiEpochsDataLoader if use_multi_epoch_loader else torch.utils.data.DataLoader)( + dataset, + batch_size=batch_size, + shuffle=shuffle and (sampler is None), + sampler=sampler, + pin_memory=True, + num_workers=num_workers, + drop_last=drop_last, + persistent_workers=persistent_workers if num_workers > 0 else False + ) + return data_loader + + +def get_train_dataloader( + cfg, shuffle=True, drop_last=True, subset_indices=None, seed=0, preemptable=False): + r"""Return dataset objects for the training and validation sets. + Args: + cfg (obj): Global configuration file. + shuffle (bool): Whether to shuffle the data + drop_last (bool): Whether to drop the last batch is the number of samples is smaller than the batch size + subset_indices (sequence): Indices of the subset to use. + seed (int): random seed. + preemptable (bool): Flag for preemption handling + Returns: + train_data_loader (obj): Train data loader. + """ + train_dataset = _get_train_dataset_objects(cfg, subset_indices=subset_indices) + train_data_loader = _get_data_loader( + cfg, train_dataset, cfg.data.train.batch_size, not_distributed=False, + shuffle=shuffle, drop_last=drop_last, seed=seed, + use_multi_epoch_loader=cfg.data.use_multi_epoch_loader, + preemptable=preemptable + ) + return train_data_loader + + +def get_val_dataloader(cfg, subset_indices=None, seed=0): + r"""Return dataset objects for the training and validation sets. + Args: + cfg (obj): Global configuration file. + subset_indices (sequence): Indices of the subset to use. + seed (int): random seed. + Returns: + val_data_loader (obj): Val data loader. + """ + val_dataset = _get_val_dataset_objects(cfg, subset_indices=subset_indices) + not_distributed = getattr(cfg.data, 'val_data_loader_not_distributed', False) + # We often use a folder of images to represent a video. As doing evaluation, we like the images to preserve the + # original order. As a result, we do not want to distribute images from the same video to different GPUs. + not_distributed = 'video' in cfg.data.type or not_distributed + drop_last = getattr(cfg.data.val, 'drop_last', False) + # Validation loader need not have preemption handling. + val_data_loader = _get_data_loader( + cfg, val_dataset, cfg.data.val.batch_size, not_distributed=not_distributed, + shuffle=False, drop_last=drop_last, seed=seed, + preemptable=False + ) + return val_data_loader + + +def get_test_dataloader(cfg, subset_indices=None): + r"""Return dataset objects for testing + + Args: + cfg (obj): Global configuration file. + subset_indices (sequence): Indices of the subset to use. + Returns: + (obj): Test data loader. It may not contain the ground truth. + """ + test_dataset = _get_test_dataset_object(cfg, subset_indices=subset_indices) + not_distributed = getattr( + cfg.test_data, 'val_data_loader_not_distributed', False) + not_distributed = 'video' in cfg.test_data.type or not_distributed + test_data_loader = _get_data_loader( + cfg, test_dataset, cfg.test_data.test.batch_size, not_distributed=not_distributed, + shuffle=False) + return test_data_loader diff --git a/neuralangelo-main/imaginaire/datasets/utils/sampler.py b/neuralangelo-main/imaginaire/datasets/utils/sampler.py new file mode 100644 index 0000000000000000000000000000000000000000..87049b68aecb8cc9810c872d94c65ede403e5e22 --- /dev/null +++ b/neuralangelo-main/imaginaire/datasets/utils/sampler.py @@ -0,0 +1,117 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import math +import torch.distributed as dist +import torch + +from torch.utils.data import Sampler +from typing import TypeVar + +T_co = TypeVar('T_co', covariant=True) + + +class DistributedSamplerPreemptable(Sampler[T_co]): + r"""Sampler that supports loading from an iteration. + This is very useful for preemptable jobs. + + Args: + dataset (torch.utils.data.Dataset): Dataset object + num_replicas (int): Number of replicas to the distribute the dataloader over. + This is typically the world size in DDP jobs. + rank (int): Rank of the current process. + shuffle (bool): Whether to shuffle the dataloader in each epoch. + seed (int): Random seed used for shuffling the dataloader. + drop_last (bool): Whether to drop the last batch. + """ + + def __init__(self, dataset, num_replicas=None, rank=None, shuffle=True, + seed=0, drop_last=False): + + if num_replicas is None: + if not dist.is_available(): + raise RuntimeError("Requires distributed package to be available") + num_replicas = dist.get_world_size() + if rank is None: + if not dist.is_available(): + raise RuntimeError("Requires distributed package to be available") + rank = dist.get_rank() + if rank >= num_replicas or rank < 0: + raise ValueError( + "Invalid rank {}, rank should be in the interval" + " [0, {}]".format(rank, num_replicas - 1)) + self.dataset = dataset + self.num_replicas = num_replicas + self.rank = rank + self.epoch = 0 + + # start_index is the index to begin the dataloader from. + self.start_index = 0 + + self.drop_last = drop_last + # If the dataset length is evenly divisible by # of replicas, then there + # is no need to drop any data, since the dataset will be split equally. + if self.drop_last and len(self.dataset) % self.num_replicas != 0: # type: ignore[arg-type] + # Split to nearest available length that is evenly divisible. + # This is to ensure each rank receives the same amount of data when + # using this Sampler. + self.num_samples = math.ceil( + (len(self.dataset) - self.num_replicas) / self.num_replicas # type: ignore[arg-type] + ) + else: + self.num_samples = math.ceil(len(self.dataset) / self.num_replicas) # type: ignore[arg-type] + self.total_size = self.num_samples * self.num_replicas + self.shuffle = shuffle + self.seed = seed + + def __iter__(self): + if self.shuffle: + # deterministically shuffle based on epoch and seed + g = torch.Model() + g.manual_seed(self.seed + self.epoch) + indices = torch.randperm(len(self.dataset), generator=g).tolist() # type: ignore[arg-type] + else: + indices = list(range(len(self.dataset))) # type: ignore[arg-type] + + if not self.drop_last: + # add extra samples to make it evenly divisible + padding_size = self.total_size - len(indices) + if padding_size <= len(indices): + indices += indices[:padding_size] + else: + indices += (indices * math.ceil(padding_size / len(indices)))[:padding_size] + else: + # remove tail of data to make it evenly divisible. + indices = indices[:self.total_size] + assert len(indices) == self.total_size + + # subsample + indices = indices[self.rank:self.total_size:self.num_replicas] + assert len(indices) == self.num_samples + + # assert self.start_index < len(indices) + if self.start_index >= len(indices): + print('(Warning): Start index is less than len of dataloader. Goint to the last batch of dataset instead') + # This is hardcoded to go one batch before. + self.start_index = len(indices) - 64 + indices = indices[self.start_index:] + + return iter(indices) + + def __len__(self): + return self.num_samples + + def set_epoch(self, epoch): + self.epoch = epoch + + def set_iteration(self, start_index): + self.start_index = start_index diff --git a/neuralangelo-main/imaginaire/models/base.py b/neuralangelo-main/imaginaire/models/base.py new file mode 100644 index 0000000000000000000000000000000000000000..2f003db93365c0d681a0eba8de9a922e5cecf435 --- /dev/null +++ b/neuralangelo-main/imaginaire/models/base.py @@ -0,0 +1,30 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import torch + + +class Model(torch.nn.Module): + + def __init__(self, cfg_model, cfg_data): + super().__init__() + + def get_param_groups(self, cfg_optim): + """Allow the network to use different hyperparameters (e.g., learning rate) for different parameters. + Returns: + PyTorch parameter group (list or generator). See the PyTorch documentation for details. + """ + return self.parameters() + + def device(self): + """Return device on which model resides.""" + return next(self.parameters()).device diff --git a/neuralangelo-main/imaginaire/models/utils/init_weight.py b/neuralangelo-main/imaginaire/models/utils/init_weight.py new file mode 100644 index 0000000000000000000000000000000000000000..bcbcbba477bdea1677a872c7e2f5e2c792cbc934 --- /dev/null +++ b/neuralangelo-main/imaginaire/models/utils/init_weight.py @@ -0,0 +1,91 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import torch +from torch.nn import init + + +def weights_init(init_type, gain, bias=None): + r"""Initialize weights in the network. + + Args: + init_type (str): The name of the initialization scheme. + gain (float): The parameter that is required for the initialization + scheme. + bias (object): If not ``None``, specifies the initialization parameter + for bias. + + Returns: + (obj): init function to be applied. + """ + + def init_func(m): + r"""Init function + + Args: + m: module to be weight initialized. + """ + class_name = m.__class__.__name__ + if hasattr(m, 'weight') and ( + class_name.find('Conv') != -1 or + class_name.find('Linear') != -1 or + class_name.find('Embedding') != -1): + lr_mul = getattr(m, 'lr_mul', 1.) + gain_final = gain / lr_mul + if init_type == 'normal': + init.normal_(m.weight.data, 0.0, gain_final) + elif init_type == 'xavier': + init.xavier_normal_(m.weight.data, gain=gain_final) + elif init_type == 'xavier_uniform': + init.xavier_uniform_(m.weight.data, gain=gain_final) + elif init_type == 'kaiming': + init.kaiming_normal_(m.weight.data, a=0, mode='fan_in') + with torch.no_grad(): + m.weight.data *= gain_final + elif init_type == 'kaiming_linear': + init.kaiming_normal_( + m.weight.data, a=0, mode='fan_in', nonlinearity='linear' + ) + with torch.no_grad(): + m.weight.data *= gain_final + elif init_type == 'orthogonal': + init.orthogonal_(m.weight.data, gain=gain_final) + elif init_type == 'none': + pass + else: + raise NotImplementedError( + 'initialization method [%s] is ' + 'not implemented' % init_type) + if hasattr(m, 'bias') and m.bias is not None: + if init_type == 'none': + pass + elif bias is not None: + bias_type = getattr(bias, 'type', 'normal') + if bias_type == 'normal': + bias_gain = getattr(bias, 'gain', 0.5) + init.normal_(m.bias.data, 0.0, bias_gain) + else: + raise NotImplementedError( + 'initialization method [%s] is ' + 'not implemented' % bias_type) + else: + init.constant_(m.bias.data, 0.0) + return init_func + + +def weights_rescale(): + def init_func(m): + if hasattr(m, 'init_gain'): + for name, p in m.named_parameters(): + if 'output_scale' not in name: + p.data.mul_(m.init_gain) + return init_func diff --git a/neuralangelo-main/imaginaire/models/utils/model_average.py b/neuralangelo-main/imaginaire/models/utils/model_average.py new file mode 100644 index 0000000000000000000000000000000000000000..c7082b293883a3cdf38430498a00cfc9d6537f19 --- /dev/null +++ b/neuralangelo-main/imaginaire/models/utils/model_average.py @@ -0,0 +1,113 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import copy + +import torch +from torch import nn +from imaginaire.utils.misc import requires_grad + + +def reset_batch_norm(m): + r"""Reset batch norm statistics + + Args: + m: Pytorch module + """ + if hasattr(m, 'reset_running_stats'): + m.reset_running_stats() + + +def calibrate_batch_norm_momentum(m): + r"""Calibrate batch norm momentum + + Args: + m: Pytorch module + """ + if hasattr(m, 'reset_running_stats'): + # if m._get_name() == 'SyncBatchNorm': + if 'BatchNorm' in m._get_name(): + m.momentum = 1.0 / float(m.num_batches_tracked + 1) + + +class ModelAverage(nn.Module): + r"""In this model average implementation, the spectral layers are + absorbed in the model parameter by default. If such options are + turned on, be careful with how you do the training. Remember to + re-estimate the batch norm parameters before using the model. + + Args: + module (torch nn module): Torch network. + beta (float): Moving average weights. How much we weight the past. + start_iteration (int): From which iteration, we start the update. + """ + def __init__(self, module, beta=0.9999, start_iteration=0): + super(ModelAverage, self).__init__() + + self.module = module + # A shallow copy creates a new object which stores the reference of + # the original elements. + # A deep copy creates a new object and recursively adds the copies of + # nested objects present in the original elements. + self._averaged_model = copy.deepcopy(self.module).to('cuda') + self.stream = torch.cuda.Stream() + + self.beta = beta + + self.start_iteration = start_iteration + # This buffer is to track how many iterations has the model been + # trained for. We will ignore the first $(start_iterations) and start + # the averaging after. + self.register_buffer('num_updates_tracked', + torch.tensor(0, dtype=torch.long)) + self.num_updates_tracked = self.num_updates_tracked.to('cuda') + self.averaged_model.eval() + + # Averaged model does not require grad. + requires_grad(self.averaged_model, False) + + @property + def averaged_model(self): + self.stream.synchronize() + return self._averaged_model + + def forward(self, *inputs, **kwargs): + r"""PyTorch module forward function overload.""" + return self.module(*inputs, **kwargs) + + @torch.no_grad() + def update_average(self): + r"""Update the moving average.""" + self.stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(self.stream): + self.num_updates_tracked += 1 + if self.num_updates_tracked <= self.start_iteration: + beta = 0. + else: + beta = self.beta + source_dict = self.module.state_dict() + target_dict = self._averaged_model.state_dict() + source_list = [] + target_list = [] + for key in target_dict: + if 'num_batches_tracked' in key: + continue + source_list.append(source_dict[key].data) + target_list.append(target_dict[key].data.float()) + + torch._foreach_mul_(target_list, beta) + torch._foreach_add_(target_list, source_list, alpha=1 - beta) + + def __repr__(self): + r"""Returns a string that holds a printable representation of an + object""" + return self.module.__repr__() diff --git a/neuralangelo-main/imaginaire/trainers/base.py b/neuralangelo-main/imaginaire/trainers/base.py new file mode 100644 index 0000000000000000000000000000000000000000..1c3edf1541e289eacff7c3a9cd511c4d6bb34f13 --- /dev/null +++ b/neuralangelo-main/imaginaire/trainers/base.py @@ -0,0 +1,685 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import importlib +import json +import os +import threading +import time +import wandb +from tqdm import tqdm +import inspect + +import torch +from torch.autograd import profiler +from torch.cuda.amp import GradScaler, autocast + +from imaginaire.datasets.utils.get_dataloader import get_train_dataloader, get_val_dataloader, get_test_dataloader +from imaginaire.models.utils.init_weight import weights_init, weights_rescale +from imaginaire.trainers.utils.get_trainer import _calculate_model_size, get_optimizer, get_scheduler, wrap_model + +from imaginaire.utils.misc import to_cuda, requires_grad, to_cpu, Timer +from imaginaire.utils.distributed import master_only_print as print +from imaginaire.utils.distributed import is_master, get_rank +from imaginaire.utils.set_random_seed import set_random_seed + + +class BaseTrainer(object): + r"""Base trainer. We expect that all trainers inherit this class. + + Args: + cfg (obj): Global configuration. + is_inference (bool): if True, load the test dataloader and run in inference mode. + """ + + def __init__(self, cfg, is_inference=True, seed=0): + super().__init__() + print('Setup trainer.') + self.cfg = cfg + torch.cuda.set_device(cfg.local_rank) + # Create objects for the networks, optimizers, and schedulers. + self.model = self.setup_model(cfg, seed=seed) + if not is_inference: + self.optim = self.setup_optimizer(cfg, self.model, seed=seed) + self.sched = self.setup_scheduler(cfg, self.optim) + else: + self.optim = None + self.sched = None + self.model = self.wrap_model(cfg, self.model) + # Data loaders & inference mode. + self.is_inference = is_inference + # Initialize automatic mixed precision training. + self.init_amp() + # Initialize loss functions. + self.init_losses(cfg) + + self.checkpointer = Checkpointer(cfg, self.model, self.optim, self.sched) + self.timer = Timer(cfg) + + # -------- The initialization steps below can be skipped during inference. -------- + if self.is_inference: + return + + # Initialize logging attributes. + self.init_logging_attributes() + # Initialize validation parameters. + self.init_val_parameters() + # AWS credentials. + if hasattr(cfg, 'aws_credentials_file'): + with open(cfg.aws_credentials_file) as fin: + self.credentials = json.load(fin) + else: + self.credentials = None + if 'TORCH_HOME' not in os.environ: + os.environ['TORCH_HOME'] = os.path.join(os.environ['HOME'], ".cache") + + def set_data_loader(self, cfg, split, shuffle=True, drop_last=True, seed=0): + """Set the data loader corresponding to the indicated split. + Args: + split (str): Must be either 'train', 'val', or 'test'. + shuffle (bool): Whether to shuffle the data (only applies to the training set). + drop_last (bool): Whether to drop the last batch if it is not full (only applies to the training set). + seed (int): Random seed. + """ + assert (split in ["train", "val", "test"]) + if split == "train": + self.train_data_loader = get_train_dataloader(cfg, shuffle=shuffle, drop_last=drop_last, seed=seed) + elif split == "val": + self.eval_data_loader = get_val_dataloader(cfg, seed=seed) + elif split == "test": + self.eval_data_loader = get_test_dataloader(cfg) + + def setup_model(self, cfg, seed=0): + r"""Return the networks. We will first set the random seed to a fixed value so that each GPU copy will be + initialized to have the same network weights. We will then use different random seeds for different GPUs. + After this we will wrap the network with a moving average model if applicable. + + The following objects are constructed as class members: + - model (obj): Model object (historically: generator network object). + + Args: + cfg (obj): Global configuration. + seed (int): Random seed. + """ + # We first set the random seed for all the process so that we initialize each copy of the network the same. + set_random_seed(seed, by_rank=False) + # Construct networks + lib_model = importlib.import_module(cfg.model.type) + model = lib_model.Model(cfg.model, cfg.data) + print('model parameter count: {:,}'.format(_calculate_model_size(model))) + print(f'Initialize model weights using type: {cfg.trainer.init.type}, gain: {cfg.trainer.init.gain}') + init_bias = getattr(cfg.trainer.init, 'bias', None) + init_gain = cfg.trainer.init.gain or 1. + model.apply(weights_init(cfg.trainer.init.type, init_gain, init_bias)) + model.apply(weights_rescale()) + model = model.to('cuda') + # Different GPU copies of the same model will receive noises initialized with different random seeds + # (if applicable) thanks to the set_random_seed command (GPU #K has random seed = args.seed + K). + set_random_seed(seed, by_rank=True) + return model + + def setup_optimizer(self, cfg, model, seed=0): + r"""Return the optimizers. + + The following objects are constructed as class members: + - optim (obj): Model optimizer object. + + Args: + cfg (obj): Global configuration. + seed (int): Random seed. + """ + optim = get_optimizer(cfg.optim, model) + self.optim_zero_grad_kwargs = {} + if 'set_to_none' in inspect.signature(optim.zero_grad).parameters: + self.optim_zero_grad_kwargs['set_to_none'] = True + return optim + + def setup_scheduler(self, cfg, optim): + r"""Return the schedulers. + + The following objects are constructed as class members: + - sched (obj): Model optimizer scheduler object. + + Args: + cfg (obj): Global configuration. + """ + return get_scheduler(cfg.optim, optim) + + def wrap_model(self, cfg, model): + # Moving average model & data distributed data parallel wrapping. + model = wrap_model(cfg, model) + # Get actual modules from wrappers. + if cfg.trainer.ema_config.enabled: + # Two wrappers (DDP + model average). + self.model_module = model.module.module + else: + # One wrapper (DDP) + self.model_module = model.module + return model + + def init_amp(self): + r"""Initialize automatic mixed precision training.""" + + if getattr(self.cfg.trainer, 'allow_tf32', True): + print("Allow TensorFloat32 operations on supported devices") + torch.backends.cudnn.allow_tf32 = True + torch.backends.cuda.matmul.allow_tf32 = True + else: + torch.backends.cudnn.allow_tf32 = False + torch.backends.cuda.matmul.allow_tf32 = False + + if self.cfg.trainer.amp_config.enabled: + print("Using automatic mixed precision training.") + + # amp scaler can be used without mixed precision training + if hasattr(self.cfg.trainer, 'scaler_config'): + scaler_kwargs = vars(self.cfg.trainer.scaler_config) + scaler_kwargs['enabled'] = self.cfg.trainer.amp_config.enabled and \ + getattr(self.cfg.trainer.scaler_config, 'enabled', True) + else: + scaler_kwargs = vars(self.cfg.trainer.amp_config) # backward compatibility + scaler_kwargs.pop('dtype', None) + scaler_kwargs.pop('cache_enabled', None) + + self.scaler = GradScaler(**scaler_kwargs) + + def init_losses(self, cfg): + r"""Initialize loss functions. All loss names have weights. Some have criterion modules.""" + self.losses = dict() + + # Mapping from loss names to criterion modules. + self.criteria = torch.nn.ModuleDict() + # Mapping from loss names to loss weights. + self.weights = dict() + + self._init_loss(cfg) # this should be implemented by children classes + + for loss_name, loss_weight in self.weights.items(): + print("Loss {:<20} Weight {}".format(loss_name, loss_weight)) + if loss_name in self.criteria.keys() and self.criteria[loss_name] is not None: + self.criteria[loss_name].to('cuda') + + def init_logging_attributes(self): + r"""Initialize logging attributes.""" + self.current_iteration = 0 + self.current_epoch = 0 + self.start_iteration_time = None + self.start_epoch_time = None + self.elapsed_iteration_time = 0 + if self.cfg.speed_benchmark: + self.timer.reset() + + def init_val_parameters(self): + r"""Initialize validation parameters.""" + if self.cfg.metrics_iter is None: + self.cfg.metrics_iter = self.cfg.checkpoint.save_iter + if self.cfg.metrics_epoch is None: + self.cfg.metrics_epoch = self.cfg.checkpoint.save_epoch + + def init_wandb(self, cfg, wandb_id=None, project="", run_name=None, mode="online", resume="allow", use_group=False): + r"""Initialize Weights & Biases (wandb) logger. + + Args: + cfg (obj): Global configuration. + wandb_id (str): A unique ID for this run, used for resuming. + project (str): The name of the project where you're sending the new run. + If the project is not specified, the run is put in an "Uncategorized" project. + run_name (str): name for each wandb run (useful for logging changes) + mode (str): online/offline/disabled + """ + if is_master(): + print('Initialize wandb') + if not wandb_id: + wandb_path = os.path.join(cfg.logdir, "wandb_id.txt") + if self.checkpointer.resume and os.path.exists(wandb_path): + with open(wandb_path, "r") as f: + wandb_id = f.read() + else: + wandb_id = wandb.util.generate_id() + with open(wandb_path, "w") as f: + f.write(wandb_id) + if use_group: + group, name = cfg.logdir.split("/")[-2:] + else: + group, name = None, os.path.basename(cfg.logdir) + + if run_name is not None: + name = run_name + + wandb.init(id=wandb_id, + project=project, + config=cfg, + group=group, + name=name, + dir=cfg.logdir, + resume=resume, + settings=wandb.Settings(start_method="fork"), + mode=mode) + wandb.config.update({'dataset': cfg.data.name}) + if self.model_module is not None: + wandb.watch(self.model_module) + + def start_of_epoch(self, current_epoch): + r"""Things to do before an epoch. + + Args: + current_epoch (int): Current number of epoch. + """ + self._start_of_epoch(current_epoch) + self.current_epoch = current_epoch + self.start_epoch_time = time.time() + + def start_of_iteration(self, data, current_iteration): + r"""Things to do before an iteration. + + Args: + data (dict): Data used for the current iteration. + current_iteration (int): Current number of iteration. + """ + data = self._start_of_iteration(data, current_iteration) + data = to_cuda(data) + self.current_iteration = current_iteration + self.model.train() + self.start_iteration_time = time.time() + return data + + def end_of_iteration(self, data, current_epoch, current_iteration): + r"""Things to do after an iteration. + + Args: + data (dict): Data used for the current iteration. + current_epoch (int): Current number of epoch. + current_iteration (int): Current number of iteration. + """ + self.current_iteration = current_iteration + self.current_epoch = current_epoch + + # Accumulate time + self.elapsed_iteration_time += time.time() - self.start_iteration_time + # Logging. + if current_iteration % self.cfg.logging_iter == 0: + avg_time = self.elapsed_iteration_time / self.cfg.logging_iter + self.timer.time_iteration = avg_time + print('Iteration: {}, average iter time: {:6f}.'.format(current_iteration, avg_time)) + self.elapsed_iteration_time = 0 + + if self.cfg.speed_benchmark: + # only needed when analyzing computation bottleneck. + self.timer._print_speed_benchmark(avg_time) + + self._end_of_iteration(data, current_epoch, current_iteration) + + # Save everything to the checkpoint by time period. + if self.checkpointer.reached_checkpointing_period(self.timer): + self.checkpointer.save(current_epoch, current_iteration) + self.timer.checkpoint_tic() # reset timer + + # Save everything to the checkpoint. + if current_iteration % self.cfg.checkpoint.save_iter == 0: + self.checkpointer.save(current_epoch, current_iteration) + + # Save everything to the checkpoint using the name 'latest_checkpoint.pt'. + if current_iteration % self.cfg.checkpoint.save_latest_iter == 0: + if current_iteration >= self.cfg.checkpoint.save_latest_iter: + self.checkpointer.save(current_epoch, current_iteration, True) + + # Update the learning rate policy for the generator if operating in the iteration mode. + if self.cfg.optim.sched.iteration_mode: + self.sched.step() + + # This iteration was successfully finished. Reset timeout counter. + self.timer.reset_timeout_counter() + + def end_of_epoch(self, data, current_epoch, current_iteration): + r"""Things to do after an epoch. + + Args: + data (dict): Data used for the current iteration. + + current_epoch (int): Current number of epoch. + current_iteration (int): Current number of iteration. + """ + # Update the learning rate policy for the generator if operating in the epoch mode. + self.current_iteration = current_iteration + self.current_epoch = current_epoch + if not self.cfg.optim.sched.iteration_mode: + self.sched.step() + elapsed_epoch_time = time.time() - self.start_epoch_time + # Logging. + print('Epoch: {}, total time: {:6f}.'.format(current_epoch, elapsed_epoch_time)) + self.timer.time_epoch = elapsed_epoch_time + self._end_of_epoch(data, current_epoch, current_iteration) + + # Save everything to the checkpoint. + if current_epoch % self.cfg.checkpoint.save_epoch == 0: + self.checkpointer.save(current_epoch, current_iteration) + + def _extra_step(self, data): + pass + + def _start_of_epoch(self, current_epoch): + r"""Operations to do before starting an epoch. + + Args: + current_epoch (int): Current number of epoch. + """ + pass + + def _start_of_iteration(self, data, current_iteration): + r"""Operations to do before starting an iteration. + + Args: + data (dict): Data used for the current iteration. + current_iteration (int): Current epoch number. + Returns: + (dict): Data used for the current iteration. They might be + processed by the custom _start_of_iteration function. + """ + return data + + def _end_of_iteration(self, data, current_epoch, current_iteration): + r"""Operations to do after an iteration. + + Args: + data (dict): Data used for the current iteration. + current_epoch (int): Current number of epoch. + current_iteration (int): Current epoch number. + """ + pass + + def _end_of_epoch(self, data, current_epoch, current_iteration): + r"""Operations to do after an epoch. + + Args: + data (dict): Data used for the current iteration. + current_epoch (int): Current number of epoch. + current_iteration (int): Current epoch number. + """ + pass + + def _get_visualizations(self, data): + r"""Compute visualization outputs. + + Args: + data (dict): Data used for the current iteration. + """ + return None + + def _init_loss(self, cfg): + r"""Every trainer should implement its own init loss function.""" + raise NotImplementedError + + def train_step(self, data, last_iter_in_epoch=False): + r"""One training step. + + Args: + data (dict): Data used for the current iteration. + """ + # Set requires_grad flags. + requires_grad(self.model_module, True) + + # Compute the loss. + self.timer._time_before_forward() + + autocast_dtype = getattr(self.cfg.trainer.amp_config, 'dtype', 'float16') + autocast_dtype = torch.bfloat16 if autocast_dtype == 'bfloat16' else torch.float16 + amp_kwargs = { + 'enabled': self.cfg.trainer.amp_config.enabled, + 'dtype': autocast_dtype + } + with autocast(**amp_kwargs): + total_loss = self.model_forward(data) + # Scale down the loss w.r.t. gradient accumulation iterations. + total_loss = total_loss / float(self.cfg.trainer.grad_accum_iter) + + # Backpropagate the loss. + self.timer._time_before_backward() + self.scaler.scale(total_loss).backward() + + self._extra_step(data) + + # Perform an optimizer step. This enables gradient accumulation when grad_accum_iter is not 1. + if (self.current_iteration + 1) % self.cfg.trainer.grad_accum_iter == 0 or last_iter_in_epoch: + self.timer._time_before_step() + self.scaler.step(self.optim) + self.scaler.update() + # Zero out the gradients. + self.optim.zero_grad(**self.optim_zero_grad_kwargs) + + # Update model average. + self.timer._time_before_model_avg() + if self.cfg.trainer.ema_config.enabled: + self.model.module.update_average() + + self._detach_losses() + self.timer._time_before_leave_gen() + + def model_forward(self, data): + r"""Every trainer should implement its own model forward.""" + raise NotImplementedError + + def train(self, cfg, data_loader, single_gpu=False, profile=False, show_pbar=False): + r"""Generic training loop. Main structure in a nutshell: + for epoch in [start_epoch, end_epoch]: + for batch in dataset (one epoch): + train_step(batch) + + Args: + cfg (obj): Global configuration. + data_loader (torch.utils.data.DataLoader): PyTorch dataloader. + single_gpu (bool): Use only a single GPU. + profile (bool): Enable profiling. + show_pbar (bool): Whether to show the progress bar + """ + start_epoch = self.checkpointer.resume_epoch or self.current_epoch # The epoch to start with. + current_iteration = self.checkpointer.resume_iteration or self.current_iteration # The starting iteration. + + self.timer.checkpoint_tic() # start timer + self.timer.reset_timeout_counter() + for current_epoch in range(start_epoch, cfg.max_epoch): + if not single_gpu: + data_loader.sampler.set_epoch(current_epoch) + self.start_of_epoch(current_epoch) + if show_pbar: + data_loader_wrapper = tqdm(data_loader, desc=f"Training epoch {current_epoch + 1}", leave=False) + else: + data_loader_wrapper = data_loader + for it, data in enumerate(data_loader_wrapper): + with profiler.profile(enabled=profile, + use_cuda=True, + profile_memory=True, + record_shapes=True) as prof: + data = self.start_of_iteration(data, current_iteration) + + self.train_step(data, last_iter_in_epoch=(it == len(data_loader) - 1)) + + current_iteration += 1 + if show_pbar: + data_loader_wrapper.set_postfix(iter=current_iteration) + if it == len(data_loader) - 1: + self.end_of_iteration(data, current_epoch + 1, current_iteration) + else: + self.end_of_iteration(data, current_epoch, current_iteration) + if current_iteration >= cfg.max_iter: + print('Done with training!!!') + return + if profile: + print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20)) + prof.export_chrome_trace(os.path.join(cfg.logdir, "trace.json")) + + self.end_of_epoch(data, current_epoch + 1, current_iteration) + print('Done with training!!!') + + def test(self, data_loader, output_dir, inference_args, show_pbar=False): + r"""Compute results images and save the results in the specified folder. + Args: + data_loader (torch.utils.data.DataLoader): PyTorch dataloader. + output_dir (str): Target location for saving the output image. + """ + pass + + def _get_total_loss(self): + r"""Return the total loss to be backpropagated. + """ + total_loss = torch.tensor(0., device=torch.device('cuda')) + # Iterates over all possible losses. + for loss_name in self.weights: + if loss_name in self.losses: + # Multiply it with the corresponding weight and add it to the total loss. + total_loss += self.losses[loss_name] * self.weights[loss_name] + self.losses['total'] = total_loss # logging purpose + return total_loss + + def _detach_losses(self): + r"""Detach all logging variables to prevent potential memory leak.""" + for loss_name in self.losses: + self.losses[loss_name] = self.losses[loss_name].detach() + + def finalize(self, cfg): + # Finish the W&B logger. + if is_master(): + wandb.finish() + + +class Checkpointer(object): + + def __init__(self, cfg, model, optim=None, sched=None): + self.model = model + self.optim = optim + self.sched = sched + self.logdir = cfg.logdir + self.save_period = cfg.checkpoint.save_period + self.strict_resume = cfg.checkpoint.strict_resume + self.iteration_mode = cfg.optim.sched.iteration_mode + self.resume = False + self.resume_epoch = self.resume_iteration = None + + def save(self, current_epoch, current_iteration, latest=False): + r"""Save network weights, optimizer parameters, scheduler parameters to a checkpoint. + + Args: + current_epoch (int): Current epoch. + current_iteration (int): Current iteration. + latest (bool): If ``True``, save it using the name 'latest_checkpoint.pt'. + """ + checkpoint_file = 'latest_checkpoint.pt' if latest else \ + f'epoch_{current_epoch:05}_iteration_{current_iteration:09}_checkpoint.pt' + if is_master(): + save_dict = to_cpu(self._collect_state_dicts()) + save_dict.update( + epoch=current_epoch, + iteration=current_iteration, + ) + # Run the checkpoint saver in a separate thread. + threading.Thread( + target=self._save_worker, daemon=False, args=(save_dict, checkpoint_file, get_rank())).start() + checkpoint_path = self._get_full_path(checkpoint_file) + return checkpoint_path + + def _save_worker(self, save_dict, checkpoint_file, rank=0): + checkpoint_path = self._get_full_path(checkpoint_file) + # Save to local disk. + os.makedirs(os.path.dirname(checkpoint_path), exist_ok=True) + torch.save(save_dict, checkpoint_path) + if rank == 0: + self.write_latest_checkpoint_file(checkpoint_file) + print('Saved checkpoint to {}'.format(checkpoint_path)) + + def _collect_state_dicts(self): + r"""Collect all the state dicts from network modules to be saved.""" + return dict( + model=self.model.state_dict(), + optim=self.optim.state_dict(), + sched=self.sched.state_dict(), + ) + + def load(self, checkpoint_path=None, resume=False, load_opt=True, load_sch=True, **kwargs): + r"""Load network weights, optimizer parameters, scheduler parameters from a checkpoint. + Args: + checkpoint_path (str): Path to the checkpoint (local file or S3 key). + resume (bool): if False, only the model weights are loaded. If True, the metadata (epoch/iteration) and + optimizer/scheduler (optional) are also loaded. + load_opt (bool): Whether to load the optimizer state dict (resume should be True). + load_sch (bool): Whether to load the scheduler state dict (resume should be True). + """ + # Priority: (1) checkpoint_path (2) latest_path (3) train from scratch. + self.resume = resume + # If checkpoint path were not specified, try to load the latest one from the same run. + if resume and checkpoint_path is None: + latest_checkpoint_file = self.read_latest_checkpoint_file() + if latest_checkpoint_file is not None: + checkpoint_path = self._get_full_path(latest_checkpoint_file) + # Load checkpoint. + if checkpoint_path is not None: + self._check_checkpoint_exists(checkpoint_path) + self.checkpoint_path = checkpoint_path + state_dict = torch.load(checkpoint_path, map_location=lambda storage, loc: storage) + print(f"Loading checkpoint (local): {checkpoint_path}") + # Load the state dicts. + print('- Loading the model...') + self.model.load_state_dict(state_dict['model'], strict=self.strict_resume) + if resume: + try: + self.resume_epoch = state_dict['epoch'] + self.resume_iteration = state_dict['iteration'] + except Exception: # TODO: for backward compatibility, should be removed eventually. + self.resume_epoch = state_dict['current_epoch'] + self.resume_iteration = state_dict['current_iteration'] + self.sched.last_epoch = self.resume_iteration if self.iteration_mode else self.resume_epoch + if load_opt: + print('- Loading the optimizer...') + self.optim.load_state_dict(state_dict['optim']) + if load_sch: + print('- Loading the scheduler...') + self.sched.load_state_dict(state_dict['sched']) + print(f"Done with loading the checkpoint (epoch {self.resume_epoch}, iter {self.resume_iteration}).") + else: + print('Done with loading the checkpoint.') + else: + # Checkpoint not found and not specified. We will train everything from scratch. + print('Training from scratch.') + torch.cuda.empty_cache() + + def _get_full_path(self, file): + return os.path.join(self.logdir, file) + + def _get_latest_pointer_path(self): + return self._get_full_path('latest_checkpoint.txt') + + def read_latest_checkpoint_file(self): + checkpoint_file = None + latest_path = self._get_latest_pointer_path() + if os.path.exists(latest_path): + checkpoint_file = open(latest_path).read().strip() + if checkpoint_file.startswith("latest_checkpoint:"): # TODO: for backward compatibility, to be removed + checkpoint_file = checkpoint_file.split(' ')[-1] + return checkpoint_file + + def write_latest_checkpoint_file(self, checkpoint_file): + latest_path = self._get_latest_pointer_path() + content = f"{checkpoint_file}\n" + with open(latest_path, "w") as file: + file.write(content) + + def _check_checkpoint_exists(self, checkpoint_path): + if not os.path.exists(checkpoint_path): + raise FileNotFoundError(f'File not found (local): {checkpoint_path}') + + def reached_checkpointing_period(self, timer): + save_now = torch.cuda.BoolTensor([False]) + if is_master(): + if timer.checkpoint_toc() > self.save_period: + save_now.fill_(True) + if save_now: + if is_master(): + print('checkpointing period!') + return save_now diff --git a/neuralangelo-main/imaginaire/trainers/utils/get_trainer.py b/neuralangelo-main/imaginaire/trainers/utils/get_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..61fc06c9de13817db047853b1d008499ac4c0834 --- /dev/null +++ b/neuralangelo-main/imaginaire/trainers/utils/get_trainer.py @@ -0,0 +1,223 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import importlib +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.optim import lr_scheduler +from imaginaire.models.utils.model_average import ModelAverage + + +def get_trainer(cfg, is_inference=True, seed=0): + """Return the trainer object. + + Args: + cfg (Config): Loaded config object. + is_inference (bool): Inference mode. + + Returns: + (obj): Trainer object. + """ + trainer_lib = importlib.import_module(cfg.trainer.type) + trainer = trainer_lib.Trainer(cfg, is_inference=is_inference, seed=seed) + return trainer + + +def wrap_model(cfg, model): + r"""Wrap the networks with AMP DDP and (optionally) model average. + + Args: + cfg (obj): Global configuration. + model (obj): Model object. + + Returns: + (dict): + - model (obj): Model object. + """ + # Apply model average wrapper. + if cfg.trainer.ema_config.enabled: + model = ModelAverage(model, + cfg.trainer.ema_config.beta, + cfg.trainer.ema_config.start_iteration, + ) + model = _wrap_model(cfg, model) + return model + + +class WrappedModel(nn.Module): + r"""Dummy wrapping the module. + """ + + def __init__(self, module): + super(WrappedModel, self).__init__() + self.module = module + + def forward(self, *args, **kwargs): + r"""PyTorch module forward function overload.""" + return self.module(*args, **kwargs) + + +def _wrap_model(cfg, model): + r"""Wrap a model for distributed data parallel training. + + Args: + model (obj): PyTorch network model. + + Returns: + (obj): Wrapped PyTorch network model. + """ + # Apply DDP wrapper. + if dist.is_available() and dist.is_initialized(): + model = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[cfg.local_rank], + output_device=cfg.local_rank, + find_unused_parameters=cfg.trainer.ddp_config.find_unused_parameters, + static_graph=cfg.trainer.ddp_config.static_graph, + broadcast_buffers=False, + ) + else: + model = WrappedModel(model) + return model + + +def _calculate_model_size(model): + r"""Calculate number of parameters in a PyTorch network. + + Args: + model (obj): PyTorch network. + + Returns: + (int): Number of parameters. + """ + return sum(p.numel() for p in model.parameters() if p.requires_grad) + + +def get_optimizer(cfg_optim, model): + r"""Return the optimizer object. + + Args: + cfg_optim (obj): Config for the specific optimization module (gen/dis). + model (obj): PyTorch network object. + + Returns: + (obj): Pytorch optimizer + """ + if hasattr(model, 'get_param_groups'): + # Allow the network to use different hyperparameters (e.g., learning rate) for different parameters. + params = model.get_param_groups(cfg_optim) + else: + params = model.parameters() + + try: + # Try the PyTorch optimizer class first. + optimizer_class = getattr(torch.optim, cfg_optim.type) + except AttributeError: + raise NotImplementedError(f"Optimizer {cfg_optim.type} is not yet implemented.") + optimizer_kwargs = cfg_optim.params + + # We will try to use fuse optimizers by default. + try: + from apex.optimizers import FusedAdam, FusedSGD + fused_opt = cfg_optim.fused_opt + except (ImportError, ModuleNotFoundError): + fused_opt = False + + if fused_opt: + if cfg_optim.type == 'Adam': + optimizer_class = FusedAdam + optimizer_kwargs['adam_w_mode'] = False + elif cfg_optim.type == 'AdamW': + optimizer_class = FusedAdam + optimizer_kwargs['adam_w_mode'] = True + elif cfg_optim.type == 'SGD': + optimizer_class = FusedSGD + if cfg_optim.type in ["RAdam", "RMSprop"]: + optimizer_kwargs["foreach"] = fused_opt + + optim = optimizer_class(params, **optimizer_kwargs) + + return optim + + +def get_scheduler(cfg_optim, optim): + """Return the scheduler object. + + Args: + cfg_optim (obj): Config for the specific optimization module (gen/dis). + optim (obj): PyTorch optimizer object. + + Returns: + (obj): Scheduler + """ + if cfg_optim.sched.type == 'step': + scheduler = lr_scheduler.StepLR(optim, + step_size=cfg_optim.sched.step_size, + gamma=cfg_optim.sched.gamma) + elif cfg_optim.sched.type == 'constant': + scheduler = lr_scheduler.LambdaLR(optim, lambda x: 1) + elif cfg_optim.sched.type == 'linear_warmup': + scheduler = lr_scheduler.LambdaLR( + optim, lambda x: x * 1.0 / cfg_optim.sched.warmup if x < cfg_optim.sched.warmup else 1.0) + elif cfg_optim.sched.type == 'cosine_warmup': + + warmup_scheduler = lr_scheduler.LinearLR( + optim, + start_factor=1.0 / cfg_optim.sched.warmup, + end_factor=1.0, + total_iters=cfg_optim.sched.warmup + ) + T_max_val = cfg_optim.sched.decay_steps - cfg_optim.sched.warmup + cosine_lr_scheduler = lr_scheduler.CosineAnnealingLR( + optim, + T_max=T_max_val, + eta_min=getattr(cfg_optim.sched, 'eta_min', 0), + ) + scheduler = lr_scheduler.SequentialLR( + optim, + schedulers=[warmup_scheduler, cosine_lr_scheduler], + milestones=[cfg_optim.sched.warmup] + ) + + elif cfg_optim.sched.type == 'linear': + # Start linear decay from here. + decay_start = cfg_optim.sched.decay_start + # End linear decay here. + # Continue to train using the lowest learning rate till the end. + decay_end = cfg_optim.sched.decay_end + # Lowest learning rate multiplier. + decay_target = cfg_optim.sched.decay_target + + def sch(x): + decay = ((x - decay_start) * decay_target + decay_end - x) / (decay_end - decay_start) + return min(max(decay, decay_target), 1.) + + scheduler = lr_scheduler.LambdaLR(optim, lambda x: sch(x)) + elif cfg_optim.sched.type == 'step_with_warmup': + # The step_size and gamma follows the signature of lr_scheduler.StepLR. + step_size = cfg_optim.sched.step_size, + gamma = cfg_optim.sched.gamma + # An additional parameter defines the warmup iteration. + warmup_step_size = cfg_optim.sched.warmup_step_size + + def sch(x): + lr_after_warmup = gamma ** (warmup_step_size // step_size) + if x < warmup_step_size: + return x / warmup_step_size * lr_after_warmup + else: + return gamma ** (x // step_size) + + scheduler = lr_scheduler.LambdaLR(optim, lambda x: sch(x)) + else: + return NotImplementedError('Learning rate policy {} not implemented.'.format(cfg_optim.sched.type)) + return scheduler diff --git a/neuralangelo-main/imaginaire/trainers/utils/logging.py b/neuralangelo-main/imaginaire/trainers/utils/logging.py new file mode 100644 index 0000000000000000000000000000000000000000..c3ecf81ee47d7e8627814a77e238deaa6fdba14c --- /dev/null +++ b/neuralangelo-main/imaginaire/trainers/utils/logging.py @@ -0,0 +1,72 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import datetime +import os + +import torch.distributed as dist + +from imaginaire.utils.distributed import is_master, broadcast_object_list +from imaginaire.utils.distributed import master_only_print as print +from imaginaire.trainers.utils.meters import set_summary_writer + + +def get_date_uid(): + """Generate a unique id based on date. + Returns: + str: Return uid string, e.g. '20171122171307111552'. + """ + return str(datetime.datetime.now().strftime("%Y_%m%d_%H%M_%S")) + + +def init_logging(config_path, logdir, makedir=True): + r"""Create log directory for storing checkpoints and output images. + + Args: + config_path (str): Path to the configuration file. + logdir (str or None): Log directory name + makedir (bool): Make a new dir or not + Returns: + str: Return log dir + """ + def _create_logdir(_config_path, _logdir, _root_dir): + config_file = os.path.basename(_config_path) + date_uid = get_date_uid() + # example: logs/2019_0125_1047_58_spade_cocostuff + _log_file = '_'.join([date_uid, os.path.splitext(config_file)[0]]) + if _logdir is None: + _logdir = os.path.join(_root_dir, _log_file) + if makedir: + print('Make folder {}'.format(_logdir)) + os.makedirs(_logdir, exist_ok=True) + _tensorboard_dir = os.path.join(_logdir, 'tensorboard') + os.makedirs(_tensorboard_dir, exist_ok=True) + set_summary_writer(_tensorboard_dir) + return _logdir + + root_dir = 'logs' + if dist.is_available(): + if dist.is_initialized(): + message = [None] + if is_master(): + logdir = _create_logdir(config_path, logdir, root_dir) + message = [logdir] + + # Send logdir from master to all workers. + message = broadcast_object_list(message=message, src=0) + logdir = message[0] + else: + logdir = _create_logdir(config_path, logdir, root_dir) + else: + logdir = _create_logdir(config_path, logdir, root_dir) + + return logdir diff --git a/neuralangelo-main/imaginaire/trainers/utils/meters.py b/neuralangelo-main/imaginaire/trainers/utils/meters.py new file mode 100644 index 0000000000000000000000000000000000000000..9dbf2854cb053b757e5d21ff1680c3c06bb4bc9a --- /dev/null +++ b/neuralangelo-main/imaginaire/trainers/utils/meters.py @@ -0,0 +1,147 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import math +import torch +import wandb +from torch.utils.tensorboard import SummaryWriter + +from imaginaire.utils.distributed import master_only, dist_all_reduce_tensor, \ + is_master, get_rank + +from imaginaire.utils.distributed import master_only_print as print + +LOG_WRITER = None +LOG_DIR = None + + +@torch.no_grad() +def sn_reshape_weight_to_matrix(weight): + r"""Reshape weight to obtain the matrix form. + + Args: + weight (Parameters): pytorch layer parameter tensor. + """ + weight_mat = weight + height = weight_mat.size(0) + return weight_mat.reshape(height, -1) + + +@torch.no_grad() +def get_weight_stats(mod): + r"""Get weight state + + Args: + mod: Pytorch module + """ + if mod.weight_orig.grad is not None: + grad_norm = mod.weight_orig.grad.data.norm().item() + else: + grad_norm = 0. + weight_norm = mod.weight_orig.data.norm().item() + weight_mat = sn_reshape_weight_to_matrix(mod.weight_orig) + sigma = torch.sum(mod.weight_u * torch.mv(weight_mat, mod.weight_v)) + return grad_norm, weight_norm, sigma + + +@master_only +def set_summary_writer(log_dir): + r"""Set summary writer + + Args: + log_dir (str): Log directory. + """ + global LOG_DIR, LOG_WRITER + LOG_DIR = log_dir + LOG_WRITER = SummaryWriter(log_dir=log_dir) + + +def write_summary(name, summary, step, hist=False): + """Utility function for write summary to log_writer. + """ + global LOG_WRITER + lw = LOG_WRITER + if lw is None: + raise Exception("Log writer not set.") + if hist: + lw.add_histogram(name, summary, step) + else: + lw.add_scalar(name, summary, step) + + +class Meter(object): + """Meter is to keep track of statistics along steps. + Meters write values for purpose like printing average values. + Meters can be flushed to log files (i.e. TensorBoard for now) + regularly. + + Args: + name (str): the name of meter + reduce (bool): If ``True``, perform a distributed reduce for the log + values across all GPUs. + """ + + def __init__(self, name, reduce=True): + self.name = name + self.reduce = reduce + self.values = [] + + def reset(self): + r"""Reset the meter values""" + if not self.reduce and get_rank() != 0: + return + self.values = [] + + def write(self, value): + r"""Record the value""" + if not self.reduce and get_rank() != 0: + return + if value is not None: + self.values.append(value) + + def flush(self, step): + r"""Write the value in the tensorboard. + + Args: + step (int): Epoch or iteration number. + """ + if not self.reduce and get_rank() != 0: + return + values = torch.tensor(self.values, device="cuda") + if self.reduce: + values = dist_all_reduce_tensor(values) + + if not all(math.isfinite(x) for x in values): + print("meter {} contained a nan or inf.".format(self.name)) + filtered_values = list(filter(lambda x: math.isfinite(x), self.values)) + if float(len(filtered_values)) != 0: + value = float(sum(filtered_values)) / float(len(filtered_values)) + if is_master(): + write_summary(self.name, value, step) + wandb.log({self.name: value}, step=step) + self.reset() + + @master_only + def write_image(self, img_grid, step): + r"""Write the value in the tensorboard. + + Args: + img_grid: + step (int): Epoch or iteration number. + """ + if not self.reduce and get_rank() != 0: + return + global LOG_WRITER + lw = LOG_WRITER + if lw is None: + raise Exception("Log writer not set.") + lw.add_image("Visualizations", img_grid, step) diff --git a/neuralangelo-main/imaginaire/utils/cudnn.py b/neuralangelo-main/imaginaire/utils/cudnn.py new file mode 100644 index 0000000000000000000000000000000000000000..dcfb85ce16372a47c2941f0829625fd3e4578286 --- /dev/null +++ b/neuralangelo-main/imaginaire/utils/cudnn.py @@ -0,0 +1,30 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import torch.backends.cudnn as cudnn + +from imaginaire.utils.distributed import master_only_print as print + + +def init_cudnn(deterministic, benchmark): + r"""Initialize the cudnn module. The two things to consider is whether to + use cudnn benchmark and whether to use cudnn deterministic. If cudnn + benchmark is set, then the cudnn deterministic is automatically false. + + Args: + deterministic (bool): Whether to use cudnn deterministic. + benchmark (bool): Whether to use cudnn benchmark. + """ + cudnn.deterministic = deterministic + cudnn.benchmark = benchmark + print('cudnn benchmark: {}'.format(benchmark)) + print('cudnn deterministic: {}'.format(deterministic)) diff --git a/neuralangelo-main/imaginaire/utils/distributed.py b/neuralangelo-main/imaginaire/utils/distributed.py new file mode 100644 index 0000000000000000000000000000000000000000..51e81949dab2e2872f12341de833ac174c3e01ac --- /dev/null +++ b/neuralangelo-main/imaginaire/utils/distributed.py @@ -0,0 +1,153 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import functools +import ctypes + +import torch +import torch.distributed as dist +from contextlib import contextmanager + + +def init_dist(local_rank, backend='nccl', **kwargs): + r"""Initialize distributed training""" + if dist.is_available(): + if dist.is_initialized(): + return torch.cuda.current_device() + torch.cuda.set_device(local_rank) + dist.init_process_group(backend=backend, init_method='env://', **kwargs) + + # Increase the L2 fetch granularity for faster speed. + _libcudart = ctypes.CDLL('libcudart.so') + # Set device limit on the current device + # cudaLimitMaxL2FetchGranularity = 0x05 + pValue = ctypes.cast((ctypes.c_int * 1)(), ctypes.POINTER(ctypes.c_int)) + _libcudart.cudaDeviceSetLimit(ctypes.c_int(0x05), ctypes.c_int(128)) + _libcudart.cudaDeviceGetLimit(pValue, ctypes.c_int(0x05)) + # assert pValue.contents.value == 128 + + +def get_rank(): + r"""Get rank of the thread.""" + rank = 0 + if dist.is_available(): + if dist.is_initialized(): + rank = dist.get_rank() + return rank + + +def get_world_size(): + r"""Get world size. How many GPUs are available in this job.""" + world_size = 1 + if dist.is_available(): + if dist.is_initialized(): + world_size = dist.get_world_size() + return world_size + + +def broadcast_object_list(message, src=0): + r"""Broadcast object list from the master to the others""" + # Send logdir from master to all workers. + if dist.is_available(): + if dist.is_initialized(): + torch.distributed.broadcast_object_list(message, src=src) + return message + + +def master_only(func): + r"""Apply this function only to the master GPU.""" + @functools.wraps(func) + def wrapper(*args, **kwargs): + r"""Simple function wrapper for the master function""" + if get_rank() == 0: + return func(*args, **kwargs) + else: + return None + return wrapper + + +def is_master(): + r"""check if current process is the master""" + return get_rank() == 0 + + +def is_dist(): + return dist.is_initialized() + + +def barrier(): + if is_dist(): + dist.barrier() + + +@contextmanager +def master_first(): + if not is_master(): + barrier() + yield + if dist.is_initialized() and is_master(): + barrier() + + +def is_local_master(): + return torch.cuda.current_device() == 0 + + +@master_only +def master_only_print(*args): + r"""master-only print""" + print(*args) + + +def dist_reduce_tensor(tensor, rank=0, reduce='mean'): + r""" Reduce to rank 0 """ + world_size = get_world_size() + if world_size < 2: + return tensor + with torch.no_grad(): + dist.reduce(tensor, dst=rank) + if get_rank() == rank: + if reduce == 'mean': + tensor /= world_size + elif reduce == 'sum': + pass + else: + raise NotImplementedError + return tensor + + +def dist_all_reduce_tensor(tensor, reduce='mean'): + r""" Reduce to all ranks """ + world_size = get_world_size() + if world_size < 2: + return tensor + with torch.no_grad(): + dist.all_reduce(tensor) + if reduce == 'mean': + tensor /= world_size + elif reduce == 'sum': + pass + else: + raise NotImplementedError + return tensor + + +def dist_all_gather_tensor(tensor): + r""" gather to all ranks """ + world_size = get_world_size() + if world_size < 2: + return [tensor] + tensor_list = [ + torch.ones_like(tensor) for _ in range(dist.get_world_size())] + with torch.no_grad(): + dist.all_gather(tensor_list, tensor) + return tensor_list diff --git a/neuralangelo-main/imaginaire/utils/gpu_affinity.py b/neuralangelo-main/imaginaire/utils/gpu_affinity.py new file mode 100644 index 0000000000000000000000000000000000000000..081f3b3eb9f4d0a1aafa67f374c105b137c62a0c --- /dev/null +++ b/neuralangelo-main/imaginaire/utils/gpu_affinity.py @@ -0,0 +1,78 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import math +import os +# pynvml is a python bindings to the NVIDIA Management Library +# https://developer.nvidia.com/nvidia-management-library-nvml +# An API for monitoring and managing various states of the NVIDIA GPU devices. +# It provides direct access to the queries and commands exposed via nvidia-smi. + +import pynvml + +pynvml.nvmlInit() + + +def system_get_driver_version(): + r"""Get Driver Version""" + return pynvml.nvmlSystemGetDriverVersion() + + +def device_get_count(): + r"""Get number of devices""" + return pynvml.nvmlDeviceGetCount() + + +class Device(object): + r"""Device used for nvml.""" + _nvml_affinity_elements = math.ceil(os.cpu_count() / 64) + + def __init__(self, device_idx): + super().__init__() + self.handle = pynvml.nvmlDeviceGetHandleByIndex(device_idx) + + def get_name(self): + r"""Get obect name""" + return pynvml.nvmlDeviceGetName(self.handle) + + def get_cpu_affinity(self): + r"""Get CPU affinity""" + affinity_string = '' + for j in pynvml.nvmlDeviceGetCpuAffinity(self.handle, Device._nvml_affinity_elements): + # assume nvml returns list of 64 bit ints + affinity_string = '{:064b}'.format(j) + affinity_string + affinity_list = [int(x) for x in affinity_string] + affinity_list.reverse() # so core 0 is in 0th element of list + + return [i for i, e in enumerate(affinity_list) if e != 0] + + +def set_affinity(gpu_id=None): + r"""Set GPU affinity + + Args: + gpu_id (int): Which gpu device. + """ + if gpu_id is None: + gpu_id = int(os.getenv('LOCAL_RANK', 0)) + + dev = Device(gpu_id) + # os.sched_setaffinity() method in Python is used to set the CPU affinity mask of a process indicated + # by the specified process id. + # A process’s CPU affinity mask determines the set of CPUs on which it is eligible to run. + # Syntax: os.sched_setaffinity(pid, mask) + # pid=0 means the current process + os.sched_setaffinity(0, dev.get_cpu_affinity()) + + # list of ints + # representing the logical cores this process is now affinitied with + return os.sched_getaffinity(0) diff --git a/neuralangelo-main/imaginaire/utils/misc.py b/neuralangelo-main/imaginaire/utils/misc.py new file mode 100644 index 0000000000000000000000000000000000000000..2328f565cee58d51533288379632e5c38e26bc30 --- /dev/null +++ b/neuralangelo-main/imaginaire/utils/misc.py @@ -0,0 +1,376 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import collections +import functools +import os +import signal +import time +from collections import OrderedDict + +import torch +import torch.nn.functional as F +import wandb + +from imaginaire.utils.distributed import is_master, master_only + +string_classes = (str, bytes) + +from imaginaire.utils.termcolor import alert, PP # noqa + + +def santize_args(name, locals_fn): + args = {k: v for k, v in locals_fn.items()} + if 'kwargs' in args and args['kwargs']: + unused = PP(args['kwargs']) + alert(f'{name}: Unused kwargs\n{unused}') + + keys_to_remove = ['self', 'kwargs'] + for k in keys_to_remove: + args.pop(k, None) + alert(f'{name}: Used args\n{PP(args)}', 'green') + return args + + +def split_labels(labels, label_lengths): + r"""Split concatenated labels into their parts. + + Args: + labels (torch.Tensor): Labels obtained through concatenation. + label_lengths (OrderedDict): Containing order of labels & their lengths. + + Returns: + + """ + assert isinstance(label_lengths, OrderedDict) + start = 0 + outputs = {} + for data_type, length in label_lengths.items(): + end = start + length + if labels.dim() == 5: + outputs[data_type] = labels[:, :, start:end] + elif labels.dim() == 4: + outputs[data_type] = labels[:, start:end] + elif labels.dim() == 3: + outputs[data_type] = labels[start:end] + start = end + return outputs + + +def requires_grad(model, require=True): + r""" Set a model to require gradient or not. + + Args: + model (nn.Module): Neural network model. + require (bool): Whether the network requires gradient or not. + + Returns: + + """ + for p in model.parameters(): + p.requires_grad = require + + +def to_device(data, device): + r"""Move all tensors inside data to device. + + Args: + data (dict, list, or tensor): Input data. + device (str): 'cpu' or 'cuda'. + """ + if isinstance(device, str): + device = torch.device(device) + assert isinstance(device, torch.device) + + if isinstance(data, torch.Tensor): + data = data.to(device, non_blocking=True) + return data + elif isinstance(data, collections.abc.Mapping): + return type(data)({key: to_device(data[key], device) for key in data}) + elif isinstance(data, collections.abc.Sequence) and not isinstance(data, string_classes): + return type(data)([to_device(d, device) for d in data]) + else: + return data + + +def to_cuda(data): + r"""Move all tensors inside data to gpu. + + Args: + data (dict, list, or tensor): Input data. + """ + return to_device(data, 'cuda') + + +def to_cpu(data): + r"""Move all tensors inside data to cpu. + + Args: + data (dict, list, or tensor): Input data. + """ + return to_device(data, 'cpu') + + +def to_half(data): + r"""Move all floats to half. + + Args: + data (dict, list or tensor): Input data. + """ + if isinstance(data, torch.Tensor) and torch.is_floating_point(data): + data = data.half() + return data + elif isinstance(data, collections.abc.Mapping): + return type(data)({key: to_half(data[key]) for key in data}) + elif isinstance(data, collections.abc.Sequence) and not isinstance(data, string_classes): + return type(data)([to_half(d) for d in data]) + else: + return data + + +def to_float(data): + r"""Move all halfs to float. + + Args: + data (dict, list or tensor): Input data. + """ + if isinstance(data, torch.Tensor) and torch.is_floating_point(data): + data = data.float() + return data + elif isinstance(data, collections.abc.Mapping): + return type(data)({key: to_float(data[key]) for key in data}) + elif isinstance(data, collections.abc.Sequence) and not isinstance(data, string_classes): + return type(data)([to_float(d) for d in data]) + else: + return data + + +def slice_tensor(data, start, end): + r"""Slice all tensors from start to end. + Args: + data (dict, list or tensor): Input data. + """ + if isinstance(data, torch.Tensor): + data = data[start:end] + return data + elif isinstance(data, collections.abc.Mapping): + return type(data)({key: slice_tensor(data[key], start, end) for key in data}) + elif isinstance(data, collections.abc.Sequence) and not isinstance(data, string_classes): + return type(data)([slice_tensor(d, start, end) for d in data]) + else: + return data + + +def get_and_setattr(cfg, name, default): + r"""Get attribute with default choice. If attribute does not exist, set it + using the default value. + + Args: + cfg (obj) : Config options. + name (str) : Attribute name. + default (obj) : Default attribute. + + Returns: + (obj) : Desired attribute. + """ + if not hasattr(cfg, name) or name not in cfg.__dict__: + setattr(cfg, name, default) + return getattr(cfg, name) + + +def get_nested_attr(cfg, attr_name, default): + r"""Iteratively try to get the attribute from cfg. If not found, return + default. + + Args: + cfg (obj): Config file. + attr_name (str): Attribute name (e.g. XXX.YYY.ZZZ). + default (obj): Default return value for the attribute. + + Returns: + (obj): Attribute value. + """ + names = attr_name.split('.') + atr = cfg + for name in names: + if not hasattr(atr, name): + return default + atr = getattr(atr, name) + return atr + + +def gradient_norm(model): + r"""Return the gradient norm of model. + + Args: + model (PyTorch module): Your network. + + """ + total_norm = 0 + for p in model.parameters(): + if p.grad is not None: + param_norm = p.grad.norm(2) + total_norm += param_norm.item() ** 2 + return total_norm ** (1. / 2) + + +def random_shift(x, offset=0.05, mode='bilinear', padding_mode='reflection'): + r"""Randomly shift the input tensor. + + Args: + x (4D tensor): The input batch of images. + offset (int): The maximum offset ratio that is between [0, 1]. + The maximum shift is offset * image_size for each direction. + mode (str): The resample mode for 'F.grid_sample'. + padding_mode (str): The padding mode for 'F.grid_sample'. + + Returns: + x (4D tensor) : The randomly shifted image. + """ + assert x.dim() == 4, "Input must be a 4D tensor." + batch_size = x.size(0) + theta = torch.eye(2, 3, device=x.device).unsqueeze(0).repeat( + batch_size, 1, 1) + theta[:, :, 2] = 2 * offset * torch.rand(batch_size, 2) - offset + grid = F.affine_grid(theta, x.size()) + x = F.grid_sample(x, grid, mode=mode, padding_mode=padding_mode, align_corners=False) + return x + + +# def truncated_gaussian(threshold, size, seed=None, device=None): +# r"""Apply the truncated gaussian trick to trade diversity for quality +# +# Args: +# threshold (float): Truncation threshold. +# size (list of integer): Tensor size. +# seed (int): Random seed. +# device: +# """ +# state = None if seed is None else np.random.RandomState(seed) +# values = truncnorm.rvs(-threshold, threshold, +# size=size, random_state=state) +# return torch.tensor(values, device=device).float() + + +def apply_imagenet_normalization(input): + r"""Normalize using ImageNet mean and std. + + Args: + input (4D tensor NxCxHxW): The input images, assuming to be [-1, 1]. + + Returns: + Normalized inputs using the ImageNet normalization. + """ + # normalize the input back to [0, 1] + normalized_input = (input + 1) / 2 + # normalize the input using the ImageNet mean and std + mean = normalized_input.new_tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1) + std = normalized_input.new_tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) + output = (normalized_input - mean) / std + return output + + +def alarm_handler(timeout_period, signum, frame): + # What to do when the process gets stuck. For now, we simply end the process. + error_message = f"Timeout error! More than {timeout_period} seconds have passed since the last iteration. Most " \ + f"likely the process has been stuck due to node failure or PBSS error." + ngc_job_id = os.environ.get('NGC_JOB_ID', None) + if ngc_job_id is not None: + error_message += f" Failed NGC job ID: {ngc_job_id}." + # Let's reserve `wandb.alert` for this purpose. + wandb.alert(title="Timeout error!", text=error_message, level=wandb.AlertLevel.ERROR) + exit() + + +class Timer(object): + + def __init__(self, cfg): + self.cfg = cfg + self.time_iteration = 0 + self.time_epoch = 0 + if is_master(): + # noinspection PyTypeChecker + signal.signal(signal.SIGALRM, functools.partial(alarm_handler, self.cfg.timeout_period)) + + def reset(self): + self.accu_forw_iter_time = 0 + self.accu_loss_iter_time = 0 + self.accu_back_iter_time = 0 + self.accu_step_iter_time = 0 + self.accu_avg_iter_time = 0 + + def _time_before_forward(self): + r"""Record time before applying forward.""" + if self.cfg.speed_benchmark: + torch.cuda.synchronize() + self.forw_time = time.time() + + def _time_before_loss(self): + r"""Record time before computing loss.""" + if self.cfg.speed_benchmark: + torch.cuda.synchronize() + self.loss_time = time.time() + + def _time_before_backward(self): + r"""Record time before applying backward.""" + if self.cfg.speed_benchmark: + torch.cuda.synchronize() + self.back_time = time.time() + + def _time_before_step(self): + r"""Record time before updating the weights""" + if self.cfg.speed_benchmark: + torch.cuda.synchronize() + self.step_time = time.time() + + def _time_before_model_avg(self): + r"""Record time before applying model average.""" + if self.cfg.speed_benchmark: + torch.cuda.synchronize() + self.avg_time = time.time() + + def _time_before_leave_gen(self): + r"""Record forward, backward, loss, and model average time for the network update.""" + if self.cfg.speed_benchmark: + torch.cuda.synchronize() + end_time = time.time() + self.accu_forw_iter_time += self.loss_time - self.forw_time + self.accu_loss_iter_time += self.back_time - self.loss_time + self.accu_back_iter_time += self.step_time - self.back_time + self.accu_step_iter_time += self.avg_time - self.step_time + self.accu_avg_iter_time += end_time - self.avg_time + + def _print_speed_benchmark(self, avg_time): + """Prints the profiling results and resets the timers.""" + print('{:6f}'.format(avg_time)) + print('\tModel FWD time {:6f}'.format(self.accu_forw_iter_time / self.cfg.logging_iter)) + print('\tModel LOS time {:6f}'.format(self.accu_loss_iter_time / self.cfg.logging_iter)) + print('\tModel BCK time {:6f}'.format(self.accu_back_iter_time / self.cfg.logging_iter)) + print('\tModel STP time {:6f}'.format(self.accu_step_iter_time / self.cfg.logging_iter)) + print('\tModel AVG time {:6f}'.format(self.accu_avg_iter_time / self.cfg.logging_iter)) + self.accu_forw_iter_time = 0 + self.accu_loss_iter_time = 0 + self.accu_back_iter_time = 0 + self.accu_step_iter_time = 0 + self.accu_avg_iter_time = 0 + + def checkpoint_tic(self): + # reset timer + self.checkpoint_start_time = time.time() + + def checkpoint_toc(self): + # return time by minutes + return (time.time() - self.checkpoint_start_time) / 60 + + @master_only + def reset_timeout_counter(self): + signal.alarm(self.cfg.timeout_period) diff --git a/neuralangelo-main/imaginaire/utils/set_random_seed.py b/neuralangelo-main/imaginaire/utils/set_random_seed.py new file mode 100644 index 0000000000000000000000000000000000000000..78c71c94f99b3203fad6353e0f7ea73e416ce3fe --- /dev/null +++ b/neuralangelo-main/imaginaire/utils/set_random_seed.py @@ -0,0 +1,36 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import random +import numpy as np +import torch + +from imaginaire.utils.distributed import get_rank +from imaginaire.utils.distributed import master_only_print as print + + +def set_random_seed(seed, by_rank=False): + r"""Set random seeds for everything, including random, numpy, torch.manual_seed, torch.cuda_manual_seed. + torch.cuda.manual_seed_all is not necessary (included in torch.manual_seed) + + Args: + seed (int): Random seed. + by_rank (bool): if true, each gpu will use a different random seed. + """ + if by_rank: + seed += get_rank() + print(f"Using random seed {seed}") + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) # sets seed on the current CPU & all GPUs + torch.cuda.manual_seed(seed) # sets seed on current GPU + # torch.cuda.manual_seed_all(seed) # included in torch.manual_seed diff --git a/neuralangelo-main/imaginaire/utils/termcolor.py b/neuralangelo-main/imaginaire/utils/termcolor.py new file mode 100644 index 0000000000000000000000000000000000000000..77de9062478326f95ac60f3616e88ac83d56b638 --- /dev/null +++ b/neuralangelo-main/imaginaire/utils/termcolor.py @@ -0,0 +1,43 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import pprint + +import termcolor + + +def red(x): return termcolor.colored(str(x), color="red") +def green(x): return termcolor.colored(str(x), color="green") +def blue(x): return termcolor.colored(str(x), color="blue") +def cyan(x): return termcolor.colored(str(x), color="cyan") +def yellow(x): return termcolor.colored(str(x), color="yellow") +def magenta(x): return termcolor.colored(str(x), color="magenta") +def grey(x): return termcolor.colored(str(x), color="grey") + + +COLORS = { + 'red': red, 'green': green, 'blue': blue, 'cyan': cyan, 'yellow': yellow, 'magenta': magenta, 'grey': grey +} + + +def PP(x): + string = pprint.pformat(x, indent=2) + if isinstance(x, dict): + string = '{\n ' + string[1:-1] + '\n}' + return string + + +def alert(x, color='red'): + color = COLORS[color] + print(color('-' * 32)) + print(color(f'* {x}')) + print(color('-' * 32)) diff --git a/neuralangelo-main/imaginaire/utils/visualization.py b/neuralangelo-main/imaginaire/utils/visualization.py new file mode 100644 index 0000000000000000000000000000000000000000..e9e3e8a85fe99dea6e2be7d198108b01727ce395 --- /dev/null +++ b/neuralangelo-main/imaginaire/utils/visualization.py @@ -0,0 +1,41 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import wandb +import torch +import torchvision + +from matplotlib import pyplot as plt +from torchvision.transforms import functional as torchvision_F + + +def wandb_image(images, from_range=(0, 1)): + images = preprocess_image(images, from_range=from_range) + image_grid = torchvision.utils.make_grid(images, nrow=1, pad_value=1) + image_grid = torchvision_F.to_pil_image(image_grid) + wandb_image = wandb.Image(image_grid) + return wandb_image + + +def preprocess_image(images, from_range=(0, 1), cmap="gray"): + min, max = from_range + images = (images - min) / (max - min) + images = images.detach().cpu().float().clamp_(min=0, max=1) + if images.shape[1] == 1: + images = get_heatmap(images[:, 0], cmap=cmap) + return images + + +def get_heatmap(gray, cmap): # [N,H,W] + color = plt.get_cmap(cmap)(gray.numpy()) + color = torch.from_numpy(color[..., :3]).permute(0, 3, 1, 2).float() # [N,3,H,W] + return color diff --git a/neuralangelo-main/neuralangelo.yaml b/neuralangelo-main/neuralangelo.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3609361dbbe5052c49a302e79cb24b5869ef662e --- /dev/null +++ b/neuralangelo-main/neuralangelo.yaml @@ -0,0 +1,26 @@ +# conda env create --file neuralangelo.yaml +name: neuralangelo +channels: + - conda-forge + - pytorch +dependencies: + # general + - gpustat + - gdown + - cudatoolkit-dev + - cmake + # python general + - python=3.8 + - pip + - numpy + - scipy + - ipython + - jupyterlab + - cython + - ninja + - diskcache + # pytorch + - pytorch + - torchvision + - pip: + - -r requirements.txt diff --git a/neuralangelo-main/projects/nerf/configs/ingp_blender.yaml b/neuralangelo-main/projects/nerf/configs/ingp_blender.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8c870c86e5bbb963f4fc18ccb5b2e912eed8171c --- /dev/null +++ b/neuralangelo-main/projects/nerf/configs/ingp_blender.yaml @@ -0,0 +1,41 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. +# ----------------------------------------------------------------------------- + +_parent_: projects/nerf/configs/nerf_blender.yaml + +max_iter: 500000 + +wandb_scalar_iter: 500 +wandb_image_iter: 10000 +validation_iter: 10000 + +model: + type: projects.nerf.models.ingp + mlp: + layers_feat: [null,64,64] + layers_rgb: [null,64,3] + voxel: + levels: + min: 4 + max: 12 + num: 16 + dict_size: 19 + dim: 4 + range: [-5,5] + init_scale: 1e-4 + sample_intvs: 256 + fine_sampling: False + +optim: + type: Adam + params: + lr: 0.01 + sched: + gamma: 1 diff --git a/neuralangelo-main/projects/nerf/configs/nerf_blender.yaml b/neuralangelo-main/projects/nerf/configs/nerf_blender.yaml new file mode 100644 index 0000000000000000000000000000000000000000..241f1f1cc04bee9468fedd632751efe4f2d9826b --- /dev/null +++ b/neuralangelo-main/projects/nerf/configs/nerf_blender.yaml @@ -0,0 +1,101 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. +# ----------------------------------------------------------------------------- + +aws_credentials_file: scripts/s3/pbss_dir.secret +logging_iter: 9999999999999 # disable the printing logger + +max_iter: 500000 + +wandb_scalar_iter: 100 +wandb_image_iter: 1000 +validation_iter: 2000 + +speed_benchmark: False + +checkpoint: + save_to_s3: False + load_from_s3: False + s3_credentials: scripts/s3/pbss_dir.secret + s3_bucket: nerf + save_iter: 10000 + +trainer: + type: projects.nerf.trainers.nerf + ema_config: + enabled: False + load_ema_checkpoint: False + loss_weight: + render: 1.0 + render_fine: 1.0 + init: + type: xavier + amp_config: + enabled: True + +model: + type: projects.nerf.models.nerf + mlp: + layers_feat: [null,256,256,256,256,256,256,256,256] + layers_rgb: [null,128,3] + skip: [4] + posenc: + L_3D: 10 + L_view: 4 + density_activ: softplus + view_dep: True + dist: + param: metric + range: [2,6] + sample_intvs: 64 + sample_stratified: True + fine_sampling: True + sample_intvs_fine: 128 + rand_rays: 1024 + density_noise_reg: + opaque_background: False + camera_ndc: False + +optim: + type: Adam + params: + lr: 0.0005 + betas: [0.9, 0.999] + sched: + iteration_mode: False + type: step + step_size: 20 + gamma: 0.97724 + +data: + type: projects.nerf.datasets.nerf_blender + use_multi_epoch_loader: True + num_workers: 4 + root: datasets/nerf-synthetic/lego + image_size: [400,400] + preload: True + bgcolor: 1 + train: + batch_size: 2 + subset: + val: + batch_size: 2 + subset: 4 + max_viz_samples: 16 + +test_data: + type: projects.nerf.datasets.nerf_blender + num_workers: 4 + root: datasets/nerf-synthetic/lego + image_size: [400,400] + preload: True + bgcolor: 1 + test: + batch_size: 2 + subset: diff --git a/neuralangelo-main/projects/nerf/configs/nerf_llff.yaml b/neuralangelo-main/projects/nerf/configs/nerf_llff.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e58c2432b3ce99f4f0c4d3e18dd9bfcfc25e6ee5 --- /dev/null +++ b/neuralangelo-main/projects/nerf/configs/nerf_llff.yaml @@ -0,0 +1,64 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. +# ----------------------------------------------------------------------------- + +_parent_: projects/nerf/configs/nerf_blender.yaml + +max_iter: 200000 + +wandb_scalar_iter: 50 +wandb_image_iter: 500 +validation_iter: 1000 + +model: + type: projects.nerf.models.nerf + dist: + param: ndc + range: [0,1] + sample_intvs: 64 + fine_sampling: True + sample_intvs_fine: 128 + rand_rays: 1024 + camera_ndc: True + +optim: + type: Adam + params: + lr: 0.0005 + betas: [0.9, 0.999] + sched: + iteration_mode: False + type: step + step_size: 100 + gamma: 0.97724 + +data: + type: projects.nerf.datasets.nerf_llff + use_multi_epoch_loader: True + num_workers: 4 + root: datasets/nerf-llff/fern + image_size: [480,640] + preload: True + val_ratio: 0.1 + train: + batch_size: 2 + subset: + val: + batch_size: 2 + subset: 4 + +test_data: + type: projects.nerf.datasets.nerf_llff + num_workers: 4 + root: datasets/nerf-llff/fern + image_size: [480,640] + preload: True + test: + batch_size: 2 + subset: diff --git a/neuralangelo-main/projects/nerf/datasets/base.py b/neuralangelo-main/projects/nerf/datasets/base.py new file mode 100644 index 0000000000000000000000000000000000000000..4889f5e54b2557f76618fb12017c54d0bfed4a65 --- /dev/null +++ b/neuralangelo-main/projects/nerf/datasets/base.py @@ -0,0 +1,55 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import torch +import tqdm +import threading +import queue + + +class Dataset(torch.utils.data.Dataset): + + def __init__(self, cfg, is_inference=False, is_test=False): + super().__init__() + self.split = "test" if is_test else "val" if is_inference else "train" + + def _preload_worker(self, data_list, load_func, q, lock, idx_tqdm): + # Keep preloading data in parallel. + while True: + idx = q.get() + data_list[idx] = load_func(idx) + with lock: + idx_tqdm.update() + q.task_done() + + def preload_threading(self, load_func, num_workers, data_str="images"): + # Use threading to preload data in parallel. + data_list = [None] * len(self) + q = queue.Queue(maxsize=len(self)) + idx_tqdm = tqdm.tqdm(range(len(self)), desc=f"preloading {data_str} ({self.split})", leave=False) + for i in range(len(self)): + q.put(i) + lock = threading.Lock() + for ti in range(num_workers): + t = threading.Thread(target=self._preload_worker, + args=(data_list, load_func, q, lock, idx_tqdm), daemon=True) + t.start() + q.join() + idx_tqdm.close() + assert all(map(lambda x: x is not None, data_list)) + return data_list + + def __getitem__(self, idx): + raise NotImplementedError + + def __len__(self): + return len(self.list) diff --git a/neuralangelo-main/projects/nerf/datasets/nerf_blender.py b/neuralangelo-main/projects/nerf/datasets/nerf_blender.py new file mode 100644 index 0000000000000000000000000000000000000000..ac5aabb1b8160d8be1db1ba7cce67f46d169746b --- /dev/null +++ b/neuralangelo-main/projects/nerf/datasets/nerf_blender.py @@ -0,0 +1,114 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import json +import numpy as np +import torch +import torchvision.transforms.functional as torchvision_F +from PIL import Image, ImageFile + +from projects.nerf.datasets import base +from projects.nerf.utils import camera + +ImageFile.LOAD_TRUNCATED_IMAGES = True + + +class Dataset(base.Dataset): + + def __init__(self, cfg, is_inference=False, is_test=False): + super().__init__(cfg, is_inference=is_inference, is_test=is_test) + cfg_data = cfg.test_data if self.split == "test" else cfg.data + data_info = cfg_data[self.split] + self.root = cfg_data.root + self.preload = cfg_data.preload + self.bgcolor = cfg_data.bgcolor + self.raw_H, self.raw_W = 800, 800 + self.H, self.W = cfg_data.image_size + meta_fname = f"{cfg_data.root}/transforms_{self.split}.json" + with open(meta_fname) as file: + self.meta = json.load(file) + self.focal = 0.5 * self.raw_W / np.tan(0.5 * self.meta["camera_angle_x"]) + self.list = self.meta["frames"] + # Consider only a subset of data. + if data_info.subset: + self.list = self.list[:data_info.subset] + # Preload dataset if possible. + if cfg_data.preload: + self.images = self.preload_threading(self.get_image, cfg_data.num_workers) + self.cameras = self.preload_threading(self.get_camera, cfg_data.num_workers, data_str="cameras") + + def __getitem__(self, idx): + """Process raw data and return processed data in a dictionary. + + Args: + idx: The index of the sample of the dataset. + Returns: A dictionary containing the data. + idx (scalar): The index of the sample of the dataset. + image (3xHxW tensor): Image with pixel values in [0,1] for supervision. + intr (3x3 tensor): The camera intrinsics of `image`. + pose (3x4 tensor): The camera extrinsics [R,t] of `image`. + """ + # Keep track of sample index for convenience. + sample = dict(idx=idx) + # Get the images. + image = self.images[idx] if self.preload else self.get_image(idx) + image = self.preprocess_image(image) + # Get the cameras (intrinsics and pose). + intr, pose = self.cameras[idx] if self.preload else self.get_camera(idx) + intr, pose = self.preprocess_camera(intr, pose) + # Update the data sample. + sample.update( + image=image, + intr=intr, + pose=pose, + ) + return sample + + def get_image(self, idx): + fpath = self.list[idx]["file_path"][2:] + image_fname = f"{self.root}/{fpath}.png" + image = Image.open(image_fname) + image.load() + return image + + def preprocess_image(self, image): + # Resize the image. + image = image.resize((self.W, self.H)) + image = torchvision_F.to_tensor(image) + # Background masking. + rgb, mask = image[:3], image[3:] + if self.bgcolor is not None: + rgb = rgb * mask + self.bgcolor * (1 - mask) + return rgb + + def get_camera(self, idx): + # Camera intrinsics. + intr = torch.tensor([[self.focal, 0, self.raw_W / 2], + [0, self.focal, self.raw_H / 2], + [0, 0, 1]]).float() + # Camera pose. + pose_raw = torch.tensor(self.list[idx]["transform_matrix"], dtype=torch.float32) + pose = self.parse_raw_camera(pose_raw) + return intr, pose + + def preprocess_camera(self, intr, pose): + # Adjust the intrinsics according to the resized image. + intr = intr.clone() + intr[0] *= self.W / self.raw_W + intr[1] *= self.H / self.raw_H + return intr, pose + + def parse_raw_camera(self, pose_raw): + pose_flip = camera.pose(R=torch.diag(torch.tensor([1, -1, -1]))) + pose = camera.pose.compose([pose_flip, pose_raw[:3]]) + pose = camera.pose.invert(pose) + return pose diff --git a/neuralangelo-main/projects/nerf/datasets/nerf_llff.py b/neuralangelo-main/projects/nerf/datasets/nerf_llff.py new file mode 100644 index 0000000000000000000000000000000000000000..0a7780705cce2495d41114592588678a15c9dcdd --- /dev/null +++ b/neuralangelo-main/projects/nerf/datasets/nerf_llff.py @@ -0,0 +1,140 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import numpy as np +import torch +import torch.nn.functional as torch_F +import torchvision.transforms.functional as torchvision_F +from PIL import Image, ImageFile + +from projects.nerf.datasets import base +from projects.nerf.utils import camera + +ImageFile.LOAD_TRUNCATED_IMAGES = True + + +class Dataset(base.Dataset): + + def __init__(self, cfg, is_inference=False, is_test=False): + super().__init__(cfg, is_inference=is_inference, is_test=is_test) + cfg_data = cfg.test_data if self.split == "test" else cfg.data + data_info = cfg_data[self.split] + self.root = cfg_data.root + self.preload = cfg_data.preload + self.raw_H, self.raw_W = 3024, 4032 + self.H, self.W = cfg_data.image_size + list_fname = f"{cfg_data.root}/images.list" + image_fnames = open(list_fname).read().splitlines() + poses_raw, bounds = self.parse_cameras_and_bounds(cfg_data) + self.list = list(zip(image_fnames, poses_raw, bounds)) + # Manually split train/val subsets. + num_val_split = int(len(self) * cfg_data.val_ratio) + self.list = self.list[:-num_val_split] if self.split == "train" else self.list[-num_val_split:] + # Consider only a subset of data. + if data_info.subset: + self.list = self.list[:data_info.subset] + # Preload dataset if possible. + if cfg_data.preload: + self.images = self.preload_threading(self.get_image, cfg_data.num_workers) + self.cameras = self.preload_threading(self.get_camera, cfg_data.num_workers, data_str="cameras") + + def parse_cameras_and_bounds(self, cfg_data): + fname = f"{cfg_data.root}/poses_bounds.npy" + data = torch.tensor(np.load(fname), dtype=torch.float32) + # Parse cameras (intrinsics and poses). + cam_data = data[:, :-2].view([-1, 3, 5]) # [N,3,5] + poses_raw = cam_data[..., :4] # [N,3,4] + poses_raw[..., 0], poses_raw[..., 1] = poses_raw[..., 1], -poses_raw[..., 0] + raw_H, raw_W, self.focal = cam_data[0, :, -1] + assert self.raw_H == raw_H and self.raw_W == raw_W + # Parse depth bounds. + bounds = data[:, -2:] # [N,2] + scale = 1. / (bounds.min() * 0.75) # Not sure how this was determined? + poses_raw[..., 3] *= scale + bounds *= scale + # Roughly center camera poses. + poses_raw = self.center_camera_poses(poses_raw) + return poses_raw, bounds + + def center_camera_poses(self, poses): + # Compute average pose. + center = poses[..., 3].mean(dim=0) + v1 = torch_F.normalize(poses[..., 1].mean(dim=0), dim=0) + v2 = torch_F.normalize(poses[..., 2].mean(dim=0), dim=0) + v0 = v1.cross(v2) + pose_avg = torch.stack([v0, v1, v2, center], dim=-1)[None] # [1,3,4] + # Apply inverse of averaged pose. + poses = camera.pose.compose([poses, camera.pose.invert(pose_avg)]) + return poses + + def __getitem__(self, idx): + """Process raw data and return processed data in a dictionary. + + Args: + idx: The index of the sample of the dataset. + Returns: A dictionary containing the data. + idx (scalar): The index of the sample of the dataset. + image (3xHxW tensor): Image with pixel values in [0,1] for supervision. + intr (3x3 tensor): The camera intrinsics of `image`. + pose (3x4 tensor): The camera extrinsics [R,t] of `image`. + """ + # Keep track of sample index for convenience. + sample = dict(idx=idx) + # Get the images. + image = self.images[idx] if self.preload else self.get_image(idx) + image = self.preprocess_image(image) + # Get the cameras (intrinsics and pose). + intr, pose = self.cameras[idx] if self.preload else self.get_camera(idx) + intr, pose = self.preprocess_camera(intr, pose) + # Update the data sample. + sample.update( + image=image, + intr=intr, + pose=pose, + ) + return sample + + def get_image(self, idx): + image_fname = f"{self.root}/images/{self.list[idx][0]}" + image = Image.open(image_fname) + image.load() + return image + + def preprocess_image(self, image): + # Resize the image and convert to Pytorch. + image = image.resize((self.W, self.H)) + image = torchvision_F.to_tensor(image) + return image + + def get_camera(self, idx): + # Camera intrinsics. + intr = torch.tensor([[self.focal, 0, self.raw_W / 2], + [0, self.focal, self.raw_H / 2], + [0, 0, 1]]).float() + # Camera pose. + pose_raw = self.list[idx][1] + pose = self.parse_raw_camera(pose_raw) + return intr, pose + + def preprocess_camera(self, intr, pose): + # Adjust the intrinsics according to the resized image. + intr = intr.clone() + intr[0] *= self.W / self.raw_W + intr[1] *= self.H / self.raw_H + return intr, pose + + def parse_raw_camera(self, pose_raw): + pose_flip = camera.pose(R=torch.diag(torch.tensor([1, -1, -1]))) + pose = camera.pose.compose([pose_flip, pose_raw[:3]]) + pose = camera.pose.invert(pose) + pose = camera.pose.compose([pose_flip, pose]) + return pose diff --git a/neuralangelo-main/projects/nerf/models/ingp.py b/neuralangelo-main/projects/nerf/models/ingp.py new file mode 100644 index 0000000000000000000000000000000000000000..e3726ae6a2db0bb16ce850d27931f663f77f68ba --- /dev/null +++ b/neuralangelo-main/projects/nerf/models/ingp.py @@ -0,0 +1,77 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import numpy as np +import torch +import tinycudann as tcnn + +from projects.nerf.models import nerf + + +class Model(nerf.Model): + + def __init__(self, cfg_model, cfg_data): + super().__init__(cfg_model, cfg_data) + self.fine_sampling = False + self.density_reg = cfg_model.density_noise_reg + # Define models. + self.nerf = InstantNGP(cfg_model) + + +class InstantNGP(nerf.NeRF): + + def __init__(self, cfg_model): + self.voxel = cfg_model.voxel + super().__init__(cfg_model) + + def set_input_dims(self, cfg_model): + # Define the input encoding dimensions. + self.input_3D_dim = 3 + cfg_model.voxel.dim * cfg_model.voxel.levels.num + self.input_view_dim = 3 if cfg_model.view_dep else None + + def build_model(self, cfg_model): + super().build_model(cfg_model) + # Build the tcnn hash grid. + l_min, l_max = self.voxel.levels.min, self.voxel.levels.max + r_min, r_max = 2 ** l_min, 2 ** l_max + num_levels = self.voxel.levels.num + growth_rate = np.exp((np.log(r_max) - np.log(r_min)) / (num_levels - 1)) + config = dict( + otype="HashGrid", + n_levels=cfg_model.voxel.levels.num, + n_features_per_level=cfg_model.voxel.dim, + log2_hashmap_size=cfg_model.voxel.dict_size, + base_resolution=2 ** cfg_model.voxel.levels.min, + per_level_scale=growth_rate, + ) + self.tiny_cuda_encoding = tcnn.Encoding(3, config) + # Compute resolutions of all levels. + self.resolutions = [] + for lv in range(0, num_levels): + size = np.floor(r_min * growth_rate ** lv).astype(int) + 1 + self.resolutions.append(size) + + def forward(self, points_3D, ray_unit, density_reg=None): + return super().forward(points_3D, ray_unit, density_reg) + + def _encode_3D(self, points_3D): + # Tri-linear interpolate the corresponding embeddings from the dictionary. + vol_min, vol_max = self.voxel.range + points_3D_normalized = (points_3D - vol_min) / (vol_max - vol_min) # Normalize to [0,1]. + tcnn_input = points_3D_normalized.view(-1, 3) + tcnn_output = self.tiny_cuda_encoding(tcnn_input) + points_enc = tcnn_output.view(*points_3D_normalized.shape[:-1], tcnn_output.shape[-1]) + points_enc = torch.cat([points_enc, points_3D], dim=-1) # [B,R,N,LD+3] + return points_enc + + def _encode_view(self, ray_unit): + return ray_unit diff --git a/neuralangelo-main/projects/nerf/models/nerf.py b/neuralangelo-main/projects/nerf/models/nerf.py new file mode 100644 index 0000000000000000000000000000000000000000..ac25b6115d9561b31cdae30d25828fcfcfbf48a3 --- /dev/null +++ b/neuralangelo-main/projects/nerf/models/nerf.py @@ -0,0 +1,251 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import torch +import torch.nn.functional as torch_F +import copy +from functools import partial +from collections import defaultdict + +from imaginaire.models.base import Model as BaseModel +from projects.nerf.utils import camera, render, nerf_util + + +class Model(BaseModel): + + def __init__(self, cfg_model, cfg_data): + super().__init__(cfg_model, cfg_data) + self.num_rays = cfg_model.rand_rays + self.image_size = cfg_data.image_size + self.fine_sampling = cfg_model.fine_sampling + self.stratified = cfg_model.sample_stratified + self.density_reg = cfg_model.density_noise_reg + self.opaque_background = cfg_model.opaque_background + self.bgcolor = getattr(cfg_data, "bgcolor", 1.) + # Define models. + self.nerf = NeRF(cfg_model) + if self.fine_sampling: + self.nerf_fine = NeRF(cfg_model) + # Define functions. + self.ray_generator = partial(nerf_util.ray_generator, + image_size=cfg_data.image_size, + camera_ndc=cfg_model.camera_ndc, + num_rays=cfg_model.rand_rays) + self.sample_dists = partial(nerf_util.sample_dists, + dist_range=cfg_model.dist.range, + intvs=cfg_model.sample_intvs) + self.sample_dists_from_pdf = partial(nerf_util.sample_dists_from_pdf, + intvs_fine=cfg_model.sample_intvs_fine) + self.reparametrize_dist = partial(nerf_util.reparametrize_dist, + param_type=cfg_model.dist.param) + self.get_inverse_depth = partial(nerf_util.get_inverse_depth, + camera_ndc=cfg_model.camera_ndc) + self.to_full_image = lambda vec, image_size=cfg_data.image_size: \ + vec.moveaxis(1, -1).unflatten(dim=-1, sizes=image_size) + + def forward(self, data): + # Randomly sample and render the pixels. + ray_idx = self._sample_random_rays(data) + output = self.render_pixels(data["pose"], data["intr"], full_image=False, ray_idx=ray_idx, + stratified=self.stratified, density_reg=self.density_reg) + output.update(ray_idx=ray_idx) # [B,R] + return output + + def _sample_random_rays(self, data): + batch_size = len(data["pose"]) + num_pixels = self.image_size[0] * self.image_size[1] + ray_idx = torch.rand(batch_size, num_pixels, device=data["pose"].device).argsort(dim=1)[:, :self.num_rays] + return ray_idx # [B,R] + + @torch.no_grad() + def inference(self, data): + self.eval() + # Render the full images. + output = self.render_image(data["pose"], data["intr"], stratified=False) # [B,N,C] + # Get full rendered RGB and depth images. + inv_depth = self.get_inverse_depth(output["depth"], opacity=output["opacity"]) + output.update( + rgb_map=self.to_full_image(output["rgb"]), # [B,3,H,W] + inv_depth_map=self.to_full_image(inv_depth), # [B,1,H,W] + ) + if self.fine_sampling: + inv_depth_fine = self.get_inverse_depth(output["depth_fine"], opacity=output["opacity_fine"]) + output.update( + rgb_map_fine=self.to_full_image(output["rgb_fine"]), # [B,3,H,W] + inv_depth_map_fine=self.to_full_image(inv_depth_fine), # [B,1,H,W] + ) + return output + + def render_image(self, pose, intr, stratified=False): + """ Render the rays given the camera intrinsics and poses. + Args: + pose (tensor [batch,3,4]): Camera poses ([R,t]). + intr (tensor [batch,3,3]): Camera intrinsics. + stratified (bool): Whether to stratify the depth sampling. + Returns: + output: A dictionary containing the outputs. + """ + output = defaultdict(list) + for center, ray, _ in self.ray_generator(pose, intr, full_image=True): + ray_unit = torch_F.normalize(ray, dim=-1) # [B,R,3] + output_batch = self.render_rays(center, ray_unit, stratified=stratified) + if not self.training: + depth = output_batch["dist"] / ray.norm(dim=-1, keepdim=True) + output_batch.update(depth=depth) + if self.fine_sampling: + depth_fine = output_batch["dist_fine"] / ray.norm(dim=-1, keepdim=True) + output_batch.update(depth_fine=depth_fine) + for key, value in output_batch.items(): + output[key].append(value.detach()) + # Concat each item (list) in output into one tensor. Concatenate along the ray dimension (1) + for key, value in output.items(): + output[key] = torch.cat(value, dim=1) + return output + + def render_pixels(self, pose, intr, full_image=False, ray_idx=None, stratified=False, density_reg=None): + center, ray = camera.get_center_and_ray(pose, intr, self.image_size) # [B,HW,3] + center = nerf_util.slice_by_ray_idx(center, ray_idx) # [B,R,3] + ray = nerf_util.slice_by_ray_idx(ray, ray_idx) # [B,R,3] + ray_unit = torch_F.normalize(ray, dim=-1) # [B,R,3] + output = self.render_rays(center, ray_unit, stratified=stratified, density_reg=density_reg) + return output + + def render_rays(self, center, ray_unit, sample_idx=None, stratified=False, density_reg=None): + with torch.no_grad(): + dists = self.sample_dists(ray_unit.shape[:2], stratified=stratified) # [B,R,N,1] + dists = self.reparametrize_dist(dists) # [B,R,N,1] + points = camera.get_3D_points_from_dist(center, ray_unit, dists) # [B,R,N,3] + rays_unit = ray_unit[..., None, :].expand_as(points) # [B,R,N,3] + rgbs, densities = self.nerf.forward(points, rays_unit, density_reg=density_reg) + weights = render.volume_rendering_weights_dist(densities, dists) # [B,R,N,1] + opacity = render.composite(1., weights) # [B,R,1] + rgb = render.composite(rgbs, weights) # [B,R,3] + if self.opaque_background: + rgb = rgb + self.bgcolor * (1 - opacity) + dist = render.composite(dists, weights) # [B,R,1] + # Collect output. + output = dict( + rgb=rgb, # [B,R,3] + dist=dist, # [B,R,1] + opacity=opacity, # [B,R,1] + ) + if self.fine_sampling: + with torch.no_grad(): + # Resample depth according to coarse empirical distribution. + dists_mid = 0.5 * (dists[..., :-1, :] + dists[..., 1:, :]) # [B,R,N-1,1] + dists_fine = self.sample_dists_from_pdf(dists_mid, weights=weights[..., 1:-1, 0]) # [B,R,Nf,1] + dists = torch.cat([dists, dists_fine], dim=2) # [B,R,N+Nf,1] + dists = dists.sort(dim=2).values + points = camera.get_3D_points_from_dist(center, ray_unit, dists) # [B,R,N,3] + rays_unit = ray_unit[..., None, :].expand_as(points) # [B,R,N,3] + rgbs, densities = self.nerf_fine.forward(points, rays_unit, density_reg=density_reg) + weights = render.volume_rendering_weights_dist(densities, dists) # [B,R,N,1] + opacity = render.composite(1., weights) # [B,R,1] + rgb = render.composite(rgbs, weights) # [B,R,3] + if self.opaque_background: + rgb = rgb + self.bgcolor * (1 - opacity) + dist = render.composite(dists, weights) # [B,R,1] + # Collect output. + output.update( + rgb_fine=rgb, # [B,R,3] + dist_fine=dist, # [B,R,1] + opacity_fine=opacity, # [B,R,1] + ) + return output + + +class NeRF(torch.nn.Module): + + def __init__(self, cfg_model): + super().__init__() + self.view_dep = cfg_model.view_dep + self.posenc = cfg_model.posenc + # Define learnable parameters. + self.set_input_dims(cfg_model) + self.build_model(cfg_model) + + def set_input_dims(self, cfg_model): + # Define the input encoding dimensions. + self.input_3D_dim = 3 + 6 * cfg_model.posenc.L_3D if cfg_model.posenc.L_3D else 3 + if cfg_model.view_dep: + self.input_view_dim = 3 + 6 * cfg_model.posenc.L_view if cfg_model.posenc.L_view else 3 + else: + self.input_view_dim = None + + def build_model(self, cfg_model): + # Point-wise feature. + layers_feat = copy.copy(cfg_model.mlp.layers_feat) + layers_feat[0] = self.input_3D_dim + layers_feat[-1] += 1 # For predicting the volume density. + self.mlp_feat = nerf_util.MLPwithSkipConnection(layers_feat, skip_connection=cfg_model.mlp.skip) + # RGB prediction. + layers_rgb = copy.copy(cfg_model.mlp.layers_rgb) + layers_rgb[0] = cfg_model.mlp.layers_feat[-1] + (self.input_view_dim if cfg_model.view_dep else 0) + self.mlp_rgb = nerf_util.MLPwithSkipConnection(layers_rgb) + self.density_activ = dict( + relu=torch_F.relu, + relu_=torch_F.relu_, + abs=torch.abs, + abs_=torch.abs_, + sigmoid=torch.sigmoid, + sigmoid_=torch.sigmoid_, + exp=torch.exp, + exp_=torch.exp_, + softplus=torch_F.softplus, + identity=lambda x: x, + )[cfg_model.density_activ] + + def forward(self, points_3D, ray_unit, density_reg=None): + """ Forward pass the NeRF (MLP). + Args: + points_3D (tensor [batch,...,3]): 3D points in world space. + ray_unit (tensor [batch,...,3]): Unit ray direction in world space. + density_reg (float or None): Density regularization. + Returns: + rgb (tensor [batch,...,3]): Predicted RGB values in [0,1]. + density (tensor [batch,...,3]): Predicted volume density values. + """ + density, feat = self.get_density(points_3D, density_reg=density_reg) + rgb = self.get_color(feat, ray_unit) + return rgb, density + + def get_density(self, points_3D, density_reg=None): + points_enc = self._encode_3D(points_3D) + out = self.mlp_feat(points_enc) + density, feat = out[..., 0], out[..., 1:].relu_() + if density_reg is not None: + density = density + torch.randn_like(density) * density_reg + density = self.density_activ(density) + return density, feat + + def get_color(self, feat, ray_unit=None): + if self.view_dep: + ray_enc = self._encode_view(ray_unit) + feat = torch.cat([feat, ray_enc], dim=-1) + rgb = self.mlp_rgb(feat).sigmoid_() # [B,...,3] + return rgb + + def _encode_3D(self, points_3D): + if self.posenc.L_3D: + points_enc = nerf_util.positional_encoding(points_3D, num_freq_bases=self.posenc.L_3D) + points_enc = torch.cat([points_3D, points_enc], dim=-1) # [B,...,6L+3] + else: + points_enc = points_3D + return points_enc + + def _encode_view(self, ray_unit): + if self.posenc.L_view: + ray_enc = nerf_util.positional_encoding(ray_unit, num_freq_bases=self.posenc.L_view) + ray_enc = torch.cat([ray_unit, ray_enc], dim=-1) # [B,...,6L+3] + else: + ray_enc = ray_unit + return ray_enc diff --git a/neuralangelo-main/projects/nerf/trainers/base.py b/neuralangelo-main/projects/nerf/trainers/base.py new file mode 100644 index 0000000000000000000000000000000000000000..b63ac650961256f7aa4d89dd0179e00c1fd1c446 --- /dev/null +++ b/neuralangelo-main/projects/nerf/trainers/base.py @@ -0,0 +1,159 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import torch +import wandb +from imaginaire.trainers.base import BaseTrainer +from imaginaire.utils.distributed import is_master, master_only +from tqdm import tqdm + +from projects.nerf.utils.misc import collate_test_data_batches, get_unique_test_data, trim_test_samples + + +class BaseTrainer(BaseTrainer): + """ + A customized BaseTrainer. + """ + + def __init__(self, cfg, is_inference=True, seed=0): + super().__init__(cfg, is_inference=is_inference, seed=seed) + self.metrics = dict() + # The below configs should be properly overridden. + cfg.setdefault("wandb_scalar_iter", 9999999999999) + cfg.setdefault("wandb_image_iter", 9999999999999) + cfg.setdefault("validation_epoch", 9999999999999) + cfg.setdefault("validation_iter", 9999999999999) + + def init_losses(self, cfg): + super().init_losses(cfg) + self.weights = {key: value for key, value in cfg.trainer.loss_weight.items() if value} + + def _end_of_iteration(self, data, current_epoch, current_iteration): + # Log to wandb. + if current_iteration % self.cfg.wandb_scalar_iter == 0: + # Compute the elapsed time (as in the original base trainer). + self.timer.time_iteration = self.elapsed_iteration_time / self.cfg.wandb_scalar_iter + self.elapsed_iteration_time = 0 + # Log scalars. + self.log_wandb_scalars(data, mode="train") + # Exit if the training loss has gone to NaN/inf. + if is_master() and self.losses["total"].isnan(): + self.finalize(self.cfg) + raise ValueError("Training loss has gone to NaN!!!") + if is_master() and self.losses["total"].isinf(): + self.finalize(self.cfg) + raise ValueError("Training loss has gone to infinity!!!") + if current_iteration % self.cfg.wandb_image_iter == 0: + self.log_wandb_images(data, mode="train") + # Run validation on val set. + if current_iteration % self.cfg.validation_iter == 0: + data_all = self.test(self.eval_data_loader, mode="val") + # Log the results to W&B. + if is_master(): + self.log_wandb_scalars(data_all, mode="val") + self.log_wandb_images(data_all, mode="val", max_samples=self.cfg.data.val.max_viz_samples) + + def _end_of_epoch(self, data, current_epoch, current_iteration): + # Run validation on val set. + if current_epoch % self.cfg.validation_epoch == 0: + data_all = self.test(self.eval_data_loader, mode="val") + # Log the results to W&B. + if is_master(): + self.log_wandb_scalars(data_all, mode="val") + self.log_wandb_images(data_all, mode="val", max_samples=self.cfg.data.val.max_viz_samples) + + @master_only + def log_wandb_scalars(self, data, mode=None): + scalars = dict() + # Log scalars (basic info & losses). + if mode == "train": + scalars.update({"optim/lr": self.sched.get_last_lr()[0]}) + scalars.update({"time/iteration": self.timer.time_iteration}) + scalars.update({"time/epoch": self.timer.time_epoch}) + scalars.update({f"{mode}/loss/{key}": value for key, value in self.losses.items()}) + scalars.update(iteration=self.current_iteration, epoch=self.current_epoch) + wandb.log(scalars, step=self.current_iteration) + + @master_only + def log_wandb_images(self, data, mode=None, max_samples=None): + trim_test_samples(data, max_samples=max_samples) + + def model_forward(self, data): + # Model forward. + output = self.model(data) # data = self.model(data) will not return the same data in the case of DDP. + data.update(output) + # Compute loss. + self.timer._time_before_loss() + self._compute_loss(data, mode="train") + total_loss = self._get_total_loss() + return total_loss + + def _compute_loss(self, data, mode=None): + raise NotImplementedError + + def train(self, cfg, data_loader, single_gpu=False, profile=False, show_pbar=False): + self.current_epoch = self.checkpointer.resume_epoch or self.current_epoch + self.current_iteration = self.checkpointer.resume_iteration or self.current_iteration + if ((self.current_epoch % self.cfg.validation_epoch == 0 or + self.current_iteration % self.cfg.validation_iter == 0)): + # Do an initial validation. + data_all = self.test(self.eval_data_loader, mode="val", show_pbar=show_pbar) + # Log the results to W&B. + if is_master(): + self.log_wandb_scalars(data_all, mode="val") + self.log_wandb_images(data_all, mode="val", max_samples=self.cfg.data.val.max_viz_samples) + # Train. + super().train(cfg, data_loader, single_gpu, profile, show_pbar) + + @torch.no_grad() + def test(self, data_loader, output_dir=None, inference_args=None, mode="test", show_pbar=False): + """The evaluation/inference engine. + Args: + data_loader: The data loader. + output_dir: Output directory to dump the test results. + inference_args: (unused) + mode: Evaluation mode {"val", "test"}. Can be other modes, but will only gather the data. + Returns: + data_all: A dictionary of all the data. + """ + if self.cfg.trainer.ema_config.enabled: + model = self.model.module.averaged_model + else: + model = self.model.module + model.eval() + if show_pbar: + data_loader = tqdm(data_loader, desc="Evaluating", leave=False) + data_batches = [] + for it, data in enumerate(data_loader): + data = self.start_of_iteration(data, current_iteration=self.current_iteration) + output = model.inference(data) + data.update(output) + data_batches.append(data) + # Aggregate the data from all devices and process the results. + data_gather = collate_test_data_batches(data_batches) + # Only the master process should process the results; slaves will just return. + if is_master(): + data_all = get_unique_test_data(data_gather, data_gather["idx"]) + tqdm.write(f"Evaluating with {len(data_all['idx'])} samples.") + # Validate/test. + if mode == "val": + self._compute_loss(data_all, mode=mode) + _ = self._get_total_loss() + if mode == "test": + # Dump the test results for postprocessing. + self.dump_test_results(data_all, output_dir) + return data_all + else: + return + + def dump_test_results(self, data_all, output_dir): + raise NotImplementedError diff --git a/neuralangelo-main/projects/nerf/trainers/nerf.py b/neuralangelo-main/projects/nerf/trainers/nerf.py new file mode 100644 index 0000000000000000000000000000000000000000..d32b8c9868febbaa98e4d549a2a6e2f261375ea5 --- /dev/null +++ b/neuralangelo-main/projects/nerf/trainers/nerf.py @@ -0,0 +1,109 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import torch +import torch.nn.functional as torch_F +import wandb +import skvideo.io + +from imaginaire.utils.distributed import master_only +from projects.nerf.trainers.base import BaseTrainer +from imaginaire.utils.visualization import wandb_image, preprocess_image + + +class Trainer(BaseTrainer): + + def __init__(self, cfg, is_inference=True, seed=0): + super().__init__(cfg, is_inference=is_inference, seed=seed) + self.batch_idx, _ = torch.meshgrid(torch.arange(cfg.data.train.batch_size), + torch.arange(cfg.model.rand_rays), indexing="ij") # [B,R] + self.batch_idx = self.batch_idx.cuda() + + def _init_loss(self, cfg): + self.criteria["render"] = self.criteria["render_fine"] = torch.nn.MSELoss() + + def _compute_loss(self, data, mode=None): + if mode == "train": + # Extract the corresponding sampled rays. + batch_size = len(data["image"]) + image_vec = data["image"].permute(0, 2, 3, 1).view(batch_size, -1, 3) # [B,HW,3] + image_sampled = image_vec[self.batch_idx, data["ray_idx"]] # [B,R,3] + # Compute loss only on randomly sampled rays. + self.losses["render"] = self.criteria["render"](data["rgb"], image_sampled) + self.metrics["psnr"] = -10 * torch_F.mse_loss(data["rgb"], image_sampled).log10() + if self.cfg.model.fine_sampling: + self.losses["render_fine"] = self.criteria["render_fine"](data["rgb_fine"], image_sampled) + self.metrics["psnr_fine"] = -10 * torch_F.mse_loss(data["rgb_fine"], image_sampled).log10() + else: + # Compute loss on the entire image. + self.losses["render"] = self.criteria["render"](data["rgb_map"], data["image"]) + self.metrics["psnr"] = -10 * torch_F.mse_loss(data["rgb_map"], data["image"]).log10() + if self.cfg.model.fine_sampling: + self.losses["render_fine"] = self.criteria["render_fine"](data["rgb_map_fine"], data["image"]) + self.metrics["psnr_fine"] = -10 * torch_F.mse_loss(data["rgb_map_fine"], data["image"]).log10() + + @master_only + def log_wandb_scalars(self, data, mode=None): + super().log_wandb_scalars(data, mode=mode) + scalars = {f"{mode}/PSNR/nerf": self.metrics["psnr"].detach()} + if "render_fine" in self.losses: + scalars.update({f"{mode}/PSNR/nerf_fine": self.metrics["psnr_fine"].detach()}) + wandb.log(scalars, step=self.current_iteration) + + @master_only + def log_wandb_images(self, data, mode=None, max_samples=None): + super().log_wandb_images(data, mode=mode, max_samples=max_samples) + images = {f"{mode}/image_target": wandb_image(data["image"])} + if mode == "val": + images_error = (data["rgb_map"] - data["image"]).abs() + images.update({ + f"{mode}/images": wandb_image(data["rgb_map"]), + f"{mode}/images_error": wandb_image(images_error), + f"{mode}/inv_depth": wandb_image(data["inv_depth_map"]), + }) + if self.cfg.model.fine_sampling: + images_error_fine = (data["rgb_map_fine"] - data["image"]).abs() + images.update({ + f"{mode}/images_fine": wandb_image(data["rgb_map_fine"]), + f"{mode}/images_error_fine": wandb_image(images_error_fine), + f"{mode}/inv_depth_fine": wandb_image(data["inv_depth_map_fine"]), + }) + images.update({"iteration": self.current_iteration}) + images.update({"epoch": self.current_epoch}) + wandb.log(images, step=self.current_iteration) + + def dump_test_results(self, data_all, output_dir): + results = dict( + images_target=preprocess_image(data_all["images_target"]), + image=preprocess_image(data_all["rgb_map"]), + inv_depth=preprocess_image(data_all["inv_depth_map"]), + ) + if self.cfg.model.fine_sampling: + results.update( + image_fine=preprocess_image(data_all["rgb_map_fine"]), + inv_depth_fine=preprocess_image(data_all["inv_depth_map_fine"]), + ) + # Write results as videos. + inputdict, outputdict = self._get_ffmpeg_dicts() + for key, image_list in results.items(): + print(f"writing video ({key})...") + video_fname = f"{output_dir}/{key}.mp4" + video_writer = skvideo.io.FFmpegWriter(video_fname, inputdict=inputdict, outputdict=outputdict) + for image in image_list: + image = (image * 255).byte().permute(1, 2, 0).numpy() + video_writer.writeFrame(image) + video_writer.close() + + def _get_ffmpeg_dicts(self): + inputdict = {"-r": str(30)} + outputdict = {"-crf": str(10), "-pix_fmt": "yuv420p"} + return inputdict, outputdict diff --git a/neuralangelo-main/projects/nerf/utils/camera.py b/neuralangelo-main/projects/nerf/utils/camera.py new file mode 100644 index 0000000000000000000000000000000000000000..c1a640b04ee563745fe1dc96ccb97f7dac28e7fb --- /dev/null +++ b/neuralangelo-main/projects/nerf/utils/camera.py @@ -0,0 +1,514 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import numpy as np +import torch + + +class Pose(): + """ + A class of operations on camera poses (PyTorch tensors with shape [...,3,4]). + Each [3,4] camera pose takes the form of [R|t]. + """ + + def __call__(self, R=None, t=None): + # Construct a camera pose from the given R and/or t. + assert R is not None or t is not None + if R is None: + if not isinstance(t, torch.Tensor): + t = torch.tensor(t) + R = torch.eye(3, device=t.device).repeat(*t.shape[:-1], 1, 1) + elif t is None: + if not isinstance(R, torch.Tensor): + R = torch.tensor(R) + t = torch.zeros(R.shape[:-1], device=R.device) + else: + if not isinstance(R, torch.Tensor): + R = torch.tensor(R) + if not isinstance(t, torch.Tensor): + t = torch.tensor(t) + assert R.shape[:-1] == t.shape and R.shape[-2:] == (3, 3) + R = R.float() + t = t.float() + pose = torch.cat([R, t[..., None]], dim=-1) # [...,3,4] + assert pose.shape[-2:] == (3, 4) + return pose + + def invert(self, pose, use_inverse=False): + # Invert a camera pose. + R, t = pose[..., :3], pose[..., 3:] + R_inv = R.inverse() if use_inverse else R.transpose(-1, -2) + t_inv = (-R_inv @ t)[..., 0] + pose_inv = self(R=R_inv, t=t_inv) + return pose_inv + + def compose(self, pose_list): + # Compose a sequence of poses together. + # pose_new(x) = poseN o ... o pose2 o pose1(x) + pose_new = pose_list[0] + for pose in pose_list[1:]: + pose_new = self.compose_pair(pose_new, pose) + return pose_new + + def compose_pair(self, pose_a, pose_b): + # pose_new(x) = pose_b o pose_a(x) + R_a, t_a = pose_a[..., :3], pose_a[..., 3:] + R_b, t_b = pose_b[..., :3], pose_b[..., 3:] + R_new = R_b @ R_a + t_new = (R_b @ t_a + t_b)[..., 0] + pose_new = self(R=R_new, t=t_new) + return pose_new + + def scale_center(self, pose, scale): + """Scale the camera center from the origin. + 0 = R@c+t --> c = -R^T@t (camera center in world coordinates) + 0 = R@(sc)+t' --> t' = -R@(sc) = -R@(-R^T@st) = st + """ + R, t = pose[..., :3], pose[..., 3:] + pose_new = torch.cat([R, t * scale], dim=-1) + return pose_new + + def interpolate(self, pose_a, pose_b, alpha): + """Interpolate between two poses with Slerp. + Args: + pose_a (tensor [...,3,4]): Pose at time t=0. + pose_b (tensor [...,3,4]): Pose at time t=1. + alpha (tensor [...,1]): Interpolation parameter. + Returns: + pose (tensor [...,3,4]): Pose at time t. + """ + R_a, t_a = pose_a[..., :3], pose_a[..., 3:] + R_b, t_b = pose_b[..., :3], pose_b[..., 3:] + q_a = quaternion.R_to_q(R_a) # [...,4] + q_b = quaternion.R_to_q(R_b) # [...,4] + q_intp = quaternion.interpolate(q_a, q_b, alpha) # [...,4] + R_intp = quaternion.q_to_R(q_intp) # [...,3,3] + t_intp = (1 - alpha) * t_a + alpha * t_b # [...,3] + pose_intp = torch.cat([R_intp, t_intp], dim=-1) # [...,3,4] + return pose_intp + + +class Lie(): + """ + Lie algebra for SO(3) and SE(3) operations in PyTorch. + """ + + def so3_to_SO3(self, w): # [..., 3] + wx = self.skew_symmetric(w) + theta = w.norm(dim=-1)[..., None, None] + eye = torch.eye(3, device=w.device, dtype=torch.float32) + A = self.taylor_A(theta) + B = self.taylor_B(theta) + R = eye + A * wx + B * wx @ wx + return R + + def SO3_to_so3(self, R, eps=1e-7): # [..., 3, 3] + trace = R[..., 0, 0] + R[..., 1, 1] + R[..., 2, 2] + theta = ((trace - 1) / 2).clamp(-1 + eps, 1 - eps).acos_()[ + ..., None, None] % np.pi # ln(R) will explode if theta==pi + lnR = 1 / (2 * self.taylor_A(theta) + 1e-8) * (R - R.transpose(-2, -1)) # FIXME: wei-chiu finds it weird + w0, w1, w2 = lnR[..., 2, 1], lnR[..., 0, 2], lnR[..., 1, 0] + w = torch.stack([w0, w1, w2], dim=-1) + return w + + def se3_to_SE3(self, wu): # [...,3] + w, u = wu.split([3, 3], dim=-1) + wx = self.skew_symmetric(w) + theta = w.norm(dim=-1)[..., None, None] + eye = torch.eye(3, device=w.device, dtype=torch.float32) + A = self.taylor_A(theta) + B = self.taylor_B(theta) + C = self.taylor_C(theta) + R = eye + A * wx + B * wx @ wx + V = eye + B * wx + C * wx @ wx + Rt = torch.cat([R, (V @ u[..., None])], dim=-1) + return Rt + + def SE3_to_se3(self, Rt, eps=1e-8): # [...,3,4] + R, t = Rt.split([3, 1], dim=-1) + w = self.SO3_to_so3(R) + wx = self.skew_symmetric(w) + theta = w.norm(dim=-1)[..., None, None] + eye = torch.eye(3, device=w.device, dtype=torch.float32) + A = self.taylor_A(theta) + B = self.taylor_B(theta) + invV = eye - 0.5 * wx + (1 - A / (2 * B)) / (theta ** 2 + eps) * wx @ wx + u = (invV @ t)[..., 0] + wu = torch.cat([w, u], dim=-1) + return wu + + def skew_symmetric(self, w): + w0, w1, w2 = w.unbind(dim=-1) + zero = torch.zeros_like(w0) + wx = torch.stack([torch.stack([zero, -w2, w1], dim=-1), + torch.stack([w2, zero, -w0], dim=-1), + torch.stack([-w1, w0, zero], dim=-1)], dim=-2) + return wx + + def taylor_A(self, x, nth=10): + # Taylor expansion of sin(x)/x. + ans = torch.zeros_like(x) + denom = 1. + for i in range(nth + 1): + if i > 0: + denom *= (2 * i) * (2 * i + 1) + ans = ans + (-1) ** i * x ** (2 * i) / denom + return ans + + def taylor_B(self, x, nth=10): + # Taylor expansion of (1-cos(x))/x**2. + ans = torch.zeros_like(x) + denom = 1. + for i in range(nth + 1): + denom *= (2 * i + 1) * (2 * i + 2) + ans = ans + (-1) ** i * x ** (2 * i) / denom + return ans + + def taylor_C(self, x, nth=10): + # Taylor expansion of (x-sin(x))/x**3. + ans = torch.zeros_like(x) + denom = 1. + for i in range(nth + 1): + denom *= (2 * i + 2) * (2 * i + 3) + ans = ans + (-1) ** i * x ** (2 * i) / denom + return ans + + +class Quaternion(): + + def q_to_R(self, q): # [...,4] + # https://en.wikipedia.org/wiki/Rotation_matrix#Quaternion + qa, qb, qc, qd = q.unbind(dim=-1) + R = torch.stack( + [torch.stack([1 - 2 * (qc ** 2 + qd ** 2), 2 * (qb * qc - qa * qd), 2 * (qa * qc + qb * qd)], dim=-1), + torch.stack([2 * (qb * qc + qa * qd), 1 - 2 * (qb ** 2 + qd ** 2), 2 * (qc * qd - qa * qb)], dim=-1), + torch.stack([2 * (qb * qd - qa * qc), 2 * (qa * qb + qc * qd), 1 - 2 * (qb ** 2 + qc ** 2)], dim=-1)], + dim=-2) + return R + + def R_to_q(self, R, eps=1e-6): # [...,3,3] + # https://en.wikipedia.org/wiki/Rotation_matrix#Quaternion + row0, row1, row2 = R.unbind(dim=-2) + R00, R01, R02 = row0.unbind(dim=-1) + R10, R11, R12 = row1.unbind(dim=-1) + R20, R21, R22 = row2.unbind(dim=-1) + t = R[..., 0, 0] + R[..., 1, 1] + R[..., 2, 2] + r = (1 + t + eps).sqrt() + qa = 0.5 * r + qb = (R21 - R12).sign() * 0.5 * (1 + R00 - R11 - R22 + eps).sqrt() + qc = (R02 - R20).sign() * 0.5 * (1 - R00 + R11 - R22 + eps).sqrt() + qd = (R10 - R01).sign() * 0.5 * (1 - R00 - R11 + R22 + eps).sqrt() + q = torch.stack([qa, qb, qc, qd], dim=-1) + return q + + def invert(self, q): # [...,4] + qa, qb, qc, qd = q.unbind(dim=-1) + norm = q.norm(dim=-1, keepdim=True) + q_inv = torch.stack([qa, -qb, -qc, -qd], dim=-1) / norm ** 2 + return q_inv + + def product(self, q1, q2): # [...,4] + q1a, q1b, q1c, q1d = q1.unbind(dim=-1) + q2a, q2b, q2c, q2d = q2.unbind(dim=-1) + hamil_prod = torch.stack([q1a * q2a - q1b * q2b - q1c * q2c - q1d * q2d, + q1a * q2b + q1b * q2a + q1c * q2d - q1d * q2c, + q1a * q2c - q1b * q2d + q1c * q2a + q1d * q2b, + q1a * q2d + q1b * q2c - q1c * q2b + q1d * q2a], dim=-1) + return hamil_prod + + def interpolate(self, q1, q2, alpha): # [...,4],[...,4],[...,1] + # https://en.wikipedia.org/wiki/Slerp + cos_angle = (q1 * q2).sum(dim=-1, keepdim=True) # [...,1] + flip = cos_angle < 0 + q1 = q1 * (~flip) - q1 * flip # [...,4] + theta = cos_angle.abs().acos() # [...,1] + slerp = (((1 - alpha) * theta).sin() * q1 + (alpha * theta).sin() * q2) / theta.sin() # [...,4] + return slerp + + +pose = Pose() +lie = Lie() +quaternion = Quaternion() + + +def to_hom(X): + # Get homogeneous coordinates of the input. + X_hom = torch.cat([X, torch.ones_like(X[..., :1])], dim=-1) + return X_hom + + +# Basic operations of transforming 3D points between world/camera/image coordinates. +def world2cam(X, pose): # [B,N,3] + X_hom = to_hom(X) + return X_hom @ pose.transpose(-1, -2) + + +def cam2img(X, cam_intr): + return X @ cam_intr.transpose(-1, -2) + + +def img2cam(X, cam_intr): + return X @ cam_intr.inverse().transpose(-1, -2) + + +def cam2world(X, pose): + X_hom = to_hom(X) + pose_inv = Pose().invert(pose) + return X_hom @ pose_inv.transpose(-1, -2) + + +def angle_to_rotation_matrix(a, axis): + # Get the rotation matrix from Euler angle around specific axis. + roll = dict(X=1, Y=2, Z=0)[axis] + if isinstance(a, float): + a = torch.tensor(a) + zero = torch.zeros_like(a) + eye = torch.ones_like(a) + M = torch.stack([torch.stack([a.cos(), -a.sin(), zero], dim=-1), + torch.stack([a.sin(), a.cos(), zero], dim=-1), + torch.stack([zero, zero, eye], dim=-1)], dim=-2) + M = M.roll((roll, roll), dims=(-2, -1)) + return M + + +def get_center_and_ray(pose, intr, image_size): + """ + Args: + pose (tensor [3,4]/[B,3,4]): Camera pose. + intr (tensor [3,3]/[B,3,3]): Camera intrinsics. + image_size (list of int): Image size. + Returns: + center_3D (tensor [HW,3]/[B,HW,3]): Center of the camera. + ray (tensor [HW,3]/[B,HW,3]): Ray of the camera with depth=1 (note: not unit ray). + """ + H, W = image_size + # Given the intrinsic/extrinsic matrices, get the camera center and ray directions. + with torch.no_grad(): + # Compute image coordinate grid. + y_range = torch.arange(H, dtype=torch.float32, device=pose.device).add_(0.5) + x_range = torch.arange(W, dtype=torch.float32, device=pose.device).add_(0.5) + Y, X = torch.meshgrid(y_range, x_range, indexing="ij") # [H,W] + xy_grid = torch.stack([X, Y], dim=-1).view(-1, 2) # [HW,2] + # Compute center and ray. + if len(pose.shape) == 3: + batch_size = len(pose) + xy_grid = xy_grid.repeat(batch_size, 1, 1) # [B,HW,2] + grid_3D = img2cam(to_hom(xy_grid), intr) # [HW,3]/[B,HW,3] + center_3D = torch.zeros_like(grid_3D) # [HW,3]/[B,HW,3] + # Transform from camera to world coordinates. + grid_3D = cam2world(grid_3D, pose) # [HW,3]/[B,HW,3] + center_3D = cam2world(center_3D, pose) # [HW,3]/[B,HW,3] + ray = grid_3D - center_3D # [B,HW,3] + return center_3D, ray + + +def get_3D_points_from_dist(center, ray_unit, dist, multi=True): + # Two possible use cases: (1) center + ray_unit * dist, or (2) center + ray * depth + if multi: + center, ray_unit = center[..., None, :], ray_unit[..., None, :] # [...,1,3] + # x = c+dv + points_3D = center + ray_unit * dist # [...,3]/[...,N,3] + return points_3D + + +def convert_NDC(center, ray, intr, near=1): + # Shift camera center (ray origins) to near plane (z=1). + # (Unlike conventional NDC, we assume the cameras are facing towards the +z direction.) + center = center + (near - center[..., 2:]) / ray[..., 2:] * ray + # Projection. + cx, cy, cz = center.unbind(dim=-1) # [...,R] + rx, ry, rz = ray.unbind(dim=-1) # [...,R] + scale_x = intr[..., 0, 0] / intr[..., 0, 2] # [...] + scale_y = intr[..., 1, 1] / intr[..., 1, 2] # [...] + cnx = scale_x[..., None] * (cx / cz) + cny = scale_y[..., None] * (cy / cz) + cnz = 1 - 2 * near / cz + rnx = scale_x[..., None] * (rx / rz - cx / cz) + rny = scale_y[..., None] * (ry / rz - cy / cz) + rnz = 2 * near / cz + center_ndc = torch.stack([cnx, cny, cnz], dim=-1) # [...,R,3] + ray_ndc = torch.stack([rnx, rny, rnz], dim=-1) # [...,R,3] + return center_ndc, ray_ndc + + +def convert_NDC2(center, ray, intr): + # Similar to convert_NDC() but shift the ray origins to its own image plane instead of the global near plane. + # Also this version is much more interpretable. + scale_x = intr[..., 0, 0] / intr[..., 0, 2] # [...] + scale_y = intr[..., 1, 1] / intr[..., 1, 2] # [...] + # Get the metric image plane (i.e. new "center"): (sx*cx/cz, sy*cy/cz, 1-2/cz). + center = center + ray # This is the key difference. + cx, cy, cz = center.unbind(dim=-1) # [...,R] + image_plane = torch.stack([scale_x[..., None] * cx / cz, + scale_x[..., None] * cy / cz, + 1 - 2 / cz], dim=-1) + # Get the infinity plane: (sx*rx/rz, sy*ry/rz, 1). + rx, ry, rz = ray.unbind(dim=-1) # [...,R] + inf_plane = torch.stack([scale_x[..., None] * rx / rz, + scale_y[..., None] * ry / rz, + torch.ones_like(rz)], dim=-1) + # The NDC ray is the difference between the two planes, assuming t \in [0,1]. + ndc_ray = inf_plane - image_plane + return image_plane, ndc_ray + + +def rotation_distance(R1, R2, eps=1e-7): + # http://www.boris-belousov.net/2016/12/01/quat-dist/ + R_diff = R1 @ R2.transpose(-2, -1) + trace = R_diff[..., 0, 0] + R_diff[..., 1, 1] + R_diff[..., 2, 2] + angle = ((trace - 1) / 2).clamp(-1 + eps, 1 - eps).acos_() # numerical stability near -1/+1 + return angle + + +def get_oscil_novel_view_poses(N=60, angle=0.05, dist=5): + # Create circular viewpoints (small oscillations). + theta = torch.arange(N) / N * 2 * np.pi + R_x = angle_to_rotation_matrix((theta.sin() * angle).asin(), "X") + R_y = angle_to_rotation_matrix((theta.cos() * angle).asin(), "Y") + pose_rot = pose(R=R_y @ R_x) + pose_shift = pose(t=[0, 0, dist]) + pose_oscil = pose.compose([pose.invert(pose_shift), pose_rot, pose_shift]) + return pose_oscil + + +def cross_product_matrix(x): + """Matrix form of cross product opertaion. + + param x: [3,] tensor. + return: [3, 3] tensor representing the matrix form of cross product. + """ + return torch.tensor( + [[0, -x[2], x[1]], + [x[2], 0, -x[0]], + [-x[1], x[0], 0, ]] + ) + + +def essential_matrix(poses): + """Compute Essential Matrix from a relative pose. + + param poses: [views, 3, 4] tensor representing relative poses. + return: [views, 3, 3] tensor representing Essential Matrix. + """ + r = poses[..., 0:3] + t = poses[..., 3] + tx = torch.stack([cross_product_matrix(tt) for tt in t], axis=0) + return tx @ r + + +def fundamental_matrix(poses, intr1, intr2): + """Compute Fundamental Matrix from a relative pose and intrinsics. + + param poses: [views, 3, 4] tensor representing relative poses. + intr1: [3, 3] tensor. Camera intrinsic of reference image. + intr2: [views, 3, 3] tensor. Camera Intrinsic of target image. + return: [views, 3, 3] tensor representing Fundamental Matrix. + """ + return intr2.inverse().transpose(-1, -2) @ essential_matrix(poses) @ intr1.inverse() + + +def get_ray_depth_plane_intersection(center, ray, depths): + """Compute the intersection of a ray with a depth plane. + Args: + center (tensor [B,HW,3]): Camera center of the target pose. + ray (tensor [B,HW,3]): Ray direction of the target pose. + depth (tensor [L]): The depth values from the source view (e.g. for MPI planes). + Returns: + intsc_points (tensor [B,HW,L,3]): Intersecting 3D points with the MPI. + """ + # Each 3D point x along the ray v from center c can be written as x = c+t*v. + # Plane equation: n@x = d, where normal n = (0,0,1), d = depth. + # --> t = (d-n@c)/(n@v). + # --> x = c+t*v = c+(d-n@c)/(n@v)*v. + center, ray = center[:, :, None], ray[:, :, None] # [B,HW,L,3], [B,HW,1,3] + depths = depths[None, None, :, None] # [1,1,L,1] + intsc_points = center + (depths - center[..., 2:]) / ray[..., 2:] * ray # [B,HW,L,3] + return intsc_points + + +def unit_view_vector_to_rotation_matrix(v, axes="ZYZ"): + """ + Args: + v (tensor [...,3]): Unit vectors on the view sphere. + axes: rotation axis order. + + Returns: + rotation_matrix (tensor [...,3,3]): rotation matrix R @ v + [0, 0, 1] = 0. + """ + alpha = torch.arctan2(v[..., 1], v[..., 0]) # [...] + beta = np.pi - v[..., 2].arccos() # [...] + euler_angles = torch.stack([torch.ones_like(alpha) * np.pi / 2, -beta, alpha], dim=-1) # [...,3] + rot2 = angle_to_rotation_matrix(euler_angles[..., 2], axes[2]) # [...,3,3] + rot1 = angle_to_rotation_matrix(euler_angles[..., 1], axes[1]) # [...,3,3] + rot0 = angle_to_rotation_matrix(euler_angles[..., 0], axes[0]) # [...,3,3] + rot = rot2 @ rot1 @ rot0 # [...,3,3] + return rot.transpose(-2, -1) + + +def sample_on_spherical_cap(anchor, N, max_angle): + """Sample n points on the view hemisphere within the angle to x. + Args: + anchor (tensor [...,3]): Reference 3-D unit vector on the view hemisphere. + N (int): Number of sampled points. + max_angle (float): Sampled points should have max angle to x. + Returns: + sampled_points (tensor [...,N,3]): Sampled points on the spherical caps. + """ + batch_shape = anchor.shape[:-1] + # First, sample uniformly on a unit 2D disk. + radius = torch.rand(*batch_shape, N, device=anchor.device) # [...,N] + theta = torch.rand(*batch_shape, N, device=anchor.device) * 2 * np.pi # [...,N] + x = radius.sqrt() * theta.cos() # [...,N] + y = radius.sqrt() * theta.sin() # [...,N] + # Reparametrize to a unit spherical cap with height h. + # http://marc-b-reynolds.github.io/distribution/2016/11/28/Uniform.html + h = 1 - np.cos(max_angle) # spherical cap height + k = h * radius # [...,N] + s = (h * (2 - k)).sqrt() # [...,N] + points = torch.stack([s * x, s * y, 1 - k], dim=-1) # [...,N,3] + # Transform to center around the anchor. + ref_z = torch.tensor([0., 0., 1.], device=anchor.device) + v = -anchor.cross(ref_z) # [...,3] + ss_v = lie.skew_symmetric(v) # [...,3,3] + R = torch.eye(3, device=anchor.device) + ss_v + ss_v @ ss_v / (1 + anchor @ ref_z)[..., None, None] # [...,3,3] + points = points @ R.transpose(-2, -1) # [...,N,3] + return points + + +def sample_on_spherical_cap_northern(anchor, N, max_angle, away_from=None, max_reject_count=None): + """Sample n points only the northern view hemisphere within the angle to x.""" + + def find_invalid_points(points): + southern = points[..., 2] < 0 # [...,N] + if away_from is not None: + cosine_ab = (away_from * anchor).sum(dim=-1, keepdim=True) # [...,1] + cosine_ac = (away_from[..., None, :] * points).sum(dim=-1) # [...,N] + not_outwards = cosine_ab < cosine_ac # [...,N] + invalid = southern | not_outwards + else: + invalid = southern + return invalid + + assert (anchor[..., 2] > 0).all() + assert anchor.norm(dim=-1).allclose(torch.ones_like(anchor[..., 0])) + points = sample_on_spherical_cap(anchor, N, max_angle) # [...,N,3] + invalid = find_invalid_points(points) + count = 0 + while invalid.any(): + # Reject and resample. + points_resample = sample_on_spherical_cap(anchor, N, max_angle) + points[invalid] = points_resample[invalid] + invalid = find_invalid_points(points) + count += 1 + if max_reject_count and count > max_reject_count: + points = anchor.repeat(N, 1) + return points diff --git a/neuralangelo-main/projects/nerf/utils/misc.py b/neuralangelo-main/projects/nerf/utils/misc.py new file mode 100644 index 0000000000000000000000000000000000000000..22ea7866214661d97eaa88eaad4afc6acfd79714 --- /dev/null +++ b/neuralangelo-main/projects/nerf/utils/misc.py @@ -0,0 +1,71 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import torch + +from imaginaire.utils.distributed import dist_all_gather_tensor + + +def collate_test_data_batches(data_batches): + """Aggregate the list of test data from all devices and process the results. + Args: + data_batches (list): List of (hierarchical) dictionaries, where leaf entries are tensors. + Returns: + data_gather (dict): (hierarchical) dictionaries, where leaf entries are concatenated tensors. + """ + data_gather = dict() + for key in data_batches[0].keys(): + data_list = [data[key] for data in data_batches] + if isinstance(data_batches[0][key], dict): + data_gather[key] = collate_test_data_batches(data_list) + elif isinstance(data_batches[0][key], torch.Tensor): + data_gather[key] = torch.cat(data_list, dim=0) + data_gather[key] = torch.cat(dist_all_gather_tensor(data_gather[key].contiguous()), dim=0) + else: + raise TypeError + return data_gather + + +def get_unique_test_data(data_gather, idx): + """Aggregate the list of test data from all devices and process the results. + Args: + data_gather (dict): (hierarchical) dictionaries, where leaf entries are tensors. + idx (tensor): sample indices. + Returns: + data_all (dict): (hierarchical) dictionaries, where leaf entries are tensors ordered by idx. + """ + data_all = dict() + for key, value in data_gather.items(): + if isinstance(value, dict): + data_all[key] = get_unique_test_data(value, idx) + elif isinstance(value, torch.Tensor): + data_all[key] = [] + for i in range(max(idx) + 1): + # If multiple occurrences of the same idx, just choose the first one. If no occurrence, just ignore. + matches = (idx == i).nonzero(as_tuple=True)[0] + if matches.numel() != 0: + data_all[key].append(value[matches[0]]) + data_all[key] = torch.stack(data_all[key], dim=0) + else: + raise TypeError + return data_all + + +def trim_test_samples(data, max_samples=None): + for key, value in data.items(): + if isinstance(value, dict): + data[key] = trim_test_samples(value, max_samples=max_samples) + elif isinstance(value, torch.Tensor): + if max_samples is not None: + data[key] = value[:max_samples] + else: + raise TypeError diff --git a/neuralangelo-main/projects/nerf/utils/nerf_util.py b/neuralangelo-main/projects/nerf/utils/nerf_util.py new file mode 100644 index 0000000000000000000000000000000000000000..6a58466a5c3ce7bad9b0f129a9bb241cefda107e --- /dev/null +++ b/neuralangelo-main/projects/nerf/utils/nerf_util.py @@ -0,0 +1,270 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import numpy as np +import torch +import torch.nn.functional as torch_F + +from projects.nerf.utils import camera + + +def sample_dists(ray_size, dist_range, intvs, stratified, device="cuda"): + """Sample points on ray shooting from pixels using distance. + Args: + ray_size (int [2]): Integers for [batch size, number of rays]. + range (float [2]): Range of distance (depth) [min, max] to be sampled on rays. + intvs: (int): Number of points sampled on a ray. + stratified: (bool): Use stratified sampling or constant 0.5 sampling. + Returns: + dists (tensor [batch_size, num_ray, intvs, 1]): Sampled distance for all rays in a batch. + """ + batch_size, num_rays = ray_size + dist_min, dist_max = dist_range + if stratified: + rands = torch.rand(batch_size, num_rays, intvs, 1, device=device) + else: + rands = torch.empty(batch_size, num_rays, intvs, 1, device=device).fill_(0.5) + rands += torch.arange(intvs, dtype=torch.float, device=device)[None, None, :, None] # [B,R,N,1] + dists = rands / intvs * (dist_max - dist_min) + dist_min # [B,R,N,1] + return dists + + +def sample_dists_from_pdf(bin, weights, intvs_fine): + """Sample points on ray shooting from pixels using the weights from the coarse NeRF. + Args: + bin (tensor [batch_size, num_rays, intvs]): bins of distance values from the coarse NeRF. + weights (tensor [batch_size, num_rays, intvs]): weights from the coarse NeRF. + intvs_fine: (int): Number of fine-grained points sampled on a ray. + Returns: + dists (tensor [batch_size, num_ray, intvs, 1]): Sampled distance for all rays in a batch. + """ + pdf = torch_F.normalize(weights, p=1, dim=-1) + # Get CDF from PDF (along last dimension). + cdf = pdf.cumsum(dim=-1) # [B,R,N] + cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], dim=-1) # [B,R,N+1] + # Take uniform samples. + grid = torch.linspace(0, 1, intvs_fine + 1, device=pdf.device) # [Nf+1] + unif = 0.5 * (grid[:-1] + grid[1:]).repeat(*cdf.shape[:-1], 1) # [B,R,Nf] + idx = torch.searchsorted(cdf, unif, right=True) # [B,R,Nf] \in {1...N} + # Inverse transform sampling from CDF. + low = (idx - 1).clamp(min=0) # [B,R,Nf] + high = idx.clamp(max=cdf.shape[-1] - 1) # [B,R,Nf] + dist_min = bin[..., 0].gather(dim=2, index=low) # [B,R,Nf] + dist_max = bin[..., 0].gather(dim=2, index=high) # [B,R,Nf] + cdf_low = cdf.gather(dim=2, index=low) # [B,R,Nf] + cdf_high = cdf.gather(dim=2, index=high) # [B,R,Nf] + # Linear interpolation. + t = (unif - cdf_low) / (cdf_high - cdf_low + 1e-8) # [B,R,Nf] + dists = dist_min + t * (dist_max - dist_min) # [B,R,Nf] + return dists[..., None] # [B,R,Nf,1] + + +def reparametrize_dist(dist, param_type="metric"): + """Reparametrize the sampled distance values according to param_type. + Args: + dist (tensor): Sampled distance values. + param_type (str): Reparametrization type. + Returns: + dist_new (tensor): Reparametrized distance values. + """ + return dict( + metric=dist, + ndc=dist, + inverse=1 / (dist + 1e-8), + )[param_type] + + +def ray_generator(pose, intr, image_size, num_rays, full_image=False, camera_ndc=False, + ray_indices=None): + """Yield sampled rays for coordinate-based model to predict NeRF. + Args: + pose (tensor [bs,3,4]): Camera poses ([R,t]). + intr (tensor [bs,3,3]): Camera intrinsics. + image_size: (tensor [bs,2]): Image size [height, width]. + num_rays (int): Number of rays to sample (random rays unless full_image=True). + full_image (bool): Sample rays from the full image. + camera_ndc (bool): Use normalized device coordinate for camera. + Returns: + center_slice (tensor [bs, ray, 3]): Sampled 3-D center in the world coordinate. + ray_slice (tensor [bs, ray, 3]): Sampled 3-D ray in the world coordinate. + ray_idx (tensor [bs, ray]): Sampled indices to index sampled pixels on images. + """ + # Create a grid of centers and rays on an image. + batch_size = pose.shape[0] + # We used to randomly sample ray indices here. Now, we assume they are pre-generated and passed in. + if ray_indices is None: + num_pixels = image_size[0] * image_size[1] + if full_image: + # Sample rays from the full image. + ray_indices = torch.arange(0, num_pixels, device=pose.device).repeat(batch_size, 1) # [B,HW] + else: + # Sample rays randomly. The below is equivalent to batched torch.randperm(). + ray_indices = torch.rand(batch_size, num_pixels, device=pose.device).argsort(dim=1)[:, :num_rays] # [B,R] + center, ray = camera.get_center_and_ray(pose, intr, image_size) # [B,HW,3] + # Convert center/ray representations to NDC if necessary. + if camera_ndc == "new": + center, ray = camera.convert_NDC2(center, ray, intr=intr) + elif camera_ndc: + center, ray = camera.convert_NDC(center, ray, intr=intr) + # Yield num_rays of sampled rays in each iteration (when random, the loop will only iterate once). + for c in range(0, ray_indices.shape[1], num_rays): + ray_idx = ray_indices[:, c:c + num_rays] # [B,R] + batch_idx = torch.arange(batch_size, device=pose.device).repeat(ray_idx.shape[1], 1).t() # [B,R] + center_slice = center[batch_idx, ray_idx] # [B,R,3] + ray_slice = ray[batch_idx, ray_idx] # [B,R,3] + yield center_slice, ray_slice, ray_idx + + +def slice_by_ray_idx(var, ray_idx): + batch_size, num_rays = ray_idx.shape[:2] + batch_idx = torch.arange(batch_size, device=ray_idx.device).repeat(num_rays, 1).t() # [B,R] + var_slice = var[batch_idx, ray_idx] # [B,R,...] + return var_slice + + +def positional_encoding(input, num_freq_bases): + """Encode input into position codes. + Args: + input (tensor [bs, ..., N]): A batch of data with N dimension. + num_freq_bases: (int): The number of frequency base of the code. + Returns: + input_enc (tensor [bs, ..., 2*N*num_freq_bases]): Positional codes for input. + """ + freq = 2 ** torch.arange(num_freq_bases, dtype=torch.float32, device=input.device) * np.pi # [L]. + spectrum = input[..., None] * freq # [B,...,N,L]. + sin, cos = spectrum.sin(), spectrum.cos() # [B,...,N,L]. + input_enc = torch.stack([sin, cos], dim=-2) # [B,...,N,2,L]. + input_enc = input_enc.view(*input.shape[:-1], -1) # [B,...,2NL]. + return input_enc + + +def get_inverse_depth(depth, opacity=None, camera_ndc=False, eps=1e-10): + # Compute inverse depth for visualization. + if opacity is not None: + return (1 - depth) / opacity if camera_ndc else 1 / (depth / opacity + eps) + else: + return (1 - depth) if camera_ndc else 1 / (depth + eps) + + +class MLPwithSkipConnection(torch.nn.Module): + + def __init__(self, layer_dims, skip_connection=[], activ=None, use_layernorm=False, use_weightnorm=False): + """Initialize a multi-layer perceptron with skip connection. + Args: + layer_dims: A list of integers representing the number of channels in each layer. + skip_connection: A list of integers representing the index of layers to add skip connection. + """ + super().__init__() + self.skip_connection = skip_connection + self.use_layernorm = use_layernorm + self.linears = torch.nn.ModuleList() + if use_layernorm: + self.layer_norm = torch.nn.ModuleList() + layer_dim_pairs = list(zip(layer_dims[:-1], layer_dims[1:])) + for li, (k_in, k_out) in enumerate(layer_dim_pairs): + if li in self.skip_connection: + k_in += layer_dims[0] + linear = torch.nn.Linear(k_in, k_out) + if use_weightnorm: + linear = torch.nn.utils.weight_norm(linear) + self.linears.append(linear) + if use_layernorm and li != len(layer_dim_pairs) - 1: + self.layer_norm.append(torch.nn.LayerNorm(k_out)) + if li == len(layer_dim_pairs) - 1: + self.linears[-1].bias.data.fill_(0.0) + self.activ = activ or torch_F.relu_ + + def forward(self, input): + feat = input + for li, linear in enumerate(self.linears): + if li in self.skip_connection: + feat = torch.cat([feat, input], dim=-1) + feat = linear(feat) + if li != len(self.linears) - 1: + if self.use_layernorm: + feat = self.layer_norm[li](feat) + feat = self.activ(feat) + return feat + + +def intersect_with_sphere(center, ray_unit, radius=1.0): + ctc = (center * center).sum(dim=-1, keepdim=True) # [...,1] + ctv = (center * ray_unit).sum(dim=-1, keepdim=True) # [...,1] + b2_minus_4ac = ctv ** 2 - (ctc - radius ** 2) + dist_near = -ctv - b2_minus_4ac.sqrt() + dist_far = -ctv + b2_minus_4ac.sqrt() + return dist_near, dist_far + + +def get_pixel_radii(intr): + # fx and fy should be very close. + focal = (intr[..., 0, 0] + intr[..., 1, 1]) / 2 + radii = 1. / focal / np.sqrt(3) + return radii + + +def contract(x, r_in=1, r_out=2, eps=1e-8): + """ Contract function in mip-NeRF 360 (eq 10). + Args: + x (tensor [...,3]): The input points. + Returns: + x_warp (tensor [...,3]): The warped points. + """ + x_norm = x.norm(dim=-1, keepdim=True) # [...,1] + scale = r_out - r_in * (r_out - r_in) / (x_norm + eps) # [...,1] + x_contract = scale * torch_F.normalize(x, dim=-1) # [...,3] + # No effect if within r_in. + inside = x_norm <= r_in + x_warp = torch.where(inside, x, x_contract) # [...,3] + return x_warp + + +def contract_jacobian(x, r_in=1, r_out=2, eps=1e-8): + """ Jacobian of the contract function in mip-NeRF 360. + Args: + x (tensor [...,3]): The input points. + Returns: + jacobian (tensor [...,3,3]): The Jacobian at the input points. + """ + x_norm = x.norm(dim=-1)[..., None, None] # [...,1,1] + x_norm_sq = (x ** 2).sum(dim=-1)[..., None, None] # [...,1,1] should be numerically more stable + b = r_in * (r_out - r_in) + scale = r_out - b / (x_norm + eps) # [...,1,1] + x_outer_prod = x[..., None] * x[..., None, :] # [...,3,3] + eye = torch.eye(3, device=x.device).repeat(*x_norm.shape) # [...,3,3] + term1 = b * x_outer_prod / (x_norm_sq ** 2 + eps) # [...,3,3] + term2 = scale * (eye - x_outer_prod / x_norm_sq + eps) / (x_norm + eps) # [...,3,3] + jacobian_contract = term1 + term2 # [...,3,3] + # No effect if within r_in. + inside = x_norm <= r_in + jacobian = torch.where(inside, eye, jacobian_contract) # [...,3] + return jacobian + + +def contract_mip(mean, cov, r_in=1, r_out=2, diag=False): + """ Contraction function on mip-NeRF 360 Gaussians. + Args: + mean (tensor [...,3]): The mean values. + cov (tensor [...,3,3]): The covariance values. + r_in (float): The radius of unaffected region. + r_out (float): The radius of contracted region. r_in < r_out. + Returns: + mean_warp (tensor [...,3]): The contracted mean values. + cov_warp (tensor [...,3,3]): The contracted covariance values. + """ + mean_warp = contract(mean, r_in=r_in, r_out=r_out) # [...,3] + jacobian = contract_jacobian(mean, r_in=r_in, r_out=r_out) # [...,3,3] + if diag: + cov_warp = (jacobian * cov[..., None, :]) @ jacobian.transpose(-2, -1) + else: + cov_warp = jacobian @ cov @ jacobian.transpose(-2, -1) + return mean_warp, cov_warp diff --git a/neuralangelo-main/projects/nerf/utils/render.py b/neuralangelo-main/projects/nerf/utils/render.py new file mode 100644 index 0000000000000000000000000000000000000000..58ce996ececeba7ee4c97347b0e6eb8a736e5409 --- /dev/null +++ b/neuralangelo-main/projects/nerf/utils/render.py @@ -0,0 +1,112 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import torch +from torch.cuda.amp import autocast + + +def volume_rendering_weights(ray, densities, depths, depth_far=None): + """The volume rendering function. Details can be found in the NeRF paper. + Args: + ray (tensor [batch,ray,3]): The ray directions in world space. + densities (tensor [batch,ray,samples]): The predicted volume density samples. + depths (tensor [batch,ray,samples,1]): The corresponding depth samples. + depth_far (tensor [batch,ray,1,1]): The farthest depth for computing the last interval. + Returns: + weights (tensor [batch,ray,samples,1]): The predicted weight of each sampled point along the ray (in [0,1]). + """ + ray_length = ray.norm(dim=-1, keepdim=True) # [B,R,1] + if depth_far is None: + depth_far = torch.empty_like(depths[..., :1, :]).fill_(1e10) # [B,R,1,1] + depths_aug = torch.cat([depths, depth_far], dim=2) # [B,R,N+1,1] + dists = depths_aug * ray_length[..., None] # [B,R,N+1,1] + # Volume rendering: compute rendering weights (using quadrature). + dist_intvs = dists[..., 1:, 0] - dists[..., :-1, 0] # [B,R,N] + sigma_delta = densities * dist_intvs # [B,R,N] + sigma_delta_0 = torch.cat([torch.zeros_like(sigma_delta[..., :1]), + sigma_delta[..., :-1]], dim=2) # [B,R,N] + T = (-sigma_delta_0.cumsum(dim=2)).exp_() # [B,R,N] + alphas = 1 - (-sigma_delta).exp_() # [B,R,N] + # Compute weights for compositing samples. + weights = (T * alphas)[..., None] # [B,R,N,1] + return weights + + +def volume_rendering_weights_dist(densities, dists, dist_far=None): + """The volume rendering function. Details can be found in the NeRF paper. + Args: + densities (tensor [batch,ray,samples]): The predicted volume density samples. + dists (tensor [batch,ray,samples,1]): The corresponding distance samples. + dist_far (tensor [batch,ray,1,1]): The farthest distance for computing the last interval. + Returns: + weights (tensor [batch,ray,samples,1]): The predicted weight of each sampled point along the ray (in [0,1]). + """ + # TODO: re-consolidate!! + if dist_far is None: + dist_far = torch.empty_like(dists[..., :1, :]).fill_(1e10) # [B,R,1,1] + dists = torch.cat([dists, dist_far], dim=2) # [B,R,N+1,1] + # Volume rendering: compute rendering weights (using quadrature). + dist_intvs = dists[..., 1:, 0] - dists[..., :-1, 0] # [B,R,N] + sigma_delta = densities * dist_intvs # [B,R,N] + sigma_delta_0 = torch.cat([torch.zeros_like(sigma_delta[..., :1]), sigma_delta[..., :-1]], dim=2) # [B,R,N] + T = (-sigma_delta_0.cumsum(dim=2)).exp_() # [B,R,N] + alphas = 1 - (-sigma_delta).exp_() # [B,R,N] + # Compute weights for compositing samples. + weights = (T * alphas)[..., None] # [B,R,N,1] + return weights + + +def volume_rendering_alphas_dist(densities, dists, dist_far=None): + """The volume rendering function. Details can be found in the NeRF paper. + Args: + densities (tensor [batch,ray,samples]): The predicted volume density samples. + dists (tensor [batch,ray,samples,1]): The corresponding distance samples. + dist_far (tensor [batch,ray,1,1]): The farthest distance for computing the last interval. + Returns: + alphas (tensor [batch,ray,samples,1]): The occupancy of each sampled point along the ray (in [0,1]). + """ + if dist_far is None: + dist_far = torch.empty_like(dists[..., :1, :]).fill_(1e10) # [B,R,1,1] + dists = torch.cat([dists, dist_far], dim=2) # [B,R,N+1,1] + # Volume rendering: compute rendering weights (using quadrature). + dist_intvs = dists[..., 1:, 0] - dists[..., :-1, 0] # [B,R,N] + sigma_delta = densities * dist_intvs # [B,R,N] + alphas = 1 - (-sigma_delta).exp_() # [B,R,N] + return alphas + + +def alpha_compositing_weights(alphas): + """Alpha compositing of (sampled) MPIs given their RGBs and alphas. + Args: + alphas (tensor [batch,ray,samples]): The predicted opacity values. + Returns: + weights (tensor [batch,ray,samples,1]): The predicted weight of each MPI (in [0,1]). + """ + alphas_front = torch.cat([torch.zeros_like(alphas[..., :1]), + alphas[..., :-1]], dim=2) # [B,R,N] + with autocast(enabled=False): # TODO: may be unstable in some cases. + visibility = (1 - alphas_front).cumprod(dim=2) # [B,R,N] + weights = (alphas * visibility)[..., None] # [B,R,N,1] + return weights + + +def composite(quantities, weights): + """Composite the samples to render the RGB/depth/opacity of the corresponding pixels. + Args: + quantities (tensor [batch,ray,samples,k]): The quantity to be weighted summed. + weights (tensor [batch,ray,samples,1]): The predicted weight of each sampled point along the ray. + Returns: + quantity (tensor [batch,ray,k]): The expected (rendered) quantity. + """ + # Integrate RGB and depth weighted by probability. + quantity = (quantities * weights).sum(dim=2) # [B,R,K] + return quantity diff --git a/neuralangelo-main/projects/nerf/utils/visualize.py b/neuralangelo-main/projects/nerf/utils/visualize.py new file mode 100644 index 0000000000000000000000000000000000000000..f870afa51b4d39fdc36e549119c058ba493d7c7d --- /dev/null +++ b/neuralangelo-main/projects/nerf/utils/visualize.py @@ -0,0 +1,120 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import numpy as np +import torch +import matplotlib.pyplot as plt +import k3d + +from projects.nerf.utils import camera + + +def get_camera_mesh(pose, depth=1): + vertices = torch.tensor([[-0.5, -0.5, 1], + [0.5, -0.5, 1], + [0.5, 0.5, 1], + [-0.5, 0.5, 1], + [0, 0, 0]]) * depth # [6,3] + faces = torch.tensor([[0, 1, 2], + [0, 2, 3], + [0, 1, 4], + [1, 2, 4], + [2, 3, 4], + [3, 0, 4]]) # [6,3] + vertices = camera.cam2world(vertices[None], pose) # [N,6,3] + wireframe = vertices[:, [0, 1, 2, 3, 0, 4, 1, 2, 4, 3]] # [N,10,3] + return vertices, faces, wireframe + + +def merge_meshes(vertices, faces): + mesh_N, vertex_N = vertices.shape[:2] + faces_merged = torch.cat([faces + i * vertex_N for i in range(mesh_N)], dim=0) + vertices_merged = vertices.view(-1, vertices.shape[-1]) + return vertices_merged, faces_merged + + +def merge_wireframes(wireframe): + wf_first, wf_last, wf_dummy = wireframe[:, :1], wireframe[:, -1:], wireframe[:, :1] * np.nan + wireframe_merged = torch.cat([wf_first, wireframe, wf_last, wf_dummy], dim=1) + return wireframe_merged + + +def get_xyz_indicators(pose, length=0.1): + xyz = torch.eye(4, 3)[None] * length + xyz = camera.cam2world(xyz, pose) + return xyz + + +def merge_xyz_indicators(xyz): # [N,4,3] + xyz = xyz[:, [[-1, 0], [-1, 1], [-1, 2]]] # [N,3,2,3] + xyz_0, xyz_1 = xyz.unbind(dim=2) # [N,3,3] + xyz_dummy = xyz_0 * np.nan + xyz_merged = torch.stack([xyz_0, xyz_0, xyz_1, xyz_1, xyz_dummy], dim=2) # [N,3,5,3] + return xyz_merged + + +def k3d_visualize_pose(poses, vis_depth=0.5, xyz_length=0.1, center_size=0.1, xyz_width=0.02): + # poses has shape [N,3,4] potentially in sequential order + N = len(poses) + centers_cam = torch.zeros(N, 1, 3) + centers_world = camera.cam2world(centers_cam, poses) + centers_world = centers_world[:, 0] + # Get the camera wireframes. + vertices, faces, wireframe = get_camera_mesh(poses, depth=vis_depth) + xyz = get_xyz_indicators(poses, length=xyz_length) + vertices_merged, faces_merged = merge_meshes(vertices, faces) + wireframe_merged = merge_wireframes(wireframe) + xyz_merged = merge_xyz_indicators(xyz) + # Set the color map for the camera trajectory and the xyz indicators. + color_map = plt.get_cmap("gist_rainbow") + center_color = [] + vertices_merged_color = [] + wireframe_color = [] + xyz_color = [] + x_hex, y_hex, z_hex = int(255) << 16, int(255) << 8, int(255) + for i in range(N): + # Set the camera pose colors (with a smooth gradient color map). + r, g, b, _ = color_map(i / (N - 1)) + r, g, b = r * 0.8, g * 0.8, b * 0.8 + pose_rgb_hex = (int(r * 255) << 16) + (int(g * 255) << 8) + int(b * 255) + center_color += [pose_rgb_hex] + vertices_merged_color += [pose_rgb_hex] * 5 + wireframe_color += [pose_rgb_hex] * 13 + # Set the xyz indicator colors. + xyz_color += [x_hex] * 5 + [y_hex] * 5 + [z_hex] * 5 + # Plot in K3D. + plot = k3d.plot(name="poses", + height=800, + camera_rotate_speed=5.0, + camera_zoom_speed=3.0, + camera_pan_speed=1.0, + ) + plot += k3d.points(centers_world, + colors=center_color, + point_size=center_size, + shader="3d", + ) + plot += k3d.mesh(vertices_merged, faces_merged, + colors=vertices_merged_color, + side="double", + opacity=0.05, + ) + plot += k3d.line(wireframe_merged, + colors=wireframe_color, + shader="simple", + ) + plot += k3d.line(xyz_merged, + colors=xyz_color, + shader="thick", + width=xyz_width, + ) + return plot diff --git a/neuralangelo-main/projects/neuralangelo/configs/base.yaml b/neuralangelo-main/projects/neuralangelo/configs/base.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c559ba750388b2ca640a6bf68a894c4b0c095a39 --- /dev/null +++ b/neuralangelo-main/projects/neuralangelo/configs/base.yaml @@ -0,0 +1,147 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. +# ----------------------------------------------------------------------------- + +logging_iter: 9999999999999 # disable the printing logger + +max_iter: 500000 + +wandb_scalar_iter: 100 +wandb_image_iter: 10000 +validation_iter: 5000 +speed_benchmark: False + +checkpoint: + save_iter: 20000 + +trainer: + type: projects.neuralangelo.trainer + ema_config: + enabled: False + load_ema_checkpoint: False + loss_weight: + render: 1.0 + eikonal: 0.1 + curvature: 5e-4 + init: + type: none + amp_config: + enabled: False + depth_vis_scale: 0.5 + +model: + type: projects.neuralangelo.model + object: + sdf: + mlp: + num_layers: 2 + hidden_dim: 256 + skip: [] + activ: softplus + activ_params: + beta: 100 + geometric_init: True + weight_norm: True + out_bias: 0.5 + inside_out: False + encoding: + type: hashgrid + levels: 16 + hashgrid: + min_logres: 5 + max_logres: 11 + dict_size: 22 + dim: 8 + range: [-2,2] + coarse2fine: + enabled: True + init_active_level: 4 + step: 5000 + gradient: + mode: numerical + taps: 4 + rgb: + mlp: + num_layers: 4 + hidden_dim: 256 + skip: [] + activ: relu_ + activ_params: {} + weight_norm: True + mode: idr + encoding_view: + type: spherical + levels: 3 + s_var: + init_val: 3. + anneal_end: 0.1 + background: + enabled: True + white: False + mlp: + num_layers: 8 + hidden_dim: 256 + skip: [4] + num_layers_rgb: 2 + hidden_dim_rgb: 128 + skip_rgb: [] + activ: relu + activ_params: {} + activ_density: softplus + activ_density_params: {} + view_dep: True + encoding: + type: fourier + levels: 10 + encoding_view: + type: spherical + levels: 3 + render: + rand_rays: 512 + num_samples: + coarse: 64 + fine: 16 + background: 32 + num_sample_hierarchy: 4 + stratified: True + appear_embed: + enabled: False + dim: 8 + +optim: + type: AdamW + params: + lr: 1e-3 + weight_decay: 1e-3 + sched: + iteration_mode: True + type: two_steps_with_warmup + warm_up_end: 5000 + two_steps: [300000,400000] + gamma: 10.0 + +data: + type: projects.nerf.datasets.nerf_blender + root: datasets/nerf-synthetic/lego + use_multi_epoch_loader: True + num_workers: 4 + preload: True + num_images: # The number of training images. + train: + image_size: [800,800] + batch_size: 2 + subset: + val: + image_size: [400,400] + batch_size: 2 + subset: 4 + max_viz_samples: 16 + readjust: + center: [0.,0.,0.] + scale: 1. diff --git a/neuralangelo-main/projects/neuralangelo/configs/custom/template.yaml b/neuralangelo-main/projects/neuralangelo/configs/custom/template.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3750d56c383d9eb4f4dfb37382d3783be1ac7dc9 --- /dev/null +++ b/neuralangelo-main/projects/neuralangelo/configs/custom/template.yaml @@ -0,0 +1,40 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. +# ----------------------------------------------------------------------------- + +_parent_: projects/neuralangelo/configs/base.yaml + +model: + object: + sdf: + mlp: + inside_out: False + encoding: + coarse2fine: + init_active_level: 8 + appear_embed: + enabled: True + dim: 8 + +data: + type: projects.neuralangelo.data + root: # The root path of the dataset. + num_images: # The number of training images. + train: + image_size: [1200,1600] + batch_size: 1 + subset: + val: + image_size: [300,400] + batch_size: 1 + subset: 1 + max_viz_samples: + readjust: + center: [0.,0.,0.] + scale: 1. diff --git a/neuralangelo-main/projects/neuralangelo/configs/dtu.yaml b/neuralangelo-main/projects/neuralangelo/configs/dtu.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1ed831849969a75e9b4865f48ef81a9aa713f054 --- /dev/null +++ b/neuralangelo-main/projects/neuralangelo/configs/dtu.yaml @@ -0,0 +1,37 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. +# ----------------------------------------------------------------------------- + +_parent_: projects/neuralangelo/configs/base.yaml + +model: + object: + sdf: + mlp: + inside_out: False + encoding: + coarse2fine: + init_active_level: 4 + s_var: + init_val: 1.4 + appear_embed: + enabled: False + +data: + type: projects.neuralangelo.data + root: datasets/dtu/dtu_scan24 + train: + image_size: [1200,1600] + batch_size: 1 + subset: + val: + image_size: [300,400] + batch_size: 1 + subset: 1 + max_viz_samples: 16 diff --git a/neuralangelo-main/projects/neuralangelo/configs/tnt.yaml b/neuralangelo-main/projects/neuralangelo/configs/tnt.yaml new file mode 100644 index 0000000000000000000000000000000000000000..aa175301e22fb8291ffd68a338bfe21e9e63a330 --- /dev/null +++ b/neuralangelo-main/projects/neuralangelo/configs/tnt.yaml @@ -0,0 +1,37 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. +# ----------------------------------------------------------------------------- + +_parent_: projects/neuralangelo/configs/base.yaml + +model: + object: + sdf: + mlp: + inside_out: False # True for Meetingroom. + encoding: + coarse2fine: + init_active_level: 8 + appear_embed: + enabled: True + dim: 8 + +data: + type: projects.neuralangelo.data + root: datasets/tanks_and_temples/Barn + num_images: 410 # The number of training images. + train: + image_size: [835,1500] + batch_size: 1 + subset: + val: + image_size: [300,540] + batch_size: 1 + subset: 1 + max_viz_samples: 16 diff --git a/neuralangelo-main/projects/neuralangelo/data.py b/neuralangelo-main/projects/neuralangelo/data.py new file mode 100644 index 0000000000000000000000000000000000000000..b09bc863f5e1083956545612c77c45e60cf4b6d7 --- /dev/null +++ b/neuralangelo-main/projects/neuralangelo/data.py @@ -0,0 +1,131 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import json +import numpy as np +import torch +import torchvision.transforms.functional as torchvision_F +from PIL import Image, ImageFile + +from projects.nerf.datasets import base +from projects.nerf.utils import camera +from projects.neuralangelo.utils.misc import gl_to_cv + +ImageFile.LOAD_TRUNCATED_IMAGES = True + + +class Dataset(base.Dataset): + + def __init__(self, cfg, is_inference=False): + super().__init__(cfg, is_inference=is_inference, is_test=False) + cfg_data = cfg.data + self.root = cfg_data.root + self.preload = cfg_data.preload + self.H, self.W = cfg_data.val.image_size if is_inference else cfg_data.train.image_size + meta_fname = f"{cfg_data.root}/transforms.json" + with open(meta_fname) as file: + self.meta = json.load(file) + self.list = self.meta["frames"] + if cfg_data[self.split].subset: + subset = cfg_data[self.split].subset + subset_idx = np.linspace(0, len(self.list), subset+1)[:-1].astype(int) + self.list = [self.list[i] for i in subset_idx] + self.num_rays = cfg.model.render.rand_rays + self.readjust = getattr(cfg_data, "readjust", None) + # Preload dataset if possible. + if cfg_data.preload: + self.images = self.preload_threading(self.get_image, cfg_data.num_workers) + self.cameras = self.preload_threading(self.get_camera, cfg_data.num_workers, data_str="cameras") + + def __getitem__(self, idx): + """Process raw data and return processed data in a dictionary. + + Args: + idx: The index of the sample of the dataset. + Returns: A dictionary containing the data. + idx (scalar): The index of the sample of the dataset. + image (R tensor): Image idx for per-image embedding. + image (Rx3 tensor): Image with pixel values in [0,1] for supervision. + intr (3x3 tensor): The camera intrinsics of `image`. + pose (3x4 tensor): The camera extrinsics [R,t] of `image`. + """ + # Keep track of sample index for convenience. + sample = dict(idx=idx) + # Get the images. + image, image_size_raw = self.images[idx] if self.preload else self.get_image(idx) + image = self.preprocess_image(image) + # Get the cameras (intrinsics and pose). + intr, pose = self.cameras[idx] if self.preload else self.get_camera(idx) + intr, pose = self.preprocess_camera(intr, pose, image_size_raw) + # Pre-sample ray indices. + if self.split == "train": + ray_idx = torch.randperm(self.H * self.W)[:self.num_rays] # [R] + image_sampled = image.flatten(1, 2)[:, ray_idx].t() # [R,3] + sample.update( + ray_idx=ray_idx, + image_sampled=image_sampled, + intr=intr, + pose=pose, + ) + else: # keep image during inference + sample.update( + image=image, + intr=intr, + pose=pose, + ) + return sample + + def get_image(self, idx): + fpath = self.list[idx]["file_path"] + image_fname = f"{self.root}/{fpath}" + image = Image.open(image_fname) + image.load() + image_size_raw = image.size + return image, image_size_raw + + def preprocess_image(self, image): + # Resize the image. + image = image.resize((self.W, self.H)) + image = torchvision_F.to_tensor(image) + rgb = image[:3] + return rgb + + def get_camera(self, idx): + # Camera intrinsics. + intr = torch.tensor([[self.meta["fl_x"], self.meta["sk_x"], self.meta["cx"]], + [self.meta["sk_y"], self.meta["fl_y"], self.meta["cy"]], + [0, 0, 1]]).float() + # Camera pose. + c2w_gl = torch.tensor(self.list[idx]["transform_matrix"], dtype=torch.float32) + c2w = gl_to_cv(c2w_gl) + if not self.meta['centered']: + # center scene + center = np.array(self.meta["sphere_center"]) + if self.readjust: + center += np.array(getattr(self.readjust, "center", [0])) + c2w[:3, -1] -= center + if not self.meta['scaled']: + # scale scene + scale = np.array(self.meta["sphere_radius"]) + if self.readjust: + scale *= getattr(self.readjust, "scale", 1.) + c2w[:3, -1] /= scale + w2c = camera.Pose().invert(c2w[:3]) + return intr, w2c + + def preprocess_camera(self, intr, pose, image_size_raw): + # Adjust the intrinsics according to the resized image. + intr = intr.clone() + raw_W, raw_H = image_size_raw + intr[0] *= self.W / raw_W + intr[1] *= self.H / raw_H + return intr, pose diff --git a/neuralangelo-main/projects/neuralangelo/model.py b/neuralangelo-main/projects/neuralangelo/model.py new file mode 100644 index 0000000000000000000000000000000000000000..685e60b76c94c09665d843f2b50e5a878053cd04 --- /dev/null +++ b/neuralangelo-main/projects/neuralangelo/model.py @@ -0,0 +1,298 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +from functools import partial +import torch +import torch.nn.functional as torch_F +from collections import defaultdict + +from imaginaire.models.base import Model as BaseModel +from projects.nerf.utils import nerf_util, camera, render +from projects.neuralangelo.utils import misc +from projects.neuralangelo.utils.modules import NeuralSDF, NeuralRGB, BackgroundNeRF + + +class Model(BaseModel): + + def __init__(self, cfg_model, cfg_data): + super().__init__(cfg_model, cfg_data) + self.cfg_render = cfg_model.render + self.white_background = cfg_model.background.white + self.with_background = cfg_model.background.enabled + self.with_appear_embed = cfg_model.appear_embed.enabled + self.anneal_end = cfg_model.object.s_var.anneal_end + self.outside_val = 1000. * (-1 if cfg_model.object.sdf.mlp.inside_out else 1) + self.image_size_train = cfg_data.train.image_size + self.image_size_val = cfg_data.val.image_size + # Define models. + self.build_model(cfg_model, cfg_data) + # Define functions. + self.ray_generator = partial(nerf_util.ray_generator, + camera_ndc=False, + num_rays=cfg_model.render.rand_rays) + self.sample_dists_from_pdf = partial(nerf_util.sample_dists_from_pdf, + intvs_fine=cfg_model.render.num_samples.fine) + self.to_full_val_image = partial(misc.to_full_image, image_size=cfg_data.val.image_size) + + def build_model(self, cfg_model, cfg_data): + # appearance encoding + if cfg_model.appear_embed.enabled: + assert cfg_data.num_images is not None + self.appear_embed = torch.nn.Embedding(cfg_data.num_images, cfg_model.appear_embed.dim) + if cfg_model.background.enabled: + self.appear_embed_outside = torch.nn.Embedding(cfg_data.num_images, cfg_model.appear_embed.dim) + else: + self.appear_embed_outside = None + else: + self.appear_embed = self.appear_embed_outside = None + self.neural_sdf = NeuralSDF(cfg_model.object.sdf) + self.neural_rgb = NeuralRGB(cfg_model.object.rgb, feat_dim=cfg_model.object.sdf.mlp.hidden_dim, + appear_embed=cfg_model.appear_embed) + if cfg_model.background.enabled: + self.background_nerf = BackgroundNeRF(cfg_model.background, appear_embed=cfg_model.appear_embed) + else: + self.background_nerf = None + self.s_var = torch.nn.Parameter(torch.tensor(cfg_model.object.s_var.init_val, dtype=torch.float32)) + + def forward(self, data): + # Randomly sample and render the pixels. + output = self.render_pixels(data["pose"], data["intr"], image_size=self.image_size_train, + stratified=self.cfg_render.stratified, sample_idx=data["idx"], + ray_idx=data["ray_idx"]) + return output + + @torch.no_grad() + def inference(self, data): + self.eval() + # Render the full images. + output = self.render_image(data["pose"], data["intr"], image_size=self.image_size_val, + stratified=False, sample_idx=data["idx"]) # [B,N,C] + # Get full rendered RGB and depth images. + rot = data["pose"][..., :3, :3] # [B,3,3] + normal_cam = -output["gradient"] @ rot.transpose(-1, -2) # [B,HW,3] + output.update( + rgb_map=self.to_full_val_image(output["rgb"]), # [B,3,H,W] + opacity_map=self.to_full_val_image(output["opacity"]), # [B,1,H,W] + depth_map=self.to_full_val_image(output["depth"]), # [B,1,H,W] + normal_map=self.to_full_val_image(normal_cam), # [B,3,H,W] + ) + return output + + def render_image(self, pose, intr, image_size, stratified=False, sample_idx=None): + """ Render the rays given the camera intrinsics and poses. + Args: + pose (tensor [batch,3,4]): Camera poses ([R,t]). + intr (tensor [batch,3,3]): Camera intrinsics. + stratified (bool): Whether to stratify the depth sampling. + sample_idx (tensor [batch]): Data sample index. + Returns: + output: A dictionary containing the outputs. + """ + output = defaultdict(list) + for center, ray, _ in self.ray_generator(pose, intr, image_size, full_image=True): + ray_unit = torch_F.normalize(ray, dim=-1) # [B,R,3] + output_batch = self.render_rays(center, ray_unit, sample_idx=sample_idx, stratified=stratified) + if not self.training: + dist = render.composite(output_batch["dists"], output_batch["weights"]) # [B,R,1] + depth = dist / ray.norm(dim=-1, keepdim=True) + output_batch.update(depth=depth) + for key, value in output_batch.items(): + if value is not None: + output[key].append(value.detach()) + # Concat each item (list) in output into one tensor. Concatenate along the ray dimension (1) + for key, value in output.items(): + output[key] = torch.cat(value, dim=1) + return output + + def render_pixels(self, pose, intr, image_size, stratified=False, sample_idx=None, ray_idx=None): + center, ray = camera.get_center_and_ray(pose, intr, image_size) # [B,HW,3] + center = nerf_util.slice_by_ray_idx(center, ray_idx) # [B,R,3] + ray = nerf_util.slice_by_ray_idx(ray, ray_idx) # [B,R,3] + ray_unit = torch_F.normalize(ray, dim=-1) # [B,R,3] + output = self.render_rays(center, ray_unit, sample_idx=sample_idx, stratified=stratified) + return output + + def render_rays(self, center, ray_unit, sample_idx=None, stratified=False): + with torch.no_grad(): + near, far, outside = self.get_dist_bounds(center, ray_unit) + app, app_outside = self.get_appearance_embedding(sample_idx, ray_unit.shape[1]) + output_object = self.render_rays_object(center, ray_unit, near, far, outside, app, stratified=stratified) + if self.with_background: + output_background = self.render_rays_background(center, ray_unit, far, app_outside, stratified=stratified) + # Concatenate object and background samples. + rgbs = torch.cat([output_object["rgbs"], output_background["rgbs"]], dim=2) # [B,R,No+Nb,3] + dists = torch.cat([output_object["dists"], output_background["dists"]], dim=2) # [B,R,No+Nb,1] + alphas = torch.cat([output_object["alphas"], output_background["alphas"]], dim=2) # [B,R,No+Nb] + else: + rgbs = output_object["rgbs"] # [B,R,No,3] + dists = output_object["dists"] # [B,R,No,1] + alphas = output_object["alphas"] # [B,R,No] + weights = render.alpha_compositing_weights(alphas) # [B,R,No+Nb,1] + # Compute weights and composite samples. + rgb = render.composite(rgbs, weights) # [B,R,3] + if self.white_background: + opacity_all = render.composite(1., weights) # [B,R,1] + rgb = rgb + (1 - opacity_all) + # Collect output. + output = dict( + rgb=rgb, # [B,R,3] + opacity=output_object["opacity"], # [B,R,1]/None + outside=outside, # [B,R,1] + dists=dists, # [B,R,No+Nb,1] + weights=weights, # [B,R,No+Nb,1] + gradient=output_object["gradient"], # [B,R,3]/None + gradients=output_object["gradients"], # [B,R,No,3] + hessians=output_object["hessians"], # [B,R,No,3]/None + ) + return output + + def render_rays_object(self, center, ray_unit, near, far, outside, app, stratified=False): + with torch.no_grad(): + dists = self.sample_dists_all(center, ray_unit, near, far, stratified=stratified) # [B,R,N,3] + points = camera.get_3D_points_from_dist(center, ray_unit, dists) # [B,R,N,3] + sdfs, feats = self.neural_sdf.forward(points) # [B,R,N,1],[B,R,N,K] + sdfs[outside[..., None].expand_as(sdfs)] = self.outside_val + # Compute 1st- and 2nd-order gradients. + rays_unit = ray_unit[..., None, :].expand_as(points).contiguous() # [B,R,N,3] + gradients, hessians = self.neural_sdf.compute_gradients(points, compute_hessian=self.training, sdf=sdfs) + normals = torch_F.normalize(gradients, dim=-1) # [B,R,N,3] + rgbs = self.neural_rgb.forward(points, normals, rays_unit, feats, app=app) # [B,R,N,3] + # SDF volume rendering. + alphas = self.compute_neus_alphas(ray_unit, sdfs, gradients, dists, dist_far=far[..., None], + progress=self.progress) # [B,R,N] + if not self.training: + weights = render.alpha_compositing_weights(alphas) # [B,R,N,1] + opacity = render.composite(1., weights) # [B,R,1] + gradient = render.composite(gradients, weights) # [B,R,3] + else: + opacity = None + gradient = None + # Collect output. + output = dict( + rgbs=rgbs, # [B,R,N,3] + sdfs=sdfs[..., 0], # [B,R,N] + dists=dists, # [B,R,N,1] + alphas=alphas, # [B,R,N] + opacity=opacity, # [B,R,3]/None + gradient=gradient, # [B,R,3]/None + gradients=gradients, # [B,R,N,3] + hessians=hessians, # [B,R,N,3]/None + ) + return output + + def render_rays_background(self, center, ray_unit, far, app_outside, stratified=False): + with torch.no_grad(): + dists = self.sample_dists_background(ray_unit, far, stratified=stratified) + points = camera.get_3D_points_from_dist(center, ray_unit, dists) # [B,R,N,3] + rays_unit = ray_unit[..., None, :].expand_as(points) # [B,R,N,3] + rgbs, densities = self.background_nerf.forward(points, rays_unit, app_outside) # [B,R,N,3] + alphas = render.volume_rendering_alphas_dist(densities, dists) # [B,R,N] + # Collect output. + output = dict( + rgbs=rgbs, # [B,R,3] + dists=dists, # [B,R,N,1] + alphas=alphas, # [B,R,N] + ) + return output + + @torch.no_grad() + def get_dist_bounds(self, center, ray_unit): + dist_near, dist_far = nerf_util.intersect_with_sphere(center, ray_unit, radius=1.) + dist_near.relu_() # Distance (and thus depth) should be non-negative. + outside = dist_near.isnan() + dist_near[outside], dist_far[outside] = 1, 1.2 # Dummy distances. Density will be set to 0. + return dist_near, dist_far, outside + + def get_appearance_embedding(self, sample_idx, num_rays): + if self.with_appear_embed: + # Object appearance embedding. + num_samples_all = self.cfg_render.num_samples.coarse + \ + self.cfg_render.num_samples.fine * self.cfg_render.num_sample_hierarchy + app = self.appear_embed(sample_idx)[:, None, None] # [B,1,1,C] + app = app.expand(-1, num_rays, num_samples_all, -1) # [B,R,N,C] + # Background appearance embedding. + if self.with_background: + app_outside = self.appear_embed_outside(sample_idx)[:, None, None] # [B,1,1,C] + app_outside = app_outside.expand(-1, num_rays, self.cfg_render.num_samples.background, -1) # [B,R,N,C] + else: + app_outside = None + else: + app = app_outside = None + return app, app_outside + + @torch.no_grad() + def sample_dists_all(self, center, ray_unit, near, far, stratified=False): + dists = nerf_util.sample_dists(ray_unit.shape[:2], dist_range=(near[..., None], far[..., None]), + intvs=self.cfg_render.num_samples.coarse, stratified=stratified) + if self.cfg_render.num_sample_hierarchy > 0: + points = camera.get_3D_points_from_dist(center, ray_unit, dists) # [B,R,N,3] + sdfs = self.neural_sdf.sdf(points) # [B,R,N] + for h in range(self.cfg_render.num_sample_hierarchy): + dists_fine = self.sample_dists_hierarchical(dists, sdfs, inv_s=(64 * 2 ** h)) # [B,R,Nf,1] + dists = torch.cat([dists, dists_fine], dim=2) # [B,R,N+Nf,1] + dists, sort_idx = dists.sort(dim=2) + if h != self.cfg_render.num_sample_hierarchy - 1: + points_fine = camera.get_3D_points_from_dist(center, ray_unit, dists_fine) # [B,R,Nf,3] + sdfs_fine = self.neural_sdf.sdf(points_fine) # [B,R,Nf] + sdfs = torch.cat([sdfs, sdfs_fine], dim=2) # [B,R,N+Nf] + sdfs = sdfs.gather(dim=2, index=sort_idx.expand_as(sdfs)) # [B,R,N+Nf,1] + return dists + + def sample_dists_hierarchical(self, dists, sdfs, inv_s, robust=True, eps=1e-5): + sdfs = sdfs[..., 0] # [B,R,N] + prev_sdfs, next_sdfs = sdfs[..., :-1], sdfs[..., 1:] # [B,R,N-1] + prev_dists, next_dists = dists[..., :-1, 0], dists[..., 1:, 0] # [B,R,N-1] + mid_sdfs = (prev_sdfs + next_sdfs) * 0.5 # [B,R,N-1] + cos_val = (next_sdfs - prev_sdfs) / (next_dists - prev_dists + 1e-5) # [B,R,N-1] + if robust: + prev_cos_val = torch.cat([torch.zeros_like(cos_val)[..., :1], cos_val[..., :-1]], dim=-1) # [B,R,N-1] + cos_val = torch.stack([prev_cos_val, cos_val], dim=-1).min(dim=-1).values # [B,R,N-1] + dist_intvs = dists[..., 1:, 0] - dists[..., :-1, 0] # [B,R,N-1] + est_prev_sdf = mid_sdfs - cos_val * dist_intvs * 0.5 # [B,R,N-1] + est_next_sdf = mid_sdfs + cos_val * dist_intvs * 0.5 # [B,R,N-1] + prev_cdf = (est_prev_sdf * inv_s).sigmoid() # [B,R,N-1] + next_cdf = (est_next_sdf * inv_s).sigmoid() # [B,R,N-1] + alphas = ((prev_cdf - next_cdf) / (prev_cdf + eps)).clip_(0.0, 1.0) # [B,R,N-1] + weights = render.alpha_compositing_weights(alphas) # [B,R,N-1,1] + dists_fine = self.sample_dists_from_pdf(dists, weights=weights[..., 0]) # [B,R,Nf,1] + return dists_fine + + def sample_dists_background(self, ray_unit, far, stratified=False, eps=1e-5): + inv_dists = nerf_util.sample_dists(ray_unit.shape[:2], dist_range=(1, 0), + intvs=self.cfg_render.num_samples.background, stratified=stratified) + dists = far[..., None] / (inv_dists + eps) # [B,R,N,1] + return dists + + def compute_neus_alphas(self, ray_unit, sdfs, gradients, dists, dist_far=None, progress=1., eps=1e-5): + sdfs = sdfs[..., 0] # [B,R,N] + # SDF volume rendering in NeuS. + inv_s = self.s_var.exp() + true_cos = (ray_unit[..., None, :] * gradients).sum(dim=-1, keepdim=False) # [B,R,N] + iter_cos = self._get_iter_cos(true_cos, progress=progress) # [B,R,N] + # Estimate signed distances at section points + if dist_far is None: + dist_far = torch.empty_like(dists[..., :1, :]).fill_(1e10) # [B,R,1,1] + dists = torch.cat([dists, dist_far], dim=2) # [B,R,N+1,1] + dist_intvs = dists[..., 1:, 0] - dists[..., :-1, 0] # [B,R,N] + est_prev_sdf = sdfs - iter_cos * dist_intvs * 0.5 # [B,R,N] + est_next_sdf = sdfs + iter_cos * dist_intvs * 0.5 # [B,R,N] + prev_cdf = (est_prev_sdf * inv_s).sigmoid() # [B,R,N] + next_cdf = (est_next_sdf * inv_s).sigmoid() # [B,R,N] + alphas = ((prev_cdf - next_cdf) / (prev_cdf + eps)).clip_(0.0, 1.0) # [B,R,N] + # weights = render.alpha_compositing_weights(alphas) # [B,R,N,1] + return alphas + + def _get_iter_cos(self, true_cos, progress=1.): + anneal_ratio = min(progress / self.anneal_end, 1.) + # The anneal strategy below keeps the cos value alive at the beginning of training iterations. + return -((-true_cos * 0.5 + 0.5).relu() * (1.0 - anneal_ratio) + + (-true_cos).relu() * anneal_ratio) # always non-positive diff --git a/neuralangelo-main/projects/neuralangelo/scripts/convert_data_to_json.py b/neuralangelo-main/projects/neuralangelo/scripts/convert_data_to_json.py new file mode 100644 index 0000000000000000000000000000000000000000..21b2922792d528a9ab0af1583fd1dd3c14f8a07c --- /dev/null +++ b/neuralangelo-main/projects/neuralangelo/scripts/convert_data_to_json.py @@ -0,0 +1,199 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import numpy as np +from argparse import ArgumentParser +import os +import sys +from pathlib import Path +import json +import math + +dir_path = Path(os.path.dirname(os.path.realpath(__file__))).parents[2] +sys.path.append(dir_path.__str__()) + +from third_party.colmap.scripts.python.read_write_model import read_model, qvec2rotmat # NOQA + + +def find_closest_point(p1, d1, p2, d2): + # Calculate the direction vectors of the lines + d1_norm = d1 / np.linalg.norm(d1) + d2_norm = d2 / np.linalg.norm(d2) + + # Create the coefficient matrix A and the constant vector b + A = np.vstack((d1_norm, -d2_norm)).T + b = p2 - p1 + + # Solve the linear system to find the parameters t1 and t2 + t1, t2 = np.linalg.lstsq(A, b, rcond=None)[0] + + # Calculate the closest point on each line + closest_point1 = p1 + d1_norm * t1 + closest_point2 = p2 + d2_norm * t2 + + # Calculate the average of the two closest points + closest_point = 0.5 * (closest_point1 + closest_point2) + + return closest_point + + +def bound_by_pose(images): + poses = [] + for img in images.values(): + rotation = qvec2rotmat(img.qvec) + translation = img.tvec.reshape(3, 1) + w2c = np.concatenate([rotation, translation], 1) + w2c = np.concatenate([w2c, np.array([0, 0, 0, 1])[None]], 0) + c2w = np.linalg.inv(w2c) + poses.append(c2w) + + center = np.array([0.0, 0.0, 0.0]) + for f in poses: + src_frame = f[0:3, :] + for g in poses: + tgt_frame = g[0:3, :] + p = find_closest_point(src_frame[:, 3], src_frame[:, 2], tgt_frame[:, 3], tgt_frame[:, 2]) + center += p + center /= len(poses)**2 + + radius = 0.0 + for f in poses: + radius += np.linalg.norm(f[0:3, 3]) + radius /= len(poses) + bounding_box = [ + [center[0] - radius, center[0] + radius], + [center[1] - radius, center[1] + radius], + [center[2] - radius, center[2] + radius], + ] + return center, radius, bounding_box + + +def bound_by_points(points3D): + xyzs = np.stack([point.xyz for point in points3D.values()]) + center = xyzs.mean(axis=0) + std = xyzs.std(axis=0) + radius = float(std.max() * 3) # use 3*std to define the region, equivalent to 99% percentile + bounding_box = [ + [center[0] - std[0] * 3, center[0] + std[0] * 3], + [center[1] - std[1] * 3, center[1] + std[1] * 3], + [center[2] - std[2] * 3, center[2] + std[2] * 3], + ] + return center, radius, bounding_box + + +def export_to_json(cameras, images, bounding_box, center, radius, file_path): + intrinsic_param = np.array([camera.params for camera in cameras.values()]) + fl_x = intrinsic_param[0][0] # TODO: only supports single camera for now + fl_y = intrinsic_param[0][1] + cx = intrinsic_param[0][2] + cy = intrinsic_param[0][3] + image_width = np.array([camera.width for camera in cameras.values()]) + image_height = np.array([camera.height for camera in cameras.values()]) + w = image_width[0] + h = image_height[0] + + angle_x = math.atan(w / (fl_x * 2)) * 2 + angle_y = math.atan(h / (fl_y * 2)) * 2 + + out = { + "camera_angle_x": angle_x, + "camera_angle_y": angle_y, + "fl_x": fl_x, + "fl_y": fl_y, + "sk_x": 0.0, # TODO: check if colmap has skew + "sk_y": 0.0, + "k1": 0.0, # take undistorted images only + "k2": 0.0, + "k3": 0.0, + "k4": 0.0, + "p1": 0.0, + "p2": 0.0, + "is_fisheye": False, # TODO: not supporting fish eye camera + "cx": cx, + "cy": cy, + "w": int(w), + "h": int(h), + "aabb_scale": np.exp2(np.rint(np.log2(radius))), # power of two, for INGP resolution computation + "aabb_range": bounding_box, + "sphere_center": center, + "sphere_radius": radius, + "centered": True, # flag for if a scene is centered at origin + "scaled": False, # flag for if a scene is scaled + "frames": [] + } + + flip_mat = np.array([ + [1, 0, 0, 0], + [0, -1, 0, 0], + [0, 0, -1, 0], + [0, 0, 0, 1] + ]) + + # read poses + for img in sorted(images.values()): + rotation = qvec2rotmat(img.qvec) + translation = img.tvec.reshape(3, 1) + w2c = np.concatenate([rotation, translation], 1) + w2c = np.concatenate([w2c, np.array([0, 0, 0, 1])[None]], 0) + c2w = np.linalg.inv(w2c) + c2w[:3, -1] -= center # center scene at origin, we do NOT scale the scene for compatibility with iNGP + c2w = c2w @ flip_mat # convert to GL convention used in iNGP + + frame = {"file_path": "images/" + img.name, "transform_matrix": c2w.tolist()} + out["frames"].append(frame) + + with open(file_path, "w") as outputfile: + json.dump(out, outputfile, indent=2) + + return + + +def auto_bound(args): + cameras, images, points3D = read_model(os.path.join(args.data_dir, "sparse"), ext=".bin") + + # define bounding regions based on scene type + if args.scene_type == "outdoor": + center_points, radius_points, bounding_box_points = bound_by_points(points3D) + center_pose, radius_pose, bounding_box_pose = bound_by_pose(images) + # check the differences + diff = np.linalg.norm(center_points - center_pose) + squared_radii = radius_points + radius_pose + diff = np.inf + if diff / squared_radii < 0.05: + # scene center is close to pose center with 95% overlap, use pose + center, radius, bounding_box = center_pose, radius_pose, bounding_box_pose + else: + # different results, use sfm points + center, radius, bounding_box = center_points, radius_points, bounding_box_points + elif args.scene_type == "indoor": + # use sfm points as a proxy to define bounding regions + center, radius, bounding_box = bound_by_points(points3D) + elif args.scene_type == "object": + # use poses as a proxy to define bounding regions + center, radius, bounding_box = bound_by_pose(images) + else: + raise TypeError("Unknown scene type") + + # export json file + export_to_json(cameras, images, bounding_box, list(center), radius, os.path.join(args.data_dir, "transforms.json")) + print("Writing data to json file: ", os.path.join(args.data_dir, "transforms.json")) + return + + +if __name__ == "__main__": + parser = ArgumentParser() + parser.add_argument("--data_dir", type=str, default=None, help="Path to data") + parser.add_argument("--scene_type", type=str, default="outdoor", choices=["outdoor", "indoor", "object"], + help="Select scene type. Outdoor for building-scale reconstruction; " + "indoor for room-scale reconstruction; object for object-centric scene reconstruction.") + args = parser.parse_args() + auto_bound(args) diff --git a/neuralangelo-main/projects/neuralangelo/scripts/convert_dtu_to_json.py b/neuralangelo-main/projects/neuralangelo/scripts/convert_dtu_to_json.py new file mode 100644 index 0000000000000000000000000000000000000000..aa95899ef486c03cff27c6792b515825c37a8bdc --- /dev/null +++ b/neuralangelo-main/projects/neuralangelo/scripts/convert_dtu_to_json.py @@ -0,0 +1,136 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import numpy as np +import json +from argparse import ArgumentParser +import os +import cv2 +from PIL import Image, ImageFile +from glob import glob +import math +import sys +from pathlib import Path + +dir_path = Path(os.path.dirname(os.path.realpath(__file__))).parents[2] +sys.path.append(dir_path.__str__()) +from projects.neuralangelo.utils import misc # NOQA + +ImageFile.LOAD_TRUNCATED_IMAGES = True + + +def load_K_Rt_from_P(filename, P=None): + # This function is borrowed from IDR: https://github.com/lioryariv/idr + if P is None: + lines = open(filename).read().splitlines() + if len(lines) == 4: + lines = lines[1:] + lines = [[x[0], x[1], x[2], x[3]] for x in (x.split(" ") for x in lines)] + P = np.asarray(lines).astype(np.float32).squeeze() + + out = cv2.decomposeProjectionMatrix(P) + K = out[0] + R = out[1] + t = out[2] + + K = K / K[2, 2] + intrinsics = np.eye(4) + intrinsics[:3, :3] = K + + pose = np.eye(4, dtype=np.float32) + pose[:3, :3] = R.transpose() + pose[:3, 3] = (t[:3] / t[3])[:, 0] + + return intrinsics, pose + + +def dtu_to_json(args): + assert args.dtu_path, "Provide path to DTU dataset" + scene_list = os.listdir(args.dtu_path) + + for scene in scene_list: + scene_path = os.path.join(args.dtu_path, scene) + if not os.path.isdir(scene_path) or 'scan' not in scene: + continue + + out = { + "k1": 0.0, # take undistorted images only + "k2": 0.0, + "k3": 0.0, + "k4": 0.0, + "p1": 0.0, + "p2": 0.0, + "is_fisheye": False, + "frames": [] + } + + camera_param = dict(np.load(os.path.join(scene_path, 'cameras_sphere.npz'))) + images_lis = sorted(glob(os.path.join(scene_path, 'image/*.png'))) + for idx, image in enumerate(images_lis): + image = os.path.basename(image) + + world_mat = camera_param['world_mat_%d' % idx] + scale_mat = camera_param['scale_mat_%d' % idx] + + # scale and decompose + P = world_mat @ scale_mat + P = P[:3, :4] + intrinsic_param, c2w = load_K_Rt_from_P(None, P) + c2w_gl = misc.cv_to_gl(c2w) + + frame = {"file_path": 'image/' + image, "transform_matrix": c2w_gl.tolist()} + out["frames"].append(frame) + + fl_x = intrinsic_param[0][0] + fl_y = intrinsic_param[1][1] + cx = intrinsic_param[0][2] + cy = intrinsic_param[1][2] + sk_x = intrinsic_param[0][1] + sk_y = intrinsic_param[1][0] + w, h = Image.open(os.path.join(scene_path, 'image', image)).size + + angle_x = math.atan(w / (fl_x * 2)) * 2 + angle_y = math.atan(h / (fl_y * 2)) * 2 + + scale_mat = scale_mat.astype(float) + + out.update({ + "camera_angle_x": angle_x, + "camera_angle_y": angle_y, + "fl_x": fl_x, + "fl_y": fl_y, + "cx": cx, + "cy": cy, + "sk_x": sk_x, + "sk_y": sk_y, + "w": int(w), + "h": int(h), + "aabb_scale": np.exp2(np.rint(np.log2(scale_mat[0, 0]))), # power of two, for INGP resolution computation + "sphere_center": [scale_mat[0, -1], scale_mat[1, -1], scale_mat[2, -1]], + "sphere_radius": scale_mat[0, 0], + "centered": True, + "scaled": True, + }) + + file_path = os.path.join(scene_path, 'transforms.json') + with open(file_path, "w") as outputfile: + json.dump(out, outputfile, indent=2) + print('Writing data to json file: ', file_path) + + +if __name__ == '__main__': + parser = ArgumentParser() + parser.add_argument('--dtu_path', type=str, default=None) + + args = parser.parse_args() + + dtu_to_json(args) diff --git a/neuralangelo-main/projects/neuralangelo/scripts/convert_tnt_to_json.py b/neuralangelo-main/projects/neuralangelo/scripts/convert_tnt_to_json.py new file mode 100644 index 0000000000000000000000000000000000000000..58bbc491e23f1f847f21e9775f86a2b1c0b65b1d --- /dev/null +++ b/neuralangelo-main/projects/neuralangelo/scripts/convert_tnt_to_json.py @@ -0,0 +1,250 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import os +import numpy as np +import json +import sys +from pathlib import Path +from argparse import ArgumentParser +import shutil +import trimesh + +dir_path = Path(os.path.dirname(os.path.realpath(__file__))).parents[2] +sys.path.append(dir_path.__str__()) + +from projects.neuralangelo.scripts.convert_data_to_json import export_to_json # NOQA + +from third_party.colmap.scripts.python.database import COLMAPDatabase # NOQA +from third_party.colmap.scripts.python.read_write_model import read_model, rotmat2qvec # NOQA + + +def create_init_files(pinhole_dict_file, db_file, out_dir): + # Partially adapted from https://github.com/Kai-46/nerfplusplus/blob/master/colmap_runner/run_colmap_posed.py + + if not os.path.exists(out_dir): + os.mkdir(out_dir) + + # create template + with open(pinhole_dict_file) as fp: + pinhole_dict = json.load(fp) + + template = {} + cameras_line_template = '{camera_id} RADIAL {width} {height} {f} {cx} {cy} {k1} {k2}\n' + images_line_template = '{image_id} {qw} {qx} {qy} {qz} {tx} {ty} {tz} {camera_id} {image_name}\n\n' + + for img_name in pinhole_dict: + # w, h, fx, fy, cx, cy, qvec, t + params = pinhole_dict[img_name] + w = params[0] + h = params[1] + fx = params[2] + # fy = params[3] + cx = params[4] + cy = params[5] + qvec = params[6:10] + tvec = params[10:13] + + cam_line = cameras_line_template.format( + camera_id="{camera_id}", width=w, height=h, f=fx, cx=cx, cy=cy, k1=0, k2=0) + img_line = images_line_template.format(image_id="{image_id}", qw=qvec[0], qx=qvec[1], qy=qvec[2], qz=qvec[3], + tx=tvec[0], ty=tvec[1], tz=tvec[2], camera_id="{camera_id}", + image_name=img_name) + template[img_name] = (cam_line, img_line) + + # read database + db = COLMAPDatabase.connect(db_file) + table_images = db.execute("SELECT * FROM images") + img_name2id_dict = {} + for row in table_images: + img_name2id_dict[row[1]] = row[0] + + cameras_txt_lines = [template[img_name][0].format(camera_id=1)] + images_txt_lines = [] + for img_name, img_id in img_name2id_dict.items(): + image_line = template[img_name][1].format(image_id=img_id, camera_id=1) + images_txt_lines.append(image_line) + + with open(os.path.join(out_dir, 'cameras.txt'), 'w') as fp: + fp.writelines(cameras_txt_lines) + + with open(os.path.join(out_dir, 'images.txt'), 'w') as fp: + fp.writelines(images_txt_lines) + fp.write('\n') + + # create an empty points3D.txt + fp = open(os.path.join(out_dir, 'points3D.txt'), 'w') + fp.close() + + +def convert_cam_dict_to_pinhole_dict(cam_dict, pinhole_dict_file): + # Partially adapted from https://github.com/Kai-46/nerfplusplus/blob/master/colmap_runner/run_colmap_posed.py + + print('Writing pinhole_dict to: ', pinhole_dict_file) + h = 1080 + w = 1920 + + pinhole_dict = {} + for img_name in cam_dict: + W2C = cam_dict[img_name] + + # params + fx = 0.6 * w + fy = 0.6 * w + cx = w / 2.0 + cy = h / 2.0 + + qvec = rotmat2qvec(W2C[:3, :3]) + tvec = W2C[:3, 3] + + params = [w, h, fx, fy, cx, cy, + qvec[0], qvec[1], qvec[2], qvec[3], + tvec[0], tvec[1], tvec[2]] + pinhole_dict[img_name] = params + + with open(pinhole_dict_file, 'w') as fp: + json.dump(pinhole_dict, fp, indent=2, sort_keys=True) + + +def load_COLMAP_poses(cam_file, img_dir, tf='w2c'): + # load img_dir namges + names = sorted(os.listdir(img_dir)) + + with open(cam_file) as f: + lines = f.readlines() + + # C2W + poses = {} + for idx, line in enumerate(lines): + if idx % 5 == 0: # header + img_idx, valid, _ = line.split(' ') + if valid != '-1': + poses[int(img_idx)] = np.eye(4) + poses[int(img_idx)] + else: + if int(img_idx) in poses: + num = np.array([float(n) for n in line.split(' ')]) + poses[int(img_idx)][idx % 5-1, :] = num + + if tf == 'c2w': + return poses + else: + # convert to W2C (follow nerf convention) + poses_w2c = {} + for k, v in poses.items(): + poses_w2c[names[k]] = np.linalg.inv(v) + return poses_w2c + + +def load_transformation(trans_file): + with open(trans_file) as f: + lines = f.readlines() + + trans = np.eye(4) + for idx, line in enumerate(lines): + num = np.array([float(n) for n in line.split(' ')]) + trans[idx, :] = num + + return trans + + +def align_gt_with_cam(pts, trans): + trans_inv = np.linalg.inv(trans) + pts_aligned = pts @ trans_inv[:3, :3].transpose(-1, -2) + trans_inv[:3, -1] + return pts_aligned + + +def compute_bound(pts): + bounding_box = np.array([pts.max(axis=0), pts.min(axis=0)]) + center = bounding_box.mean(axis=0) + radius = np.max(np.linalg.norm(pts - center, axis=-1)) * 1.01 + return center, radius, bounding_box.T.tolist() + + +def init_colmap(args): + assert args.tnt_path, "Provide path to Tanks and Temples dataset" + scene_list = os.listdir(args.tnt_path) + + for scene in scene_list: + scene_path = os.path.join(args.tnt_path, scene) + # extract features + os.system(f"colmap feature_extractor --database_path {scene_path}/database.db \ + --image_path {scene_path}/images \ + --ImageReader.camera_model=RADIAL \ + --SiftExtraction.use_gpu=true \ + --SiftExtraction.num_threads=32 \ + --ImageReader.single_camera=true" + ) + + # match features + os.system(f"colmap sequential_matcher \ + --database_path {scene_path}/database.db \ + --SiftMatching.use_gpu=true" + ) + + # read poses + poses = load_COLMAP_poses(os.path.join(scene_path, f'{scene}_COLMAP_SfM.log'), + os.path.join(scene_path, 'images')) + + # convert to colmap files + pinhole_dict_file = os.path.join(scene_path, 'pinhole_dict.json') + convert_cam_dict_to_pinhole_dict(poses, pinhole_dict_file) + + db_file = os.path.join(scene_path, 'database.db') + sfm_dir = os.path.join(scene_path, 'sparse') + create_init_files(pinhole_dict_file, db_file, sfm_dir) + + # bundle adjustment + os.makedirs(f'{scene_path}/sparse/0', exist_ok=True) + os.system(f"colmap point_triangulator \ + --database_path {scene_path}/database.db \ + --image_path {scene_path}/images \ + --input_path {scene_path}/sparse \ + --output_path {scene_path}/sparse/0 \ + --Mapper.tri_ignore_two_view_tracks=true" + ) + os.system(f"colmap bundle_adjuster \ + --input_path {scene_path}/sparse/0 \ + --output_path {scene_path}/sparse/0 \ + --BundleAdjustment.refine_extrinsics=false" + ) + + # undistortion + os.system(f"colmap image_undistorter \ + --image_path {scene_path}/images \ + --input_path {scene_path}/sparse/0 \ + --output_path {scene_path}/dense \ + --output_type COLMAP \ + --max_image_size 1500" + ) + shutil.rmtree(f'{scene_path}/sparse') + + # read for bounding information + trans = load_transformation(os.path.join(scene_path, f'{scene}_trans.txt')) + pts = trimesh.load(os.path.join(scene_path, f'{scene}.ply')) + pts = pts.vertices + pts_aligned = align_gt_with_cam(pts, trans) + center, radius, bounding_box = compute_bound(pts_aligned[::100]) + + # colmap to json + cameras, images, points3D = read_model(os.path.join(args.tnt_path, scene, 'dense/sparse'), ext='.bin') + export_to_json(cameras, images, bounding_box, list(center), radius, os.path.join(scene_path, 'transforms.json')) + print('Writing data to json file: ', os.path.join(scene_path, 'transforms.json')) + + +if __name__ == '__main__': + parser = ArgumentParser() + parser.add_argument('--tnt_path', type=str, default=None, help='Path to tanks and temples dataset') + + args = parser.parse_args() + + init_colmap(args) diff --git a/neuralangelo-main/projects/neuralangelo/scripts/extract_mesh.py b/neuralangelo-main/projects/neuralangelo/scripts/extract_mesh.py new file mode 100644 index 0000000000000000000000000000000000000000..935c8bd40ed275f02b1b97c9b926ddaf31328684 --- /dev/null +++ b/neuralangelo-main/projects/neuralangelo/scripts/extract_mesh.py @@ -0,0 +1,95 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import argparse +import json +import os +import sys +import numpy as np + +sys.path.append(os.getcwd()) +from imaginaire.config import Config, recursive_update_strict, parse_cmdline_arguments # noqa: E402 +from imaginaire.utils.distributed import init_dist, get_world_size, is_master, master_only_print as print # noqa: E402 +from imaginaire.utils.gpu_affinity import set_affinity # noqa: E402 +from imaginaire.trainers.utils.logging import init_logging # noqa: E402 +from imaginaire.trainers.utils.get_trainer import get_trainer # noqa: E402 +from projects.neuralangelo.utils.mesh import extract_mesh # noqa: E402 + + +def parse_args(): + parser = argparse.ArgumentParser(description="Training") + parser.add_argument("--config", required=True, help="Path to the training config file.") + parser.add_argument("--logdir", help="Dir for saving logs and models.") + parser.add_argument("--checkpoint", default="", help="Checkpoint path.") + parser.add_argument('--local_rank', type=int, default=os.getenv('LOCAL_RANK', 0)) + parser.add_argument('--single_gpu', action='store_true') + parser.add_argument("--resolution", default=512, type=int, help="Marching cubes resolution") + parser.add_argument("--block_res", default=64, type=int, help="Block-wise resolution for marching cubes") + parser.add_argument("--output_file", default="mesh.ply", type=str, help="Output file name") + args, cfg_cmd = parser.parse_known_args() + return args, cfg_cmd + + +def main(): + args, cfg_cmd = parse_args() + set_affinity(args.local_rank) + cfg = Config(args.config) + + cfg_cmd = parse_cmdline_arguments(cfg_cmd) + recursive_update_strict(cfg, cfg_cmd) + + # If args.single_gpu is set to True, we will disable distributed data parallel. + if not args.single_gpu: + # this disables nccl timeout + os.environ["NCLL_BLOCKING_WAIT"] = "0" + os.environ["NCCL_ASYNC_ERROR_HANDLING"] = "0" + cfg.local_rank = args.local_rank + init_dist(cfg.local_rank, rank=-1, world_size=-1) + print(f"Running mesh extraction with {get_world_size()} GPUs.") + + cfg.logdir = init_logging(args.config, args.logdir, makedir=True) + + # Initialize data loaders and models. + trainer = get_trainer(cfg, is_inference=True, seed=0) + # Load checkpoint. + trainer.checkpointer.load(args.checkpoint, load_opt=False, load_sch=False) + trainer.model.eval() + + # Set the coarse-to-fine levels. + trainer.current_iteration = cfg.max_iter + if cfg.model.object.sdf.encoding.coarse2fine.enabled: + trainer.model_module.neural_sdf.set_active_levels(trainer.current_iteration) + if cfg.model.object.sdf.gradient.mode == "numerical": + trainer.model_module.neural_sdf.set_normal_epsilon() + + meta_fname = f"{cfg.data.root}/transforms.json" + with open(meta_fname) as file: + meta = json.load(file) + + if "aabb_range" in meta: + bounds = (np.array(meta["aabb_range"]) - np.array(meta["sphere_center"])[..., None]) / meta["sphere_radius"] + else: + bounds = np.array([[-1.0, 1.0], [-1.0, 1.0], [-1.0, 1.0]]) + + mesh = extract_mesh(sdf_func=lambda x: -trainer.model_module.neural_sdf.sdf(x), + bounds=bounds, intv=(2.0 / args.resolution), block_res=args.block_res) + + if is_master(): + print(f"vertices: {len(mesh.vertices)}") + print(f"faces: {len(mesh.faces)}") + # center and scale + mesh.vertices = mesh.vertices * meta["sphere_radius"] + np.array(meta["sphere_center"]) + mesh.export(args.output_file) + + +if __name__ == "__main__": + main() diff --git a/neuralangelo-main/projects/neuralangelo/scripts/generate_config.py b/neuralangelo-main/projects/neuralangelo/scripts/generate_config.py new file mode 100644 index 0000000000000000000000000000000000000000..024ec530a5c08034f1f60274870b21695b7d2385 --- /dev/null +++ b/neuralangelo-main/projects/neuralangelo/scripts/generate_config.py @@ -0,0 +1,81 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import os +import sys +from argparse import ArgumentParser +from pathlib import Path +import yaml +from addict import Dict +from PIL import Image, ImageFile + +dir_path = Path(os.path.dirname(os.path.realpath(__file__))).parents[2] +sys.path.append(dir_path.__str__()) + +ImageFile.LOAD_TRUNCATED_IMAGES = True + + +def generate_config(args): + cfg = Dict() + cfg._parent_ = "projects/neuralangelo/configs/base.yaml" + num_images = len(os.listdir(os.path.join(args.data_dir, "images"))) + # model cfg + if args.auto_exposure_wb: + cfg.data.num_images = num_images + cfg.model.appear_embed.enabled = True + cfg.model.appear_embed.dim = 8 + else: + cfg.model.appear_embed.enabled = False + if args.scene_type == "outdoor": + cfg.model.object.sdf.mlp.inside_out = False + cfg.model.object.sdf.encoding.coarse2fine.init_active_level = 8 + elif args.scene_type == "indoor": + cfg.model.object.sdf.mlp.inside_out = True + cfg.model.object.sdf.encoding.coarse2fine.init_active_level = 8 + cfg.model.background.enabled = False + cfg.model.render.num_samples.background = 0 + elif args.scene_type == "object": + cfg.model.object.sdf.mlp.inside_out = False + cfg.model.object.sdf.encoding.coarse2fine.init_active_level = 4 + else: + raise TypeError("Unknown scene type") + # data config + cfg.data.type = "projects.neuralangelo.data" + cfg.data.root = args.data_dir + img = Image.open(os.path.join(args.data_dir, "images", os.listdir(os.path.join(args.data_dir, "images"))[0])) + w, h = img.size + cfg.data.train.image_size = [h, w] + short_size = args.val_short_size + cfg.data.val.image_size = [short_size, int(w/h*short_size)] if w > h else [int(h/w*short_size), short_size] + cfg.data.readjust.center = [0., 0., 0.] + cfg.data.readjust.scale = 1. + # export cfg + cfg_fname = os.path.join(dir_path, "projects/neuralangelo/configs", f"custom/{args.experiment_name}.yaml") + with open(cfg_fname, "w") as file: + yaml.safe_dump(cfg.to_dict(), file, default_flow_style=False, indent=4) + print("Config generated to file: ", cfg_fname) + return + + +if __name__ == "__main__": + parser = ArgumentParser() + parser.add_argument("--experiment_name", type=str, default="recon", help="Name of experiment") + parser.add_argument("--data_dir", type=str, default=None, help="Path to data") + parser.add_argument("--auto_exposure_wb", action="store_true", + help="Video capture with auto-exposure or white-balance") + parser.add_argument("--scene_type", type=str, default="outdoor", choices=["outdoor", "indoor", "object"], + help="Select scene type. Outdoor for building-scale reconstruction; " + "indoor for room-scale reconstruction; object for object-centric scene reconstruction.") + parser.add_argument("--val_short_size", type=int, default=300, + help="Set the short side of validation images (for saving compute when rendering val images)") + args = parser.parse_args() + generate_config(args) diff --git a/neuralangelo-main/projects/neuralangelo/scripts/preprocess.sh b/neuralangelo-main/projects/neuralangelo/scripts/preprocess.sh new file mode 100644 index 0000000000000000000000000000000000000000..204e7f4662f86e5f41d19bf226fb25cb9094420b --- /dev/null +++ b/neuralangelo-main/projects/neuralangelo/scripts/preprocess.sh @@ -0,0 +1,16 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. +# ----------------------------------------------------------------------------- + +# {experiment_name} {path_to_video} {skip_frame_rate} {scene_type} +image_path=${2%.*}_skip${3} +bash projects/neuralangelo/scripts/run_ffmpeg.sh ${2} ${3} +bash projects/neuralangelo/scripts/run_colmap.sh ${image_path} +python3 projects/neuralangelo/scripts/convert_data_to_json.py --data_dir ${image_path}/dense --scene_type ${4} +python3 projects/neuralangelo/scripts/generate_config.py --experiment_name ${1} --data_dir ${image_path}/dense --scene_type ${4} --auto_exposure_wb diff --git a/neuralangelo-main/projects/neuralangelo/scripts/preprocess_dtu.sh b/neuralangelo-main/projects/neuralangelo/scripts/preprocess_dtu.sh new file mode 100644 index 0000000000000000000000000000000000000000..4cb7acaafb7755b8ad0dd7975319bed65f37accc --- /dev/null +++ b/neuralangelo-main/projects/neuralangelo/scripts/preprocess_dtu.sh @@ -0,0 +1,21 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. +# ----------------------------------------------------------------------------- + +# usage: dtu_download.sh + +echo "Download DTU data" +mkdir -p "${1}" +curl -L -o data.zip https://www.dropbox.com/sh/w0y8bbdmxzik3uk/AAAaZffBiJevxQzRskoOYcyja?dl=1 +unzip data.zip "data_DTU.zip" +rm data.zip +unzip -q data_DTU.zip -d ${1} +rm data_DTU.zip +echo "Generate json files" +python3 projects/neuralangelo/scripts/convert_dtu_to_json.py --dtu_path ${1} diff --git a/neuralangelo-main/projects/neuralangelo/scripts/preprocess_tnt.sh b/neuralangelo-main/projects/neuralangelo/scripts/preprocess_tnt.sh new file mode 100644 index 0000000000000000000000000000000000000000..509d58a3c26b4265cee72423304e35412f9b431e --- /dev/null +++ b/neuralangelo-main/projects/neuralangelo/scripts/preprocess_tnt.sh @@ -0,0 +1,20 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. +# ----------------------------------------------------------------------------- + +# usage: tnt_download.sh + +echo "Download fixed poses for Courthouse" +gdown 10pcCwaQY6hqyiegJGdgmLp_HMFOnsmgq +gdown 19TT5aTz3z60eUVIDhFJ7EgGqpcqQnJEi +mv Courthouse_COLMAP_SfM.log ${1}/Courthouse/Courthouse_COLMAP_SfM.log +mv Courthouse_trans.log ${1}/Courthouse/Courthouse_trans.log + +echo "Compute intrinsics, undistort images and generate json files. This may take a while" +python3 projects/neuralangelo/scripts/convert_tnt_to_json.py --tnt_path ${1} diff --git a/neuralangelo-main/projects/neuralangelo/scripts/run_colmap.sh b/neuralangelo-main/projects/neuralangelo/scripts/run_colmap.sh new file mode 100644 index 0000000000000000000000000000000000000000..532f152898952abd3dd51d729828dc9306c6633c --- /dev/null +++ b/neuralangelo-main/projects/neuralangelo/scripts/run_colmap.sh @@ -0,0 +1,40 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. +# ----------------------------------------------------------------------------- + +# usage: run_colmap.sh + +mkdir -p ${1}/sparse +mkdir -p ${1}/dense + +colmap feature_extractor \ + --database_path ${1}/database.db \ + --image_path ${1}/raw_images \ + --ImageReader.camera_model=RADIAL \ + --SiftExtraction.use_gpu=true \ + --SiftExtraction.num_threads=32 \ + --ImageReader.single_camera=true # assuming single camera + +colmap sequential_matcher \ + --database_path ${1}/database.db \ + --SiftMatching.use_gpu=true + +colmap mapper \ + --database_path ${1}/database.db \ + --image_path ${1}/raw_images \ + --output_path ${1}/sparse + +colmap image_undistorter \ + --image_path ${1}/raw_images \ + --input_path ${1}/sparse/0 \ + --output_path ${1}/dense \ + --output_type COLMAP \ + --max_image_size 2000 + +rm -rf ${1}/sparse diff --git a/neuralangelo-main/projects/neuralangelo/scripts/run_ffmpeg.sh b/neuralangelo-main/projects/neuralangelo/scripts/run_ffmpeg.sh new file mode 100644 index 0000000000000000000000000000000000000000..70553f7bef1bbaf8dd0f299df91cf099e6fcfa3a --- /dev/null +++ b/neuralangelo-main/projects/neuralangelo/scripts/run_ffmpeg.sh @@ -0,0 +1,15 @@ +# ----------------------------------------------------------------------------- +# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# +# NVIDIA CORPORATION and its licensors retain all intellectual property +# and proprietary rights in and to this software, related documentation +# and any modifications thereto. Any use, reproduction, disclosure or +# distribution of this software and related documentation without an express +# license agreement from NVIDIA CORPORATION is strictly prohibited. +# ----------------------------------------------------------------------------- + +# usage: run_ffmpeg.sh + +image_path=${1%.*}_skip${2}/raw_images +mkdir -p ${image_path} +ffmpeg -i $1 -vf "select=not(mod(n\,$2))" -vsync vfr -q:v 2 ${image_path}/%06d.jpg diff --git a/neuralangelo-main/projects/neuralangelo/scripts/visualize_colmap.ipynb b/neuralangelo-main/projects/neuralangelo/scripts/visualize_colmap.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..e542d666ee0e2f7ae5e9cc2c3f192e96e630b77b --- /dev/null +++ b/neuralangelo-main/projects/neuralangelo/scripts/visualize_colmap.ipynb @@ -0,0 +1,129 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "8b8d7b17-af50-42cd-b531-ef61c49c9e61", + "metadata": {}, + "outputs": [], + "source": [ + "# Set the work directory to the imaginaire root.\n", + "import os, sys, time\n", + "import pathlib\n", + "root_dir = pathlib.Path().absolute().parents[2]\n", + "os.chdir(root_dir)\n", + "print(f\"Root Directory Path: {root_dir}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2b5b9e2f-841c-4815-92e0-0c76ed46da62", + "metadata": {}, + "outputs": [], + "source": [ + "# Import Python libraries.\n", + "import numpy as np\n", + "import torch\n", + "import k3d\n", + "import json\n", + "from collections import OrderedDict\n", + "# Import imaginaire modules.\n", + "from projects.nerf.utils import camera, visualize\n", + "from third_party.colmap.scripts.python.read_write_model import read_model" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "76033016-2d92-4a5d-9e50-3978553e8df4", + "metadata": {}, + "outputs": [], + "source": [ + "# Read the COLMAP data.\n", + "colmap_path = \"datasets/iphone/climbingnet_skip12\"\n", + "cameras, images, points_3D = read_model(path=f\"{colmap_path}/dense/sparse\", ext=\".bin\")\n", + "# Convert camera poses.\n", + "images = OrderedDict(sorted(images.items()))\n", + "qvecs = torch.from_numpy(np.stack([image.qvec for image in images.values()]))\n", + "tvecs = torch.from_numpy(np.stack([image.tvec for image in images.values()]))\n", + "Rs = camera.quaternion.q_to_R(qvecs)\n", + "poses = torch.cat([Rs, tvecs[..., None]], dim=-1) # [N,3,4]\n", + "print(f\"# images: {len(poses)}\")\n", + "# Get the sparse 3D points and the colors.\n", + "xyzs = torch.from_numpy(np.stack([point.xyz for point in points_3D.values()]))\n", + "rgbs = np.stack([point.rgb for point in points_3D.values()])\n", + "rgbs = (rgbs[:, 0] * 2**16 + rgbs[:, 1] * 2**8 + rgbs[:, 2]).astype(np.uint32)\n", + "print(f\"# points: {len(xyzs)}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b6cf60ec-fe6a-43ba-9aaf-e3c7afd88208", + "metadata": {}, + "outputs": [], + "source": [ + "# Visualize the bounding sphere.\n", + "json_fname = f\"{colmap_path}/dense/transforms.json\"\n", + "with open(json_fname) as file:\n", + " meta = json.load(file)\n", + "center = meta[\"sphere_center\"]\n", + "radius = meta[\"sphere_radius\"]\n", + "# ------------------------------------------------------------------------------------\n", + "# These variables can be adjusted to make the bounding sphere fit the region of interest.\n", + "# The adjusted values can then be set in the config as data.readjust.center and data.readjust.scale\n", + "readjust_center = np.array([0., 0., 0.])\n", + "readjust_scale = 0.25\n", + "# ------------------------------------------------------------------------------------\n", + "center += readjust_center\n", + "radius *= readjust_scale\n", + "# Make some points to hallucinate a bounding sphere.\n", + "sphere_points = np.random.randn(100000, 3)\n", + "sphere_points = sphere_points / np.linalg.norm(sphere_points, axis=-1, keepdims=True)\n", + "sphere_points = sphere_points * radius + center" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fdde170b-4546-4617-9162-a9fcb936347d", + "metadata": {}, + "outputs": [], + "source": [ + "# Visualize with K3D.\n", + "vis_scale = 0.5\n", + "plot = visualize.k3d_visualize_pose(poses,\n", + " vis_depth=(0.5 * vis_scale),\n", + " xyz_length=(0.1 * vis_scale),\n", + " center_size=(0.05 * vis_scale),\n", + " xyz_width=(0.02 * vis_scale))\n", + "plot += k3d.points(xyzs, colors=rgbs, point_size=(0.05 * vis_scale), shader=\"flat\")\n", + "plot += k3d.points(sphere_points, color=0x4488ff, point_size=0.02, shader=\"flat\")\n", + "plot.display()\n", + "plot.camera_fov = 30.0" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.9.13" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/neuralangelo-main/projects/neuralangelo/trainer.py b/neuralangelo-main/projects/neuralangelo/trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..d90849af409dae82f15e62b6b24f1de6899a054e --- /dev/null +++ b/neuralangelo-main/projects/neuralangelo/trainer.py @@ -0,0 +1,106 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import torch +import torch.nn.functional as torch_F +import wandb + +from imaginaire.utils.distributed import master_only +from imaginaire.utils.visualization import wandb_image +from projects.nerf.trainers.base import BaseTrainer +from projects.neuralangelo.utils.misc import get_scheduler, eikonal_loss, curvature_loss + + +class Trainer(BaseTrainer): + + def __init__(self, cfg, is_inference=True, seed=0): + super().__init__(cfg, is_inference=is_inference, seed=seed) + self.metrics = dict() + self.warm_up_end = cfg.optim.sched.warm_up_end + self.cfg_gradient = cfg.model.object.sdf.gradient + if cfg.model.object.sdf.encoding.type == "hashgrid" and cfg.model.object.sdf.encoding.coarse2fine.enabled: + self.c2f_step = cfg.model.object.sdf.encoding.coarse2fine.step + self.model.module.neural_sdf.warm_up_end = self.warm_up_end + + def _init_loss(self, cfg): + self.criteria["render"] = torch.nn.L1Loss() + + def setup_scheduler(self, cfg, optim): + return get_scheduler(cfg.optim, optim) + + def _compute_loss(self, data, mode=None): + if mode == "train": + # Compute loss only on randomly sampled rays. + self.losses["render"] = self.criteria["render"](data["rgb"], data["image_sampled"]) * 3 # FIXME:sumRGB?! + self.metrics["psnr"] = -10 * torch_F.mse_loss(data["rgb"], data["image_sampled"]).log10() + if "eikonal" in self.weights.keys(): + self.losses["eikonal"] = eikonal_loss(data["gradients"], outside=data["outside"]) + if "curvature" in self.weights: + self.losses["curvature"] = curvature_loss(data["hessians"], outside=data["outside"]) + else: + # Compute loss on the entire image. + self.losses["render"] = self.criteria["render"](data["rgb_map"], data["image"]) + self.metrics["psnr"] = -10 * torch_F.mse_loss(data["rgb_map"], data["image"]).log10() + + def get_curvature_weight(self, current_iteration, init_weight, decay_factor): + if "curvature" in self.weights: + if current_iteration <= self.warm_up_end: + self.weights["curvature"] = current_iteration / self.warm_up_end * init_weight + else: + self.weights["curvature"] = init_weight / decay_factor + + def _start_of_iteration(self, data, current_iteration): + model = self.model_module + self.progress = model.progress = current_iteration / self.cfg.max_iter + if self.cfg.model.object.sdf.encoding.coarse2fine.enabled: + model.neural_sdf.set_active_levels(current_iteration) + if self.cfg_gradient.mode == "numerical": + model.neural_sdf.set_normal_epsilon() + decay_factor = model.neural_sdf.growth_rate ** model.neural_sdf.add_levels # TODO: verify? + self.get_curvature_weight(current_iteration, self.cfg.trainer.loss_weight.curvature, decay_factor) + return super()._start_of_iteration(data, current_iteration) + + @master_only + def log_wandb_scalars(self, data, mode=None): + super().log_wandb_scalars(data, mode=mode) + scalars = { + f"{mode}/PSNR": self.metrics["psnr"].detach(), + f"{mode}/s-var": self.model_module.s_var.item(), + } + if "curvature" in self.weights: + scalars[f"{mode}/curvature_weight"] = self.weights["curvature"] + if "eikonal" in self.weights: + scalars[f"{mode}/eikonal_weight"] = self.weights["eikonal"] + if mode == "train" and self.cfg_gradient.mode == "numerical": + scalars[f"{mode}/epsilon"] = self.model.module.neural_sdf.normal_eps + if self.cfg.model.object.sdf.encoding.coarse2fine.enabled: + scalars[f"{mode}/active_levels"] = self.model.module.neural_sdf.active_levels + wandb.log(scalars, step=self.current_iteration) + + @master_only + def log_wandb_images(self, data, mode=None, max_samples=None): + images = {"iteration": self.current_iteration, "epoch": self.current_epoch} + if mode == "val": + images_error = (data["rgb_map"] - data["image"]).abs() + images.update({ + f"{mode}/vis/rgb_target": wandb_image(data["image"]), + f"{mode}/vis/rgb_render": wandb_image(data["rgb_map"]), + f"{mode}/vis/rgb_error": wandb_image(images_error), + f"{mode}/vis/normal": wandb_image(data["normal_map"], from_range=(-1, 1)), + f"{mode}/vis/inv_depth": wandb_image(1 / (data["depth_map"] + 1e-8) * self.cfg.trainer.depth_vis_scale), + f"{mode}/vis/opacity": wandb_image(data["opacity_map"]), + }) + wandb.log(images, step=self.current_iteration) + + def train(self, cfg, data_loader, single_gpu=False, profile=False, show_pbar=False): + self.progress = self.model_module.progress = self.current_iteration / self.cfg.max_iter + super().train(cfg, data_loader, single_gpu, profile, show_pbar) diff --git a/neuralangelo-main/projects/neuralangelo/utils/mesh.py b/neuralangelo-main/projects/neuralangelo/utils/mesh.py new file mode 100644 index 0000000000000000000000000000000000000000..1e676d7a558f8bb6342d0474d95dbd70d429a890 --- /dev/null +++ b/neuralangelo-main/projects/neuralangelo/utils/mesh.py @@ -0,0 +1,115 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import numpy as np +import trimesh +import mcubes +import torch +import torch.distributed as dist +from tqdm import tqdm + +from imaginaire.utils.distributed import get_world_size, is_master + + +@torch.no_grad() +def extract_mesh(sdf_func, bounds, intv, block_res=64): + lattice_grid = LatticeGrid(bounds, intv=intv, block_res=block_res) + data_loader = get_lattice_grid_loader(lattice_grid) + mesh_blocks = [] + if is_master(): + data_loader = tqdm(data_loader, leave=False) + for it, data in enumerate(data_loader): + xyz = data["xyz"][0] + xyz_cuda = xyz.cuda() + sdf_cuda = sdf_func(xyz_cuda)[..., 0] + sdf = sdf_cuda.cpu() + mesh = marching_cubes(sdf.numpy(), xyz.numpy(), intv) + mesh_blocks.append(mesh) + mesh_blocks_gather = [None] * get_world_size() + dist.all_gather_object(mesh_blocks_gather, mesh_blocks) + if is_master(): + mesh_blocks_all = [mesh for mesh_blocks in mesh_blocks_gather for mesh in mesh_blocks] + mesh = trimesh.util.concatenate(mesh_blocks_all) + return mesh + else: + return None + + +class LatticeGrid(torch.utils.data.Dataset): + + def __init__(self, bounds, intv, block_res=64): + super().__init__() + self.block_res = block_res + ((x_min, x_max), (y_min, y_max), (z_min, z_max)) = bounds + self.x_grid = torch.arange(x_min, x_max, intv) + self.y_grid = torch.arange(y_min, y_max, intv) + self.z_grid = torch.arange(z_min, z_max, intv) + res_x, res_y, res_z = len(self.x_grid), len(self.y_grid), len(self.z_grid) + print("Extracting surface at resolution", res_x, res_y, res_z) + self.num_blocks_x = int(np.ceil(res_x / block_res)) + self.num_blocks_y = int(np.ceil(res_y / block_res)) + self.num_blocks_z = int(np.ceil(res_z / block_res)) + + def __getitem__(self, idx): + # Keep track of sample index for convenience. + sample = dict(idx=idx) + block_idx_x = idx // (self.num_blocks_y * self.num_blocks_z) + block_idx_y = (idx // self.num_blocks_z) % self.num_blocks_y + block_idx_z = idx % self.num_blocks_z + xi = block_idx_x * self.block_res + yi = block_idx_y * self.block_res + zi = block_idx_z * self.block_res + x, y, z = torch.meshgrid(self.x_grid[xi:xi+self.block_res+1], + self.y_grid[yi:yi+self.block_res+1], + self.z_grid[zi:zi+self.block_res+1], indexing="ij") + xyz = torch.stack([x, y, z], dim=-1) + sample.update(xyz=xyz) + return sample + + def __len__(self): + return self.num_blocks_x * self.num_blocks_y * self.num_blocks_z + + +def get_lattice_grid_loader(dataset, num_workers=8): + if dist.is_initialized(): + sampler = torch.utils.data.distributed.DistributedSampler(dataset, shuffle=False) + else: + sampler = None + return torch.utils.data.DataLoader( + dataset, + batch_size=1, + shuffle=False, + sampler=sampler, + pin_memory=True, + num_workers=num_workers, + drop_last=False + ) + + +def marching_cubes(sdf, xyz, intv): + # marching cubes + V, F = mcubes.marching_cubes(sdf, 0.) + V = V * intv + xyz[0, 0, 0] + mesh = trimesh.Trimesh(V, F) + mesh = filter_points_outside_bounding_sphere(mesh) + return mesh + + +def filter_points_outside_bounding_sphere(old_mesh): + mask = np.linalg.norm(old_mesh.vertices, axis=-1) < 1.0 + indices = np.ones(len(old_mesh.vertices), dtype=int) * -1 + indices[mask] = np.arange(mask.sum()) + faces_mask = mask[old_mesh.faces[:, 0]] & mask[old_mesh.faces[:, 1]] & mask[old_mesh.faces[:, 2]] + new_faces = indices[old_mesh.faces[faces_mask]] + new_vertices = old_mesh.vertices[mask] + new_mesh = trimesh.Trimesh(new_vertices, new_faces) + return new_mesh diff --git a/neuralangelo-main/projects/neuralangelo/utils/misc.py b/neuralangelo-main/projects/neuralangelo/utils/misc.py new file mode 100644 index 0000000000000000000000000000000000000000..157b57bd86928456bf09a373cfcd7689393cedea --- /dev/null +++ b/neuralangelo-main/projects/neuralangelo/utils/misc.py @@ -0,0 +1,127 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +from functools import partial +import numpy as np +import torch +import torch.nn.functional as torch_F +import imaginaire.trainers.utils +from torch.optim import lr_scheduler + +flip_mat = np.array([ + [1, 0, 0, 0], + [0, -1, 0, 0], + [0, 0, -1, 0], + [0, 0, 0, 1] +]) + + +def cv_to_gl(cv): + gl = cv @ flip_mat # convert to GL convention used in iNGP + return gl + + +def gl_to_cv(gl): + cv = gl @ np.linalg.inv(flip_mat) # convert to CV convention used in Imaginaire + return cv + + +def get_scheduler(cfg_opt, opt): + """Return the scheduler object. + + Args: + cfg_opt (obj): Config for the specific optimization module (gen/dis). + opt (obj): PyTorch optimizer object. + + Returns: + (obj): Scheduler + """ + if cfg_opt.sched.type == 'two_steps_with_warmup': + warm_up_end = cfg_opt.sched.warm_up_end + two_steps = cfg_opt.sched.two_steps + gamma = cfg_opt.sched.gamma + + def sch(x): + if x < warm_up_end: + return x / warm_up_end + else: + if x > two_steps[1]: + return 1.0 / gamma ** 2 + elif x > two_steps[0]: + return 1.0 / gamma + else: + return 1.0 + + scheduler = lr_scheduler.LambdaLR(opt, lambda x: sch(x)) + elif cfg_opt.sched.type == 'cos_with_warmup': + alpha = cfg_opt.sched.alpha + max_iter = cfg_opt.sched.max_iter + warm_up_end = cfg_opt.sched.warm_up_end + + def sch(x): + if x < warm_up_end: + return x / warm_up_end + else: + progress = (x - warm_up_end) / (max_iter - warm_up_end) + learning_factor = (np.cos(np.pi * progress) + 1.0) * 0.5 * (1 - alpha) + alpha + return learning_factor + + scheduler = lr_scheduler.LambdaLR(opt, lambda x: sch(x)) + else: + return imaginaire.trainers.utils.get_scheduler() + return scheduler + + +def eikonal_loss(gradients, outside=None): + gradient_error = (gradients.norm(dim=-1) - 1.0) ** 2 # [B,R,N] + gradient_error = gradient_error.nan_to_num(nan=0.0, posinf=0.0, neginf=0.0) # [B,R,N] + if outside is not None: + return (gradient_error * (~outside).float()).mean() + else: + return gradient_error.mean() + + +def curvature_loss(hessian, outside=None): + laplacian = hessian.sum(dim=-1).abs() # [B,R,N] + laplacian = laplacian.nan_to_num(nan=0.0, posinf=0.0, neginf=0.0) # [B,R,N] + if outside is not None: + return (laplacian * (~outside).float()).mean() + else: + return laplacian.mean() + + +def get_activation(activ, **kwargs): + func = dict( + identity=lambda x: x, + relu=torch_F.relu, + relu_=torch_F.relu_, + abs=torch.abs, + abs_=torch.abs_, + sigmoid=torch.sigmoid, + sigmoid_=torch.sigmoid_, + exp=torch.exp, + exp_=torch.exp_, + softplus=torch_F.softplus, + silu=torch_F.silu, + silu_=partial(torch_F.silu, inplace=True), + )[activ] + return partial(func, **kwargs) + + +def to_full_image(image, image_size=None, from_vec=True): + # if from_vec is True: [B,HW,...,K] --> [B,K,H,W,...] + # if from_vec is False: [B,H,W,...,K] --> [B,K,H,W,...] + if from_vec: + assert image_size is not None + image = image.unflatten(dim=1, sizes=image_size) + image = image.moveaxis(-1, 1) + return image diff --git a/neuralangelo-main/projects/neuralangelo/utils/mlp.py b/neuralangelo-main/projects/neuralangelo/utils/mlp.py new file mode 100644 index 0000000000000000000000000000000000000000..43a4ad71c0955528b97bc1fb2f944d9b238f4afb --- /dev/null +++ b/neuralangelo-main/projects/neuralangelo/utils/mlp.py @@ -0,0 +1,84 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import numpy as np +import torch +import torch.nn.functional as torch_F + + +class MLPforNeuralSDF(torch.nn.Module): + + def __init__(self, layer_dims, skip_connection=[], activ=None, use_layernorm=False, use_weightnorm=False, + geometric_init=False, out_bias=0., invert=False): + """Initialize a multi-layer perceptron with skip connection. + Args: + layer_dims: A list of integers representing the number of channels in each layer. + skip_connection: A list of integers representing the index of layers to add skip connection. + """ + super().__init__() + self.skip_connection = skip_connection + self.use_layernorm = use_layernorm + self.linears = torch.nn.ModuleList() + if use_layernorm: + self.layer_norm = torch.nn.ModuleList() + # Hidden layers + layer_dim_pairs = list(zip(layer_dims[:-1], layer_dims[1:])) + for li, (k_in, k_out) in enumerate(layer_dim_pairs): + if li in self.skip_connection: + k_in += layer_dims[0] + linear = torch.nn.Linear(k_in, k_out) + if geometric_init: + self._geometric_init(linear, k_in, k_out, first=(li == 0), + skip_dim=(layer_dims[0] if li in self.skip_connection else 0)) + if use_weightnorm: + linear = torch.nn.utils.weight_norm(linear) + self.linears.append(linear) + if use_layernorm and li != len(layer_dim_pairs) - 1: + self.layer_norm.append(torch.nn.LayerNorm(k_out)) + if li == len(layer_dim_pairs) - 1: + self.linears[-1].bias.data.fill_(0.0) + # SDF prediction layer + self.linear_sdf = torch.nn.Linear(k_in, 1) + if geometric_init: + self._geometric_init_sdf(self.linear_sdf, k_in, out_bias=out_bias, invert=invert) + self.activ = activ or torch_F.relu_ + + def forward(self, input, with_sdf=True, with_feat=True): + feat = input + for li, linear in enumerate(self.linears): + if li in self.skip_connection: + feat = torch.cat([feat, input], dim=-1) + if li != len(self.linears) - 1 or with_feat: + feat_pre = linear(feat) + if self.use_layernorm: + feat_pre = self.layer_norm[li](feat_pre) + feat_activ = self.activ(feat_pre) + if li == len(self.linears) - 1: + out = [self.linear_sdf(feat) if with_sdf else None, + feat_activ if with_feat else None] + feat = feat_activ + return out + + def _geometric_init(self, linear, k_in, k_out, first=False, skip_dim=0): + torch.nn.init.constant_(linear.bias, 0.0) + torch.nn.init.normal_(linear.weight, 0.0, np.sqrt(2 / k_out)) + if first: + torch.nn.init.constant_(linear.weight[:, 3:], 0.0) # positional encodings + if skip_dim: + torch.nn.init.constant_(linear.weight[:, -skip_dim:], 0.0) # skip connections + + def _geometric_init_sdf(self, linear, k_in, out_bias=0., invert=False): + torch.nn.init.normal_(linear.weight, mean=np.sqrt(np.pi / k_in), std=0.0001) + torch.nn.init.constant_(linear.bias, -out_bias) + if invert: + linear.weight.data *= -1 + linear.bias.data *= -1 diff --git a/neuralangelo-main/projects/neuralangelo/utils/modules.py b/neuralangelo-main/projects/neuralangelo/utils/modules.py new file mode 100644 index 0000000000000000000000000000000000000000..034a1eec9600ff4648e5041f32c30f4c8b060be8 --- /dev/null +++ b/neuralangelo-main/projects/neuralangelo/utils/modules.py @@ -0,0 +1,311 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import torch +from functools import partial +import numpy as np +import tinycudann as tcnn + +from projects.neuralangelo.utils.spherical_harmonics import get_spherical_harmonics +from projects.neuralangelo.utils.mlp import MLPforNeuralSDF +from projects.neuralangelo.utils.misc import get_activation +from projects.nerf.utils import nerf_util + + +class NeuralSDF(torch.nn.Module): + + def __init__(self, cfg_sdf): + super().__init__() + self.cfg_sdf = cfg_sdf + encoding_dim = self.build_encoding(cfg_sdf.encoding) + input_dim = 3 + encoding_dim + self.build_mlp(cfg_sdf.mlp, input_dim=input_dim) + + def build_encoding(self, cfg_encoding): + if cfg_encoding.type == "fourier": + encoding_dim = 6 * cfg_encoding.levels + elif cfg_encoding.type == "hashgrid": + # Build the multi-resolution hash grid. + l_min, l_max = cfg_encoding.hashgrid.min_logres, cfg_encoding.hashgrid.max_logres + r_min, r_max = 2 ** l_min, 2 ** l_max + num_levels = cfg_encoding.levels + self.growth_rate = np.exp((np.log(r_max) - np.log(r_min)) / (num_levels - 1)) + config = dict( + otype="HashGrid", + n_levels=cfg_encoding.levels, + n_features_per_level=cfg_encoding.hashgrid.dim, + log2_hashmap_size=cfg_encoding.hashgrid.dict_size, + base_resolution=2 ** cfg_encoding.hashgrid.min_logres, + per_level_scale=self.growth_rate, + ) + self.tcnn_encoding = tcnn.Encoding(3, config) + self.resolutions = [] + for lv in range(0, num_levels): + size = np.floor(r_min * self.growth_rate ** lv).astype(int) + 1 + self.resolutions.append(size) + encoding_dim = cfg_encoding.hashgrid.dim * cfg_encoding.levels + else: + raise NotImplementedError("Unknown encoding type") + return encoding_dim + + def build_mlp(self, cfg_mlp, input_dim=3): + # SDF + point-wise feature + layer_dims = [input_dim] + [cfg_mlp.hidden_dim] * cfg_mlp.num_layers + activ = get_activation(cfg_mlp.activ, **cfg_mlp.activ_params) + self.mlp = MLPforNeuralSDF(layer_dims, skip_connection=cfg_mlp.skip, activ=activ, + use_weightnorm=cfg_mlp.weight_norm, geometric_init=cfg_mlp.geometric_init, + out_bias=cfg_mlp.out_bias, invert=cfg_mlp.inside_out) + + def forward(self, points_3D, with_sdf=True, with_feat=True): + points_enc = self.encode(points_3D) # [...,3+LD] + sdf, feat = self.mlp(points_enc, with_sdf=with_sdf, with_feat=with_feat) + return sdf, feat # [...,1],[...,K] + + def sdf(self, points_3D): + return self.forward(points_3D, with_sdf=True, with_feat=False)[0] + + def encode(self, points_3D): + if self.cfg_sdf.encoding.type == "fourier": + points_enc = nerf_util.positional_encoding(points_3D, num_freq_bases=self.cfg_sdf.encoding.levels) + feat_dim = 6 + elif self.cfg_sdf.encoding.type == "hashgrid": + # Tri-linear interpolate the corresponding embeddings from the dictionary. + vol_min, vol_max = self.cfg_sdf.encoding.hashgrid.range + points_3D_normalized = (points_3D - vol_min) / (vol_max - vol_min) # Normalize to [0,1]. + tcnn_input = points_3D_normalized.view(-1, 3) + tcnn_output = self.tcnn_encoding(tcnn_input) + points_enc = tcnn_output.view(*points_3D_normalized.shape[:-1], tcnn_output.shape[-1]) + feat_dim = self.cfg_sdf.encoding.hashgrid.dim + else: + raise NotImplementedError("Unknown encoding type") + # Coarse-to-fine. + if self.cfg_sdf.encoding.coarse2fine.enabled: + mask = self._get_coarse2fine_mask(points_enc, feat_dim=feat_dim) + points_enc = points_enc * mask + points_enc = torch.cat([points_3D, points_enc], dim=-1) # [B,R,N,3+LD] + return points_enc + + def set_active_levels(self, current_iter=None): + add_levels = (current_iter - self.warm_up_end) // self.cfg_sdf.encoding.coarse2fine.step + self.add_levels = min(self.cfg_sdf.encoding.levels - self.cfg_sdf.encoding.coarse2fine.init_active_level, + add_levels) + self.active_levels = self.cfg_sdf.encoding.coarse2fine.init_active_level + self.add_levels + assert self.active_levels <= self.cfg_sdf.encoding.levels + + def set_normal_epsilon(self): + epsilon_res = self.resolutions[self.active_levels - 1] + self.normal_eps = 1. / epsilon_res + + @torch.no_grad() + def _get_coarse2fine_mask(self, points_enc, feat_dim): + mask = torch.zeros_like(points_enc) + mask[..., :(self.active_levels * feat_dim)] = 1 + return mask + + def compute_gradients(self, x, compute_hessian=False, sdf=None): + # Note: hessian is not fully hessian but diagonal elements + if self.cfg_sdf.gradient.mode == "analytical": + requires_grad = x.requires_grad + with torch.enable_grad(): + # 1st-order gradient + x.requires_grad_(True) + y = self.sdf(x) + gradient = torch.autograd.grad(y.sum(), x, create_graph=True)[0] + # 2nd-order gradient (hessian) + if compute_hessian: + hessian = torch.autograd.grad(gradient.sum(), x, create_graph=True)[0] + else: + hessian = None + x.requires_grad_(requires_grad) + if not requires_grad: + gradient = gradient.detach() + if compute_hessian: + hessian = hessian.detach() + elif self.cfg_sdf.gradient.mode == "numerical": + eps = self.normal_eps + if self.cfg_sdf.gradient.taps == 6: + # 1st-order gradient + eps_x = torch.tensor([eps, 0., 0.], dtype=x.dtype, device=x.device) # [3] + eps_y = torch.tensor([0., eps, 0.], dtype=x.dtype, device=x.device) # [3] + eps_z = torch.tensor([0., 0., eps], dtype=x.dtype, device=x.device) # [3] + sdf_x_pos = self.sdf(x + eps_x) # [...,1] + sdf_x_neg = self.sdf(x - eps_x) # [...,1] + sdf_y_pos = self.sdf(x + eps_y) # [...,1] + sdf_y_neg = self.sdf(x - eps_y) # [...,1] + sdf_z_pos = self.sdf(x + eps_z) # [...,1] + sdf_z_neg = self.sdf(x - eps_z) # [...,1] + gradient_x = (sdf_x_pos - sdf_x_neg) / (2 * eps) + gradient_y = (sdf_y_pos - sdf_y_neg) / (2 * eps) + gradient_z = (sdf_z_pos - sdf_z_neg) / (2 * eps) + gradient = torch.cat([gradient_x, gradient_y, gradient_z], dim=-1) # [...,3] + # 2nd-order gradient (hessian) + if compute_hessian: + assert sdf is not None # computed when feed-forwarding through the network + hessian_xx = (sdf_x_pos + sdf_x_neg - 2 * sdf) / (eps ** 2) # [...,1] + hessian_yy = (sdf_y_pos + sdf_y_neg - 2 * sdf) / (eps ** 2) # [...,1] + hessian_zz = (sdf_z_pos + sdf_z_neg - 2 * sdf) / (eps ** 2) # [...,1] + hessian = torch.cat([hessian_xx, hessian_yy, hessian_zz], dim=-1) # [...,3] + else: + hessian = None + elif self.cfg_sdf.gradient.taps == 4: + k1 = torch.tensor([1, -1, -1], dtype=x.dtype, device=x.device) # [3] + k2 = torch.tensor([-1, -1, 1], dtype=x.dtype, device=x.device) # [3] + k3 = torch.tensor([-1, 1, -1], dtype=x.dtype, device=x.device) # [3] + k4 = torch.tensor([1, 1, 1], dtype=x.dtype, device=x.device) # [3] + sdf1 = self.sdf(x + k1 * eps) # [...,1] + sdf2 = self.sdf(x + k2 * eps) # [...,1] + sdf3 = self.sdf(x + k3 * eps) # [...,1] + sdf4 = self.sdf(x + k4 * eps) # [...,1] + gradient = (k1*sdf1 + k2*sdf2 + k3*sdf3 + k4*sdf4) / (4.0 * eps) + if compute_hessian: + # the result of 4 taps is directly trace, but we assume they are individual components + # so we use the same signature as 6 taps + hessian_xx = ((sdf1 + sdf2 + sdf3 + sdf4) / 2.0 - 2 * sdf) / eps ** 2 # [N,1] + hessian = torch.cat([hessian_xx, hessian_xx, hessian_xx], dim=-1) / 3.0 + else: + hessian = None + else: + raise ValueError("Only support 4 or 6 taps.") + return gradient, hessian + + +class NeuralRGB(torch.nn.Module): + + def __init__(self, cfg_rgb, feat_dim, appear_embed): + super().__init__() + self.cfg_rgb = cfg_rgb + self.cfg_appear_embed = appear_embed + encoding_view_dim = self.build_encoding(cfg_rgb.encoding_view) + input_base_dim = 6 if cfg_rgb.mode == "idr" else 3 + input_dim = input_base_dim + encoding_view_dim + feat_dim + (appear_embed.dim if appear_embed.enabled else 0) + self.build_mlp(cfg_rgb.mlp, input_dim=input_dim) + + def build_encoding(self, cfg_encoding_view): + if cfg_encoding_view.type == "fourier": + encoding_view_dim = 6 * cfg_encoding_view.levels + elif cfg_encoding_view.type == "spherical": + self.spherical_harmonic_encoding = partial(get_spherical_harmonics, levels=cfg_encoding_view.levels) + encoding_view_dim = (cfg_encoding_view.levels + 1) ** 2 + else: + raise NotImplementedError("Unknown encoding type") + return encoding_view_dim + + def build_mlp(self, cfg_mlp, input_dim=3): + # RGB prediction + layer_dims = [input_dim] + [cfg_mlp.hidden_dim] * cfg_mlp.num_layers + [3] + activ = get_activation(cfg_mlp.activ, **cfg_mlp.activ_params) + self.mlp = nerf_util.MLPwithSkipConnection(layer_dims, skip_connection=cfg_mlp.skip, activ=activ, + use_weightnorm=cfg_mlp.weight_norm) + + def forward(self, points_3D, normals, rays_unit, feats, app): + view_enc = self.encode_view(rays_unit) # [...,LD] + input_list = [points_3D, view_enc, normals, feats] + if app is not None: + input_list.append(app) + if self.cfg_rgb.mode == "no_view_dir": + input_list.remove(view_enc) + if self.cfg_rgb.mode == "no_normal": + input_list.remove(normals) + input_vec = torch.cat(input_list, dim=-1) + rgb = self.mlp(input_vec).sigmoid_() + return rgb # [...,3] + + def encode_view(self, rays_unit): + if self.cfg_rgb.encoding_view.type == "fourier": + view_enc = nerf_util.positional_encoding(rays_unit, num_freq_bases=self.cfg_rgb.encoding_view.levels) + elif self.cfg_rgb.encoding_view.type == "spherical": + view_enc = self.spherical_harmonic_encoding(rays_unit) + else: + raise NotImplementedError("Unknown encoding type") + return view_enc + + +class BackgroundNeRF(torch.nn.Module): + + def __init__(self, cfg_background, appear_embed): + super().__init__() + self.cfg_background = cfg_background + self.cfg_appear_embed = appear_embed + encoding_dim, encoding_view_dim = self.build_encoding(cfg_background.encoding, cfg_background.encoding_view) + input_dim = 4 + encoding_dim + input_view_dim = cfg_background.mlp.hidden_dim + encoding_view_dim + \ + (appear_embed.dim if appear_embed.enabled else 0) + self.build_mlp(cfg_background.mlp, input_dim=input_dim, input_view_dim=input_view_dim) + + def build_encoding(self, cfg_encoding, cfg_encoding_view): + # Positional encoding. + if cfg_encoding.type == "fourier": + encoding_dim = 8 * cfg_encoding.levels + else: + raise NotImplementedError("Unknown encoding type") + # View encoding. + if cfg_encoding_view.type == "fourier": + encoding_view_dim = 6 * cfg_encoding_view.levels + elif cfg_encoding_view.type == "spherical": + self.spherical_harmonic_encoding = partial(get_spherical_harmonics, levels=cfg_encoding_view.levels) + encoding_view_dim = (cfg_encoding_view.levels + 1) ** 2 + else: + raise NotImplementedError("Unknown encoding type") + return encoding_dim, encoding_view_dim + + def build_mlp(self, cfg_mlp, input_dim=3, input_view_dim=3): + activ = get_activation(cfg_mlp.activ, **cfg_mlp.activ_params) + # Point-wise feature. + layer_dims = [input_dim] + [cfg_mlp.hidden_dim] * (cfg_mlp.num_layers - 1) + [cfg_mlp.hidden_dim + 1] + self.mlp_feat = nerf_util.MLPwithSkipConnection(layer_dims, skip_connection=cfg_mlp.skip, activ=activ) + self.activ_density = get_activation(cfg_mlp.activ_density, **cfg_mlp.activ_density_params) + # RGB prediction. + layer_dims_rgb = [input_view_dim] + [cfg_mlp.hidden_dim_rgb] * (cfg_mlp.num_layers_rgb - 1) + [3] + self.mlp_rgb = nerf_util.MLPwithSkipConnection(layer_dims_rgb, skip_connection=cfg_mlp.skip_rgb, activ=activ) + + def forward(self, points_3D, rays_unit, app_outside): + points_enc = self.encode(points_3D) # [...,4+LD] + # Volume density prediction. + out = self.mlp_feat(points_enc) + density, feat = self.activ_density(out[..., 0]), self.mlp_feat.activ(out[..., 1:]) # [...],[...,K] + # RGB color prediction. + if self.cfg_background.view_dep: + view_enc = self.encode_view(rays_unit) # [...,LD] + input_list = [feat, view_enc] + if app_outside is not None: + input_list.append(app_outside) + input_vec = torch.cat(input_list, dim=-1) + rgb = self.mlp_rgb(input_vec).sigmoid_() # [...,3] + else: + raise NotImplementedError + return rgb, density + + def encode(self, points_3D): + # Reparametrize the 3D points. + # TODO: revive this. + if True: + points_3D_norm = points_3D.norm(dim=-1, keepdim=True) # [B,R,N,1] + points = torch.cat([points_3D / points_3D_norm, 1.0 / points_3D_norm], dim=-1) # [B,R,N,4] + else: + points = points_3D + # Positional encoding. + if self.cfg_background.encoding.type == "fourier": + points_enc = nerf_util.positional_encoding(points, num_freq_bases=self.cfg_background.encoding.levels) + else: + raise NotImplementedError("Unknown encoding type") + # TODO: 1/x? + points_enc = torch.cat([points, points_enc], dim=-1) # [B,R,N,4+LD] + return points_enc + + def encode_view(self, rays_unit): + if self.cfg_background.encoding_view.type == "fourier": + view_enc = nerf_util.positional_encoding(rays_unit, num_freq_bases=self.cfg_background.encoding_view.levels) + elif self.cfg_background.encoding_view.type == "spherical": + view_enc = self.spherical_harmonic_encoding(rays_unit) + else: + raise NotImplementedError("Unknown encoding type") + return view_enc diff --git a/neuralangelo-main/projects/neuralangelo/utils/spherical_harmonics.py b/neuralangelo-main/projects/neuralangelo/utils/spherical_harmonics.py new file mode 100644 index 0000000000000000000000000000000000000000..cca8cd8362170def791bbd25d55e4b911fdc7a40 --- /dev/null +++ b/neuralangelo-main/projects/neuralangelo/utils/spherical_harmonics.py @@ -0,0 +1,84 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import torch + + +SH_C0 = 0.28209479177387814 +SH_C1 = 0.4886025119029199 +SH_C2 = [ + 1.0925484305920792, + -1.0925484305920792, + 0.31539156525252005, + -1.0925484305920792, + 0.5462742152960396 +] +SH_C3 = [ + -0.5900435899266435, + 2.890611442640554, + -0.4570457994644658, + 0.3731763325901154, + -0.4570457994644658, + 1.445305721320277, + -0.5900435899266435 +] +SH_C4 = [ + 2.5033429417967046, + -1.7701307697799304, + 0.9461746957575601, + -0.6690465435572892, + 0.10578554691520431, + -0.6690465435572892, + 0.47308734787878004, + -1.7701307697799304, + 0.6258357354491761, +] + + +def get_spherical_harmonics(dirs, levels): + # Evaluate spherical harmonics bases at unit directions, without taking linear combination. + vals = torch.empty((*dirs.shape[:-1], (levels + 1) ** 2), device=dirs.device) + vals[..., 0] = SH_C0 + if levels >= 1: + x, y, z = dirs.unbind(-1) + vals[..., 1] = -SH_C1 * y + vals[..., 2] = SH_C1 * z + vals[..., 3] = -SH_C1 * x + if levels >= 2: + xx, yy, zz = x * x, y * y, z * z + xy, yz, xz = x * y, y * z, x * z + vals[..., 4] = SH_C2[0] * xy + vals[..., 5] = SH_C2[1] * yz + vals[..., 6] = SH_C2[2] * (2.0 * zz - xx - yy) + vals[..., 7] = SH_C2[3] * xz + vals[..., 8] = SH_C2[4] * (xx - yy) + if levels >= 3: + vals[..., 9] = SH_C3[0] * y * (3 * xx - yy) + vals[..., 10] = SH_C3[1] * xy * z + vals[..., 11] = SH_C3[2] * y * (4 * zz - xx - yy) + vals[..., 12] = SH_C3[3] * z * (2 * zz - 3 * xx - 3 * yy) + vals[..., 13] = SH_C3[4] * x * (4 * zz - xx - yy) + vals[..., 14] = SH_C3[5] * z * (xx - yy) + vals[..., 15] = SH_C3[6] * x * (xx - 3 * yy) + if levels >= 4: + vals[..., 16] = SH_C4[0] * xy * (xx - yy) + vals[..., 17] = SH_C4[1] * yz * (3 * xx - yy) + vals[..., 18] = SH_C4[2] * xy * (7 * zz - 1) + vals[..., 19] = SH_C4[3] * yz * (7 * zz - 3) + vals[..., 20] = SH_C4[4] * (zz * (35 * zz - 30) + 3) + vals[..., 21] = SH_C4[5] * xz * (7 * zz - 3) + vals[..., 22] = SH_C4[6] * (xx - yy) * (7 * zz - 1) + vals[..., 23] = SH_C4[7] * xz * (xx - 3 * yy) + vals[..., 24] = SH_C4[8] * (xx * (xx - 3 * yy) - yy * (3 * xx - yy)) + if levels >= 5: + raise NotImplementedError + return vals diff --git a/neuralangelo-main/requirements.txt b/neuralangelo-main/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..583c9459d5ae6c7657b426198b4508661c83818f --- /dev/null +++ b/neuralangelo-main/requirements.txt @@ -0,0 +1,34 @@ +addict +gdown +git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch +gpustat +icecream +imageio-ffmpeg +imutils +ipdb +k3d +kornia +lpips +matplotlib +mediapy +nvidia-ml-py3 +open3d +opencv-python-headless +OpenEXR +pathlib +pillow +pyequilib +pyexr +PyMCubes +pyquaternion +pyyaml +requests +scikit-image +scikit-video +scipy +seaborn +tensorboard +termcolor +tqdm +trimesh +wandb diff --git a/neuralangelo-main/third_party/.DS_Store b/neuralangelo-main/third_party/.DS_Store new file mode 100644 index 0000000000000000000000000000000000000000..f5b78d12ebebf63f46db8af11104e7220079af25 Binary files /dev/null and b/neuralangelo-main/third_party/.DS_Store differ diff --git a/neuralangelo-main/train.py b/neuralangelo-main/train.py new file mode 100644 index 0000000000000000000000000000000000000000..6f4384008dac400bb3cd9eb053c7739e85b2d649 --- /dev/null +++ b/neuralangelo-main/train.py @@ -0,0 +1,104 @@ +''' +----------------------------------------------------------------------------- +Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + +NVIDIA CORPORATION and its licensors retain all intellectual property +and proprietary rights in and to this software, related documentation +and any modifications thereto. Any use, reproduction, disclosure or +distribution of this software and related documentation without an express +license agreement from NVIDIA CORPORATION is strictly prohibited. +----------------------------------------------------------------------------- +''' + +import argparse +import os + +import imaginaire.config +from imaginaire.config import Config, recursive_update_strict, parse_cmdline_arguments +from imaginaire.utils.cudnn import init_cudnn +from imaginaire.utils.distributed import init_dist, get_world_size, master_only_print as print, is_master +from imaginaire.utils.gpu_affinity import set_affinity +from imaginaire.trainers.utils.logging import init_logging +from imaginaire.trainers.utils.get_trainer import get_trainer +from imaginaire.utils.set_random_seed import set_random_seed + + +def parse_args(): + parser = argparse.ArgumentParser(description='Training') + parser.add_argument('--config', help='Path to the training config file.', required=True) + parser.add_argument('--logdir', help='Dir for saving logs and models.', default=None) + parser.add_argument('--checkpoint', default=None, help='Checkpoint path.') + parser.add_argument('--seed', type=int, default=0, help='Random seed.') + parser.add_argument('--local_rank', type=int, default=os.getenv('LOCAL_RANK', 0)) + parser.add_argument('--single_gpu', action='store_true') + parser.add_argument('--debug', action='store_true') + parser.add_argument('--profile', action='store_true') + parser.add_argument('--show_pbar', action='store_true') + parser.add_argument('--wandb', action='store_true', help="Enable using Weights & Biases as the logger") + parser.add_argument('--wandb_name', default='default', type=str) + parser.add_argument('--resume', action='store_true') + args, cfg_cmd = parser.parse_known_args() + return args, cfg_cmd + + +def main(): + args, cfg_cmd = parse_args() + set_affinity(args.local_rank) + cfg = Config(args.config) + + cfg_cmd = parse_cmdline_arguments(cfg_cmd) + recursive_update_strict(cfg, cfg_cmd) + + # If args.single_gpu is set to True, we will disable distributed data parallel. + if not args.single_gpu: + # this disables nccl timeout + os.environ["NCLL_BLOCKING_WAIT"] = "0" + os.environ["NCCL_ASYNC_ERROR_HANDLING"] = "0" + cfg.local_rank = args.local_rank + init_dist(cfg.local_rank, rank=-1, world_size=-1) + print(f"Training with {get_world_size()} GPUs.") + + # set random seed by rank + set_random_seed(args.seed, by_rank=True) + + # Global arguments. + imaginaire.config.DEBUG = args.debug + + # Create log directory for storing training results. + cfg.logdir = init_logging(args.config, args.logdir, makedir=True) + + # Print and save final config + if is_master(): + cfg.print_config() + cfg.save_config(cfg.logdir) + + # Initialize cudnn. + init_cudnn(cfg.cudnn.deterministic, cfg.cudnn.benchmark) + + # Initialize data loaders and models. + trainer = get_trainer(cfg, is_inference=False, seed=args.seed) + trainer.set_data_loader(cfg, split="train") + trainer.set_data_loader(cfg, split="val") + trainer.checkpointer.load(args.checkpoint, args.resume, load_sch=True, load_opt=True) + + # Initialize Wandb. + trainer.init_wandb(cfg, + project=args.wandb_name, + mode="disabled" if args.debug or not args.wandb else "online", + resume=args.resume, + use_group=True) + + trainer.mode = 'train' + # Start training. + trainer.train(cfg, + trainer.train_data_loader, + single_gpu=args.single_gpu, + profile=args.profile, + show_pbar=args.show_pbar) + + # Finalize training. + trainer.finalize(cfg) + + +if __name__ == "__main__": + main()