a

README.md

@@ -1,5 +1,215 @@
-# NAISR
-NAISR: A 3D Neural Additive Model for Interpretable Shape Representation
+# <img src="NAISRicon.svg" width="35" height="35"> NAISR: A 3D Neural Additive Model for Interpretable Shape Representation
 
+Pytorch implementation for our `NAISR` paper <br>
+[NAISR: A 3D Neural Additive Model for Interpretable Shape Representation](https://arxiv.org/abs/2303.09234), ICLR 2024 **Spotlight**. <br>
+Yining Jiao, Carlton Zdanski, Julia Kimbell, Andrew Prince, Cameron Worden, Samuel Kirse, Christopher Rutter, Benjamin Shields, William Dunn, Jisan Mahmud, Marc Niethammer. <br>
+UNC-Chapel Hill
 
-Codes on the way...
+#### [Paper](https://arxiv.org/abs/2303.09234) | [Project Page](https://uncbiag.github.io/NAISR/) | [Colab Demos](https://colab.research.google.com/drive/1OudGynEydIXpAgfA9lvi5d0L86kK-uwg?usp=sharing)
+
+
+<div align=center><img src="https://github.com/uncbiag/NAISR/blob/main/figs/NAISRoverviewwide.svg?raw=true" width="1000" height="600"/></div>
+
+Please cite as:
+```
+@inproceedings{
+jiao2024naisr,
+title={\texttt{NAISR}: A 3D Neural Additive Model for Interpretable Shape Representation},
+author={Yining Jiao and Carlton Jude ZDANSKI and Julia S Kimbell and Andrew Prince and Cameron P Worden and Samuel Kirse and Christopher Rutter and Benjamin Shields and William Alexander Dunn and Jisan Mahmud and Marc Niethammer},
+booktitle={The Twelfth International Conference on Learning Representations},
+year={2024},
+url={https://openreview.net/forum?id=wg8NPfeMF9}
+}
+```
+
+>[!NOTE]
+>For the three datasets used in `NAISR` paper, we provided this [Colab demo](https://colab.research.google.com/drive/1OudGynEydIXpAgfA9lvi5d0L86kK-uwg) to investigate the shapes learned with `NAISR`. <br>
+We also provide the instructions on this page to apply `NAISR` to your customized shape analysis questions. 
+
+
+## Installation
+The code is tested with ``python=3.9``, ``torch=2.1.0``, ``torchvision=0.15.2``.
+```
+git clone https://github.com/uncbiag/NAISR
+cd NAISR
+```
+Now, create a new conda environment and install required packages accordingly.
+```
+conda create -n naisr python=3.9
+conda activate naisr
+pip install -r requirements.txt
+```
+
+
+## Data and Model Weights
+We train and test our method on three datasets: Starman, ADNI Hippocampus, and Pediatric Airways.
+
+| Dataset   |                      Description             |           Dataset Link              |   Model Link              |
+|-----------|----------------------------------------------|:------------------------------------:|:------------------------------------:|
+|Starman      |  simulated 2D Starman shapes      |  [simulation code)][StarmanData]          | [weights][StarmanModel]    | 
+|ADNI Hippocampus  |  1632 hippocampus from ADNI     |  [official site][ADNIData]              | [weights][ADNIModel]  |
+|Pedeatric Airway |  357 pediatric airway shapes      |  NA | [weights][AirwayModel]               |
+
+
+
+[ADNIData]: https://ida.loni.usc.edu/login.jsp?project=ADNI
+[StarmanData]: 
+
+[ADNIModel]: https://github.com/uncbiag/NAISR/releases/download/naisr_weights_v0/naisr_weights_adni.pth
+[StarmanModel]: https://github.com/uncbiag/NAISR/releases/download/naisr_weights_v0/naisr_weights_starman.pth
+[AirwayModel]: https://github.com/uncbiag/NAISR/releases/download/naisr_weights_v0/naisr_weights_pediatric_airway.pth
+
+
+## Visualizations of Shape Space Extrapolation 
+
+### Getting Started
+One just need to use this [Colab demo](https://colab.research.google.com/drive/1OudGynEydIXpAgfA9lvi5d0L86kK-uwg) to explore the template shape space of {starmen, hippocampi and airways} used in the `NAISR` paper. 
+
+### More functions
+Our code repo provides more functions to visualize the shape space, e.g., as a matrix for a specific case or for the template shape (as Figure 3 in the main paper),
+
+To get the matrix of the the matrix of thetemplate shape extrapolation, pls use
+```
+python evolution_shapematrix.py -e examples/hippocampus/naigsr_0920_base.json
+```
+
+To get the matrix of the shape extrapolation for a specific patient, pls use
+```
+python evolution_shapematrix_specific.py -e examples/hippocampus/naigsr_0920_base.json
+```
+
+
+
+## Customize
+One may also want to use `NAISR` on their own shape analysis problems. For this use we provide our best suggestions/instructions here through the illustration for the simulated starman dataset,
+
+### Data Preprocessing
+
+#### Alignment with Rigid Transformation
+The shapes to explore need to be registered with a rigid transformation (translation + rotation). 
+If paired point clouds are available, we recommend to use ; otherwise, we recommend to use ICP to register the point clouds. 
+In our case, we use airway landmarks to learn the rigid transformation; and ICP algorithms to register the point clouds of the hippocampi.
+
+#### SDF Extraction
+
+###
+```json
+ {
+  "Description" : [ "This experiment learns a shape representation for starman dataset." ],
+  "Device":0, 
+  "DataSource": {"train": "/home/jyn/NAISR/examples/starman/2dshape_train_with_temp.csv",
+                 "test": "/home/jyn/NAISR/examples/starman/2dshape_test_with_temp.csv"},
+  "Split": null,
+  "Network": "DeepNAIGSR",
+  "NumEpochs": 300,
+  "LoggingRoot": "/playpen-raid/jyn/NAISR/log",
+  "ExperimentName": "DeepNAIGSR_STARMAN3D_0222_256_base",
+
+  "EpochsTilCkpt": 10,
+  "StepsTilSummary": 1000,
+  "UseLBFGS": false,
+  "DoublePrecision": false,
+  "CheckpointPath": "",
+  "CodeLength": 256,
+
+  "AdditionalSnapshots" : [ 50, 100, 200, 300, 400, 500 ],
+  "LearningRateSchedule" : [
+    {
+      "Type": "Step",
+      "Initial": 0.00005,
+      "Interval": 1000,
+      "Factor": 0.5
+    },
+    {
+      "Type": "Step",
+      "Initial":  0.001,
+      "Interval": 1000,
+      "Factor": 0.5
+    }],
+  "SamplesPerScene" : 750,
+  "BatchSize": 64,
+  "DataLoaderThreads": 4,
+  "ClampingDistance": 1,
+
+  "Articulation": true,
+  "NumAtcParts": 1,
+  "TrainWithParts": false,
+  "Class": "starman",
+  "Attributes": ["cov_1", "cov_2"],
+   "TemplateAttributes":  {"cov_1": 0, "cov_2": 0},
+  "Backbone": "siren",
+  "PosEnc": false,
+  "InFeatures": 2,
+  "HiddenFeatures": 256,
+  "HidenLayers": 6,
+  "OutFeatures": 1,
+  "Loss": {
+    "whether_sdf": true, 
+    "whether_normal_constraint": true, 
+    "whether_inter_constraint": true, 
+    "whether_eikonal": true,
+    "whether_code_regularization": true}
+```
+### Training
+For training, one needs to use `train_atlas_3dnaigsr.py` with different networking settings indicated with the `json` files, e.g., for the starman dataset,
+```
+python train_atlas_3dnaigsr.py -e examples/starman/naigsr_0222_base.json
+```
+
+### Shape Reconstruction
+For testing/reconstruction without covariates, one needs to `reconstruct_atlas.py` with different network settings indicated with the `json` files, e.g., for the starman dataset.
+```
+python reconstruct_atlas.py -e examples/starman/naigsr_0222_base.json 
+```
+
+For testing/reconstruction with covariates, one needs to `reconstruct_atlas.py` with different network settings indicated with the `json` files, e.g., for the starman dataset.
+```
+python reconstruct_atlas_with_cov.py -e examples/starman/naigsr_0222_base.json 
+```
+
+
+### Shape Transport
+For shape transport without covariates, one needs to `reconstruct_atlas.py` with different network settings indicated with the `json` files, e.g., for the starman dataset.
+```
+python transport_general.py -e examples/starman/naigsr_0222_base.json 
+```
+
+For shape transport with covariates, one needs to `reconstruct_atlas.py` with different network settings indicated with the `json` files, e.g., for the starman dataset.
+```
+python transport.py -e examples/starman/naigsr_0222_base.json 
+```
+
+
+### Shape Evolution and Disentanglement
+
+One just needs to adjust this [Colab demo](https://colab.research.google.com/drive/1OudGynEydIXpAgfA9lvi5d0L86kK-uwg) with their own `NAISR` weights to explore the learned deform the template shape with the query covariates. 
+
+To get the matrix of the template shape extrapolation, pls use
+```
+python evolution_shapematrix.py -e examples/starman/naigsr_0222_base.json 
+```
+
+To get the matrix of the shape extrapolation for a specific patient, pls use
+```
+python evolution_shapematrix_specific.py -e examples/starman/naigsr_0222_base.json 
+```
+
+
+
+More instructions on the way...
+
+
+
+
+
+## Cite this work
+```
+@inproceedings{
+jiao2024naisr,
+title={\texttt{NAISR}: A 3D Neural Additive Model for Interpretable Shape Representation},
+author={Yining Jiao and Carlton Jude ZDANSKI and Julia S Kimbell and Andrew Prince and Cameron P Worden and Samuel Kirse and Christopher Rutter and Benjamin Shields and William Alexander Dunn and Jisan Mahmud and Marc Niethammer},
+booktitle={The Twelfth International Conference on Learning Representations},
+year={2024},
+url={https://openreview.net/forum?id=wg8NPfeMF9}
+}
+```

docs/index.html

@@ -104,6 +104,16 @@ <h2 class="title is-2 publication-title"><tt>NAISR</tt>: A 3D Neural Additive Mo
                   </a>
               </span>
               <!-- video Link. -->
+              <span class="link-block">
+                <a href="https://colab.research.google.com/drive/1OudGynEydIXpAgfA9lvi5d0L86kK-uwg?usp=sharing"
+                   class="external-link button is-normal is-rounded is-dark">
+                  <span class="icon">
+                      <i class="fab fa-google"></i>
+                  </span>
+                  <span>Video</span>
+                  </a>
+              </span>
+              <!-- video Link. -->
               <span class="link-block">
                 <a href="https://www.youtube.com/watch?v=8wSIIPAvlLI&t=43s"
                    class="external-link button is-normal is-rounded is-dark">

examples/starman/naigsr_0222_base.json

@@ -1,7 +1,8 @@
  {
   "Description" : [ "This experiment learns a shape representation for starman dataset." ],
   "Device":0,
-  "DataSource": {"train": "/home/jyn/NAISR/examples/starman/2dshape_train_with_temp.csv", "test": "/home/jyn/NAISR/examples/starman/2dshape_test_with_temp.csv"},
+  "DataSource": {"train": "examples/starman/2dshape_train_with_temp.csv",
+    "test": "examples/starman/2dshape_test_with_temp.csv"},
   "Split": null,
   "Network" : "DeepNAIGSR",
   "NumEpochs" : 300,

naisr/__init__.py

@@ -10,4 +10,7 @@
 #from naisr.model import  NAISiren, NAIVF, BaselineVF, FCBaseline, NAIVF_withtempl, SirenlatentVF, NAIlatentVF_withtempl, LipNAIVF_withtempl, Baseline, BaselineVF, ICVF, NAIVF_with3dtempl, NAIVF_autotempl, NAIVF_fixedtempl, NAIVF_fixed, DeepSDF #NAISR, HyperSirenBaseline, SirenBaseline,
 from naisr.diff_operators import *
 from naisr.metrics import *
-from naisr.model_naigsr import DeepNAIGSR
+from naisr.model_naigsr import DeepNAIGSR
+from naisr.starman_dataset import *
+from naisr.adni_dataset import *
+from naisr.airway_dataset import *

naisr_meshing.py

@@ -1148,10 +1148,9 @@ def convert_3d_sdf_samples_to_ply(
     print(numpy_3d_sdf_tensor.max())
     print(numpy_3d_sdf_tensor.min())
     try:
-        verts, faces, normals, values = measure.marching_cubes_lewiner(
+        verts, faces, normals, values = measure.marching_cubes(
                 numpy_3d_sdf_tensor, level=0., spacing=[voxel_size] * 3
             )
-
     except:
         path_aligned_surface = os.path.join(savedir, 'surface.stl')
         ply_filename_out = os.path.join(savedir, 'surface.ply')

publicdata/__init__.py

@@ -0,0 +1,2 @@
+from publicdata.utils_2d import *
+from publicdata.deformation import *