finalize environment

.gitignore

@@ -168,4 +168,3 @@ wandb/
 wdb_sweep.py
 
 zeus_config.yaml
-pycg

LICENSE.txt

@@ -0,0 +1,91 @@
+
+Copyright (c) 2022, NVIDIA Corporation & affiliates. All rights reserved.
+
+
+NVIDIA Source Code License for NKSR
+
+
+=======================================================================
+
+1. Definitions
+
+“Licensor” means any person or entity that distributes its Work.
+
+“Work” means (a) the original work of authorship made available under
+this license, which may include software, documentation, or other files,
+and (b) any additions to or derivative works  thereof  that are made
+available under this license.
+
+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 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 use, 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. Notwithstanding the foregoing,
+     NVIDIA Corporation 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
+MERCHANTABILITY, 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 DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN
+ADVISED OF THE POSSIBILITY OF SUCH DAMAGES.
+
+=======================================================================

README.md

@@ -1,7 +1,109 @@
 # Neural Kernel Surface Reconstruction
 
-> This repo contains the training script for NKSR. If you just want to test it in your project without re-training, please refer to https://github.com/nksr/nksr.
+![NKSR](assets/teaser.png)
 
-Alright.
+[![PyPI version](https://badge.fury.io/py/nksr.svg)](https://badge.fury.io/py/nksr)
 
-# Test data
+**Neural Kernel Surface Reconstruction**<br>
+[Jiahui Huang](https://huangjh-pub.github.io/),
+[Zan Gojcic](https://zgojcic.github.io/),
+[Matan Atzmon](https://matanatz.github.io/),
+[Or Litany](https://orlitany.github.io/), 
+[Sanja Fidler](https://www.cs.toronto.edu/~fidler/),
+[Francis Williams](https://www.fwilliams.info/) <br>
+**[Paper](https://research.nvidia.com/labs/toronto-ai/NKSR/paper.pdf), [Project Page](https://research.nvidia.com/labs/toronto-ai/NKSR/)**
+
+Abstract: *We present a novel method for reconstructing a 3D implicit surface from a large-scale, sparse, and noisy point cloud. 
+Our approach builds upon the recently introduced [Neural Kernel Fields (NKF)](https://nv-tlabs.github.io/nkf/) representation. 
+It enjoys similar generalization capabilities to NKF, while simultaneously addressing its main limitations: 
+(a) We can scale to large scenes through compactly supported kernel functions, which enable the use of memory-efficient sparse linear solvers. 
+(b) We are robust to noise, through a gradient fitting solve. 
+(c) We minimize training requirements, enabling us to learn from any dataset of dense oriented points, and even mix training data consisting of objects and scenes at different scales. 
+Our method is capable of reconstructing millions of points in a few seconds, and handling very large scenes in an out-of-core fashion. 
+We achieve state-of-the-art results on reconstruction benchmarks consisting of single objects, indoor scenes, and outdoor scenes.*
+
+For business inquiries, please visit our website and submit the
+form: [NVIDIA Research Licensing](https://www.nvidia.com/en-us/research/inquiries/)
+
+
+
+## News
+
+- 2023-06-01: Code released!
+
+## Environment setup
+
+We recommend using the latest Python and PyTorch to run our algorithm. To install all dependencies using [conda](https://www.anaconda.com/):
+
+```bash
+# Clone the repository
+git clone [email protected]:nv-tlabs/nksr
+cd nksr
+
+# Create conda environment
+conda env create
+
+# Activate it
+conda activate nksr
+```
+
+> For docker users, we suggest using a base image from [nvidia/cuda](https://hub.docker.com/r/nvidia/cuda) with tag `12.1.1-cudnn8-devel-ubuntu22.04`, and applying the above conda setup over it.
+
+## Testing NKSR on your own data
+
+We have tested our algorithm on multiple different spatial scales. It can reconstruct scenes spanning kilometers with millions of points+ on an RTX 3090 GPU.
+To use our `kitchen-sink` model (released under CC-BY-SA 4.0 license), use the following code snippet:
+
+```python
+import torch
+import nksr
+```
+
+> To prevent OOM, one last resort is to add `PYTORCH_NO_CUDA_MEMORY_CACHING=1` as environment variable!
+
+## Reproducing results from the paper
+
+Our training and inference system is based on the [Zeus Deep Learning](ZEUS_DL.md) infrastructure, supporting both tensorboard and wandb (recommended) as loggers. To config Zeus, copy the default yaml file and modify the related paths:
+
+```bash
+cp configs/default/zeus.yaml zeus_config.yaml
+```
+
+Modify the contents of `zeus_config.yaml` as needed to include your `wandb` account name and checkpoint/test results save directory.
+
+## Training
+
+NKSR
+
+## Inference
+
+You can either infer using your own trained models or our pre-trained checkpoints.
+
+```bash
+python test.py configs/shapenet/train_3k_noise.yaml --url https://nksr.s3.ap-northeast-1.amazonaws.com/snet-n3k-wnormal.pth --exec udf.enabled=False --test_print_metrics --test_n_upsample 4
+```
+
+## License
+
+Copyright &copy; 2023, NVIDIA Corporation & affiliates. All rights reserved.
+
+This work is made available under the [Nvidia Source Code License](LICENSE.txt).
+
+## Related Works
+
+NKSR is highly based on the following existing works:
+
+- Williams et al. 2021. [Neural Fields as Learnable Kernels for 3D Reconstruction](https://nv-tlabs.github.io/nkf/).
+- Huang et al. 2022. [A Neural Galerkin solver for Accurate Surface Reconstruction](https://github.com/huangjh-pub/neural-galerkin).
+
+## Citation
+
+```bibtex
+@inproceedings{huang2023nksr,
+  title={Neural Kernel Surface Reconstruction},
+  author={Huang, Jiahui and Gojcic, Zan and Atzmon, Matan and Litany, Or and Fidler, Sanja and Williams, Francis},
+  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
+  pages={4369--4379},
+  year={2023}
+}
+```

ZEUS_DL.md

@@ -71,25 +71,29 @@ python upload_ckpt.py <WDB-URL>:<CKPT-ID>
 
 ### WANDB Sweep
 
-To run hyperparameter sweep, first write a config file like `sweeps/*.yaml`. Then run on `ws`:
+To run hyperparameter sweep, first write a config file like `sweeps/*.yaml`. Then run on the local machine (master):
 ```shell
 wandb sweep sweeps/<CONFIG>.yaml
 ```
 
-Then, on NGC, run:
+Then, on the remote cluster, run:
 ```shell
 CUDA_VISIBLE_DEVICES=... wandb agent <SWEEP-ID> --count 1
 ```
 where `count` specifies how many jobs should this agent run.
 Agents will ask for job from the wandb server to run.
 It will call our training script with `--exec` set.
 
+<!---
+
 Alternatively, the following summarizes the `ngman -> jupyter -> jn -> git pull -> run agent` step chain.
 ```shell
 ngman remote sh --type 32g.7 --sweep <SWEEP-ID>
 ```
 where `SWEEP-ID` is something like `huangjh/nkfw-shapenet/j9bbtcm1`.
 
+-->
+
 ## Configuration
 
 Here are some configurations that you can tweak. Please create a file named `zeus_config.yaml` in the same folder as `train.py` to override the ones in `configs/default/zeus.yaml`.

configs/default/param.yaml

@@ -27,3 +27,6 @@ test_use_gt_structure: false
 
 # Test-time transform (for generalization experiments)
 test_transform: null
+
+# Load ckpt from URL.
+url: ''

configs/default/test.yaml

@@ -1,2 +1,3 @@
 include_configs:
   - param.yaml
+  - train.yaml

configs/default/zeus.yaml

@@ -1,8 +1,16 @@
+# Specify this if you want to use wandb
 wandb:
   user: ""
+  # Wandb checkpoint base directory
   base: "./wandb/"
+  # Optional upload path
   upload: "<REMOTE>:<PATH>"
+
+# Specify this if you want to use tensorboard
 tb:
+  # Checkpoint base directory
   base: "./tb/"
   upload: "<REMOTE>:<PATH>"
+
+# Path where inference results will be saved
 test_path: "./test/"

environment.yml

@@ -0,0 +1,45 @@
+name: nksr
+channels:
+  - pyg
+  - nvidia/label/cuda-11.8.0
+  - pytorch
+  - conda-forge
+dependencies:
+  - python=3.10                   # Fix version for reproducibility
+  - pytorch=2.0.0                 # |
+  - pytorch-lightning=1.9.4       # |
+  - libprotobuf=3.19.6    # Protobuf
+  - protobuf=3.19.6       # | 4.x has weird linking bugs...
+  - tensorboard
+  - wandb
+  - pybind11
+  - pip
+  - gitpython
+  - ca-certificates
+  - certifi
+  - openssl
+  - cuda-toolkit
+  - cuda-cudart
+  - cuda-nvcc
+  - cuda-tools
+  - parameterized
+  - gcc_linux-64=11
+  - gxx_linux-64=11
+  - cuda-toolkit
+  - setuptools
+  - cmake
+  - ninja
+  - ipython
+  - matplotlib
+  - tqdm
+  - pyg
+  - rich
+  - pandas
+  - pytorch-scatter
+  - omegaconf
+  - flatten-dict
+  - pip:
+    - python-pycg
+    - randomname
+    - open3d
+    - pykdtree

models/base_model.py

@@ -28,7 +28,6 @@
 from pycg.exp import AverageMeter
 from torch.optim.lr_scheduler import LambdaLR
 from torch.utils.data import DataLoader
-from torch.utils.tensorboard.summary import hparams
 
 
 exp.global_var_manager.register_variable('skip_backward', False)
@@ -227,6 +226,8 @@ def on_fit_start(self):
         if self.trainer.logger is None:
             return
         if self.logger_type == 'tb':
+            from torch.utils.tensorboard.summary import hparams
+
             writer = self.trainer.logger.experiment._get_file_writer()
             if writer is not None:
                 hparams_metrics = {}

models/nksr_net.py

@@ -6,6 +6,8 @@
 import numpy as np
 from nksr import NKSRNetwork, SparseFeatureHierarchy
 from nksr.fields import KernelField, NeuralField, LayerField
+from nksr.configs import load_checkpoint_from_url
+
 from pycg import exp, vis
 
 from dataset.base import DatasetSpec as DS, list_collate
@@ -23,6 +25,9 @@ class Model(BaseModel):
     def __init__(self, hparams):
         super().__init__(hparams)
         self.network = NKSRNetwork(self.hparams)
+        if self.hparams.url:
+            ckpt_data = load_checkpoint_from_url(self.hparams.url)
+            self.network.load_state_dict(ckpt_data['state_dict'])
 
     @exp.mem_profile(every=1)
     def forward(self, batch, out: dict):

test.py

@@ -37,14 +37,16 @@ def get_default_parser():
 
     known_args = parser.parse_known_args()[0]
     args_ckpt = known_args.ckpt
-    if args_ckpt.startswith("wdb:"):
-        wdb_run, args_ckpt = wdb.get_wandb_run(args_ckpt, wdb_base=zeus.config.wandb.base, default_ckpt="test_auto")
-        assert args_ckpt is not None, "Please specify checkpoint version!"
-        assert args_ckpt.exists(), "Selected checkpoint does not exist!"
-        model_args = omegaconf.OmegaConf.create(wdb.recover_from_wandb_config(wdb_run.config))
-    elif args_ckpt is not None:
-        model_yaml_path = Path(known_args.ckpt).parent.parent / "hparams.yaml"
-        model_args = exp.parse_config_yaml(model_yaml_path)
+
+    if args_ckpt is not None:
+        if args_ckpt.startswith("wdb:"):
+            wdb_run, args_ckpt = wdb.get_wandb_run(args_ckpt, wdb_base=zeus.config.wandb.base, default_ckpt="test_auto")
+            assert args_ckpt is not None, "Please specify checkpoint version!"
+            assert args_ckpt.exists(), "Selected checkpoint does not exist!"
+            model_args = omegaconf.OmegaConf.create(wdb.recover_from_wandb_config(wdb_run.config))
+        else:
+            model_yaml_path = Path(known_args.ckpt).parent.parent / "hparams.yaml"
+            model_args = exp.parse_config_yaml(model_yaml_path)
     else:
         model_args = None
     args = parser.parse_args(additional_args=model_args)