Authors: Sven Dummer, Nicola Strisciuglio, Christoph Brune.
Paper: The full paper is available here or here.
NOTE: this repository has only been tested on Linux / Ubuntu! When using Windows, we recommend using Windows Subsystem for Linux (WSL).
To install and activate the conda environment, execute:
conda env create -f environment.yml
conda activate RDA_INR
In the paper, two datasets are used: a liver dataset and a rectangles dataset.
We use the preprocessed liver data from Sun et al. [2022], whose GitHub repository can be found here. One can download the liver dataset from here. Put the data in the preparing_datasets/liver/raw_data/liver folder.
After saving the liver dataset, we process the data and transfer the processed data to the data folder:
python preparing_datasets/liver/generate_liver_data.py
python preparing_datasets/liver/transfer_processed_liver_data.py
The rectangles dataset contains randomly rotated rectangles centered around the origin with edge lengths uniformly distributed in
- One can download the rectangles dataset from here. This is the dataset that is used in the paper for training the models. We recommend saving the data in the data/Rectangles folder.
- One can generate their own rectangles data.
To generate your own rectangles dataset, execute the following code:
python making_shapes/create_rectangles_dataset.py --N_tr N_tr --N_te N_te
where N_tr denotes the number of samples in the training dataset and N_te the number of samples in the test dataset. In the paper, we use N_tr=100 and N_te=20.
NOTE: the generated data is saved in the data/Rectangles folder and does not have to be transferred to another folder!
To generate the required data from mesh datasets in .ply format, run:
python preprocess_ply_dataset.py dataset_dir
Here, "dataset_dir" denotes the directory containing all the .ply files. Particularly, the directory should include a subfolder dataset_dir/MeshSamples that contains all the .ply meshes.
After running the code, two new folders are generated:
- dataset_dir/SdfSamples: contains signed distance functions (SDFs) for each mesh / shape.
- dataset_dir/SurfaceSamples: contains pointclouds representing each mesh / shape.
For training the main model in our paper, which is the model that uses pointclouds, one has to use the train_rda_inr_pointcloud.py function. For instance, one can execute:
CUDA_VISIBLE_DEVICES=0,1 python train_rda_inr_pointcloud.py -e ${experiment_name} --batch_split n_s -c latest -d ${data_source}
For training the model that uses occupancy data, one has to use the train_rda_inr_levelset.py function. For instance, one can execute:
CUDA_VISIBLE_DEVICES=0,1 python train_rda_inr_levelset.py -e ${experiment_name} --batch_split n_s -c latest -d ${data_source}
In the above, we have:
- We use the GPU's 0 and 1 (you can change this if you want).
- "n_s" is the number of times we split the batch. E.g. if we have a batch size of 10 and n_s=2, we first process the first 5 shapes of the batch and then the second 5 shapes of the batch. The gradient of the loss function is still calculated with respect to the 10 shapes (and not with respect to 5 shapes).
- "experiment_name" should be the log path of your experiment. Here you put your specs.json with the specific training settings. For instance, you can use "experiments/liver/riem reg". This folder also contains an example specs.json file.
- "data_source" is the directory where all the preprocessed data is saved. For instance, it is "data/Rectangles" or "data/liver"
For an example of how to train a model when using a cluster with SLURM: see the scripts/train_model.sbatch file.
NOTE: the training process is memory intensive due to the calculation of the Riemannian flow regularization. The "n_s" parameter is useful in case such memory problems occur.
Execute the following code to evaluate the model:
bash scripts/Evaluate_model.sh -e $experiment_directory_input -s $test_split_input -o -n $noisy_split -I -R
This script is used to:
- Calculate the learned template shape.
- Get the reconstructions of the training data, test data, and noisy data.
- Evaluate the reconstruction metrics for the training data, test data, and noisy data.
- Save the deformations of the template to the reconstructed training shapes.
- Generate correspondences for the training shapes and test shapes.
Explanation of the variables in the above code block:
- -e: the directory of the experiment with respect to the main folder (the main folder is the folder containing the repository). E.g. use something like experiments/liver/riem reg.
- -s: the location of the (test) split .json file (with respect to the main folder, which contains the repository). E.g. use something like splits/liver_test.
- -n: the location of the split file that determines to which meshes we add noise for evaluating noisy mesh reconstruction. E.g. use something like splits/liver_test.
For an example of how to evaluate a model when using a cluster with SLURM: see the scripts/evaluate_model.sbatch file.
If you have any questions, please contact Sven Dummer (s.c.dummer@utwente.nl).
This repository is mainly based on Neural-Diffeomorphic-Flow. Furthermore, it contains aspects of DIF-NET. We would like to thank the authors for their excellent work!
@article{dummer2024rda,
title={RDA-INR: Riemannian diffeomorphic autoencoding via implicit neural representations},
author={Dummer, Sven and Strisciuglio, Nicola and Brune, Christoph},
journal={SIAM Journal on Imaging Sciences},
volume={17},
number={4},
pages={2302--2330},
year={2024},
publisher={SIAM}
}
MIT License
The raw data in the preparing_datasets/liver/raw_data/liver folder, should have the following folder structure:
preparing_datasets
│
└─── liver
│
└─── raw_data
│
└─── liver
│
└─── MeshSamples
│ │
│ └─── INHOUSE_LIVER
│ │
│ └─── liver
│ │
│ └─── shape_name_*.ply
│
└─── SdfSamples
│ │
│ └─── INHOUSE_LIVER
│ │
│ └─── liver
│ │
│ └─── shape_name_*.npz
│
└─── SurfaceSamples
│
└─── INHOUSE_LIVER
│
└─── liver
│
└─── shape_name_*.ply
where the * indicates there might be multiple shapes.
The preprocessing of this data using python preparing_datasets/liver/generate_liver_data.py
makes sure that the meshes are within the unit cube
- samples 100 000 points and corresponding normals uniformly from the mesh surface.
- scales down these uniformly sampled points and the mesh vertices (and hence the mesh) by multiplying each location coordinate by 0.75.
We save the 100 000 point-normal pairs as .npz files. These .npz files are saved in the preparing_datasets/liver/raw_data/liver/SurfaceSamples folder and replace the .ply files originally in that folder. The rescaled meshes are saved as .ply files in the preparing_datasets/liver/raw_data/liver/MeshSamples_rescaled folder.
Finally, the rescaled meshes and the sampled pointclouds need to be transferred to the data folder, which is done via python preparing_datasets/liver/transfer_processed_liver_data.py
NOTE: the above functions only preprocess the meshes and do not preprocess the SDF data that is also present in the downloaded data. Hence, the preprocessed meshes do NOT correspond to the SDF data anymore! Hence, one can not immediately apply the Occupancy Functions code to the liver dataset. This is not a problem for reproducing our results as in our paper we only apply the method that uses pointclouds when dealing with the liver dataset!
For a new dataset with name datasource_name, make sure the data/datasource_name folder follows a similar structure to DeepSDF:
RDA-INR
│
└─── data
│
└─── datasource_name
│ │
│ └─── MeshSamples
│ │ │
│ │ └─── dataset_name_*
│ │ │
│ │ └─── class_name_*
│ │ │
│ │ └─── shape_name_*.ply
│ │
│ └─── SdfSamples
│ │ │
│ │ └─── dataset_name_*
│ │ │
│ │ └─── class_name_*
│ │ │
│ │ └─── shape_name_*.npz
│ │
│ └─── SurfaceSamples
│ │
│ └─── dataset_name_*
│ │
│ └─── class_name_*
│ │
│ └─── shape_name_*.npz
...
where dataset_name_* is one of the sub-datasets of the datasource, class_name_* one of the classes in this sub-dataset (e.g. "chair", "couch", etc.), and shape_name_* is the name of a specific shape. Note that the * indicates that there might be multiple shapes, multiple classes, and multiple sub-datasets!
NOTE: Make sure the shapes are situated within the unit cube
To perform an experiment, one has to specify the experiment options in a specs.json file. You should put this file inside the directory that you assigned for the experiment. Examples of such specs.json files are given in the subfolders of the experiments folder. The train_rda_inr_pointcloud.py file requires more experiment options than the train_rda_inr_levelset.py file. Some of the options that overlap are:
- template: the implicit function neural network used for the template. Current recommended options: DIF-Net_no_final_tanh (for the pointcloud data) and DIF-Net (for the implicit function data). There are two more options, but we do not discuss them here as these are not used in the paper.
- warper: the neural ODE type used for deforming the template. Current options: Warper (used for generating the results in the paper) and WarperOT. The architecture of WarperOT is the same as Warper but it calculates the velocity field regularization differently / more efficiently. More precisely, it uses the approach in MFG-Net to calculate the Riemannian integral regularization.
-
NetworkSpecs:
-
warper_kargs (NOTE: when you construct a whole new warper, you have to supply the inputs required for THAT warper!):
-
steps: the
$K$ variable in the paper. -
killing_l2_balance_lamb: the
$\eta$ in the paper. - odefunc: the velocity vector field parameterizations. Current options: ODEFunc and ODEFuncLarger. In the paper, we use ODEFuncLarger. The ODEFuncLarger velocity vector field parameterization is an extension of ODEFunc, which in turn is used in NDF.
-
cut_off: whether to use the cut_off function
$h_\epsilon$ as discussed in Appendix D.1 / Supplementary Materials SM4.1 of the paper.
-
steps: the
- decoder_kargs: options used for the template implicit neural representation. Consult the code to see which options are available.
-
warper_kargs (NOTE: when you construct a whole new warper, you have to supply the inputs required for THAT warper!):
- ScenesPerBatch: the number of shapes in a training batch.
- SamplesPerScene: the number of implicit function values sampled per shape in the batch when training using implicit functions OR the number of points sampled from each pointcloud in the batch when training with pointclouds.
-
ClampingDistance: when not using occupancy data, the implicit function value produced by the decoder is clamped to be in
$[-d, d]$ (with"ClampingDistance": d). -
CodeBound:
"CodeBound: c"enforces the hard constraint$\lVert z \rVert_\infty \leq c$ for the learned latent codes of the training shapes.
Two additional (required) options for train_rda_inr_pointcloud.py are:
-
regularize_eikonal_warped_points: we uniformly sample points from the domain and use the eikonal loss at those locations to regularize the template INR. We use
"regularize_eikonal_warped_points": trueif we ALSO want to employ the eikonal penalty at warped pointcloud points/locations. We use"regularize_eikonal_warped_points": truein the experiments in the paper. -
far_from_zero_loss_after_warp: whether to evaluate the
$\beta$ loss term at uniformly sampled locations OR at locations after they have been warped with the chosen warper. We use"far_from_zero_loss_after_warp": truein the experiments in the paper.
Finally, we have some options related to regularization constants / losses. For both train_rda_inr_pointcloud.py and train_rda_inr_levelset.py, we have:
-
CodeRegularizationLambda: the regularization coefficient for the
$l_2$ regularization of the latent vectors. NOTE: we only have this loss component if we have"CodeRegularization": true. -
PointWiseLossWeight: this is the
$c_{PW}$ regularization constant in the paper (see Table D.1 / SM4.1 in the paper for its definition). The pointwise loss is only used if UsePointwiseLoss is set to true in the specs.json file. - PointpairLossWeight: NOT USED IN THE PAPER! It is a loss used in DIT-NET and can be activated by putting UsePointpairLoss to true in the specs.json file.
-
UseCurriculum: the name of this variable stems from a curriculum learning strategy employed in Neural-Diffeomorphic-Flow (NDF). In our code/repo, it determines whether the pointwise and pointpair losses are only calculated at the end points of the template deformation OR also after some other velocity vector field blocks. More precisely, we have
"UseCurriculum": truein the specs.json file if we use the pointwise and/or pointpair loss to regularize the deformed template at times$t \in T={\lfloor K/4 \rfloor \cdot i \mid 1 \leq i \leq K / \lfloor 0.25 K \rfloor, i \in \mathbb{N} }$ with$K$ the number of stationary velocity vector fields in our quasi-time varying velocity field parameterization. We have"UseCurriculum": falseif we only regularize the deformed final shape at time$t=t_K=1$ . Note that the UseCurriculum option only matters when training the model using the pointwise loss and/or the pairwise loss. In the experiments where we used the pointwise loss, we used"UseCurriculum": truein the specs.json file.
Additionally, for train_rda_inr_pointcloud.py we have:
-
eikonal_reg_const: this corresponds to the
$\lambda$ parameter in the paper. -
far_from_zero_const: this corresponds to the
$\beta$ parameter in the paper. -
vel_reg_const: this corresponds to the
$\sigma^2$ parameter in the paper. -
cos_sim_const: this corresponds to the
$\tau$ parameter in the paper.
Explanation of the obligatory inputs of the script:
- -e: the directory of the experiment with respect to the main folder (the main folder is the folder containing the repository). E.g. use something like experiments/liver/riem reg.
- -s: the location of the (test) split .json file (with respect to the main folder, which contains the repository). E.g. use something like splits/liver_test.
- -n: the location of the split file that determines to which meshes we add noise for evaluating noisy mesh reconstruction. E.g. use something like splits/liver_test.
Explanation of the optional inputs of the script:
-
-d: the location of the dataset we are interested in to evaluate against. If it is not specified, we are going to use the data source that is specified in the specs.json file. This option can be used to evaluate a model trained with downscaled meshes. E.g. we can reconstruct data using the downscaled meshes and then compare against the original dataset (specified via the -d option).
-
-o: if this option is specified, we are going to use the octree option for generating meshes. If it is not specified, we are not using the octree option.
-
-c: a string "0 1 2 ..." that contains the devices that are visible for CUDA. These will be put in CUDA_VISIBLE_DEVICES.
-
-r: the resolution that is used by marching cubes to generate meshes of:
- the shape reconstructions of the training data,
- the shape reconstructions of the data in the split file indicated in the -s option,
- the shape reconstructions of the generated noisy data.
The default resolution value is 256.
-
-I: this option indicates if we use the pointcloud model or not. If we specify it, we use the pointcloud model, else we do not. In case we do not use the pointcloud model, the evaluation function only:
- returns the estimated template
- returns the training reconstructions
- returns the reconstruction metrics for the training data
- returns the deformations between the template and the training data
- colors the training meshes and template to show the correspondences.
-
-R: this option indicates if we test for robustness to noise or not. If we specify it, we test for robustness.
Below we provide an extensive overview of what is contained in each experiment folder after training and evaluation (of the pointcloud model). Before training, only a specs.json file has to be present in the experiment folder.
experiment_folder
│
└─── LatentCodes
│ │
│ └─── <Epoch>.pth AND/OR latest.pth
│
└─── ModelParameters
│ │
│ └─── <Epoch>.pth AND/OR latest.pth
│
└─── OptimizerParameters
│ │
│ └─── <Epoch>.pth AND/OR latest.pth
│
└─── wandb
│
└─── TrainingMeshes
│ │
│ └─── <Epoch> (only those epochs at which you evaluated the model)
│ │
│ └─── template.ply
│ │
│ └─── template_color_coded.ply
│ │
│ └─── dataset_name_*
│ │
│ └─── class_name_*
│ │
│ └─── shape_name_*.ply
│ │
│ └─── shape_name_*_color_coded.ply
│ │
│ └─── shape_name_*_interpolation_meshes
│ │
│ └─── mesh_i.ply
│
└─── DefaultReconstruction
│ │
│ └─── <Epoch> (only those epochs at which you evaluated the model)
│ │
│ └─── Codes
│ │ │
│ │ └─── dataset_name_*
│ │ │
│ │ └─── class_name_*
│ │ │
│ │ └─── shape_name_*.pth
│ │
│ └─── Meshes (same structure as the TrainingMeshes/<Epoch> folder without the shape_name_*_interpolation_meshes folders)
│
└─── NoisyReconstructions_0.0050
│ │
│ └─── Same structure as DefaultReconstruction (but with the with the shape_name_*_interpolation_meshes folders as
│ subfolders in the Meshes folder. The shape_name_*_interpolation_meshes folders contain the colored deformation
│ from the template to the reconstructed (noisy) shapes. Moreover, the NoisyReconstructions_0.0050 folder has a
│ subfolder called gtNoisyMeshes that contains the noisy variants of the used ground truth meshes)
│
└─── NoisyReconstructions_0.0100
│ │
│ └─── Same structure as NoisyReconstruction_0.0050
│
└─── Evaluation
│ │
│ └─── <Epoch>
│ │
│ └─── .csv / .xlsx files containing reconstruction metric statistics
│
└─── specs.json
│
└─── Logs.pth
│
└─── code_bk_*.tar.gz
The folders contain the following:
- LatentCodes, ModelParameters, and OptimizerParameters: these folders contain .pth files that save the latent codes of the training shapes, the model parameters used, and the optimizer parameters used FOR the model obtained at a certain epoch.
- wandb: a folder containing information about the weights and biases session used during training.
- TrainingMeshes: contains the template mesh and the reconstructed training meshes. Moreover, it contains the deformations of the template shape to the reconstructed training shapes.
- DefaultReconstructions: contains the found latent codes for the test shapes AND it contains the reconstructed test meshes.
- NoisyReconstructions_0.0050 and NoisyReconstructions_0.0100: first of all, they have the gtNoisyMeshes folder as subfolder. This subfolder contains the ground truth noisy shapes. For the rest, the NoisyReconstructions_* folders contain the same as DefaultReconstructions, but instead for the noisy meshes in the gtNoisyMeshes subfolder. In one case, we added Gaussian noise with 0.0050 standard deviation and in the other case noise with 0.0100 standard deviation.
- Evaluation: contains .csv / .xlsx files containing reconstruction metric statistics when reconstructing the training data, test data, and the noisy data.
- specs.json determines the specific training parameters. This is the only file that must be present BEFORE training!
- Logs.pth is a logging of the training.
- code_bk_.tar.gz* contains the used code.
NOTE: the .ply files with '_color_coded' in their name are colored meshes where the colors show the correspondence/matching to the template shape (and hence to the other shapes).