Topology-Preserving Shape Reconstruction and Registration via Neural Diffeomorphic Flow [CVPR 2022]

Shanlin Sun, Kun Han, Deying Kong, Hao Tang, Xiangyi Yan, Xiaohui Xie.

Full paper is available here.

Data Preprocess

• given the center-cropped mask volumes, generate organ shapes via marching cube
• apply minor laplacian smoothing
• sample points with sdf values as DeepSDF suggests

The preprocesed data is here, including four datasets used in the paper -- Pancreas CT, MMWHS, Liver and Lung. Download them to your local path and this is will be the "Data Source" in the following steps.

Installation

conda create --name ndf --file requirements.txt
source activate ndf

The following steps should be executed in order as is shown in "scripts/pancreas_experiments.sh"

Train

CUDA_VISIBLE_DEVICES=0 python train_ndf.py -e ${experiment_name} --debug --batch_split 2 -c latest -d ${data_source}

Here, "experiment_name" should be the log path of your experiment and "data_source" indicates the directory path where you save all preprocessed data.

Template generation

CUDA_VISIBLE_DEVICES=0 python generate_template_mesh.py -e ${experiment_name} --debug 

Optimize for Unseen Shape Objects

CUDA_VISIBLE_DEVICES=0 python reconstruct_ndf.py -e ${experiment_name} -c latest -s ${data_split} -d ${data_source} --iters 2400 --octree --skip --debug

Here, "data split" is a json file containing the test cases. You can find some split samples under the folder examples/splits.

Shape Reconstruction

For seen cases

CUDA_VISIBLE_DEVICES=0 python generate_training_meshes.py -e ${experiment_name} --debug --start_id 0 --end_id ${number_of_samples-1} --octree --keep_normalization

For unseen cases

Shape reconstruction is embedded in "Optimize for unseen shape objects"

Shape Registration

label the colors of vertices in the reconstructed shape mesh

# seen cases
CUDA_VISIBLE_DEVICES=0 python generate_meshes_correspondence.py -e ${experiment_name} -c ${num_clusters} --debug --start_id 0 --end_id ${number_of_samples-1} --test_or_train train
# unseen cases
CUDA_VISIBLE_DEVICES=0 python generate_meshes_correspondence.py -e ${experiment_name} -c ${num_clusters} --debug --start_id 0 --end_id ${number_of_samples-1} --test_or_train test

Dense point correspondence across each shape object and the corresponding learned template

# seen cases
CUDA_VISIBLE_DEVICES=${GPU_ID} python generate_meshes_topology_correspondence.py -e ${experiment_name} --debug --start_id 0 --end_id ${number_of_samples-1} --num_clusters ${num_clusters} --bp_or_ode ode --test_or_train train
# unseen cases
CUDA_VISIBLE_DEVICES=${GPU_ID} python generate_meshes_topology_correspondence.py -e ${experiment_name} --debug --start_id 0 --end_id ${number_of_samples-1} --num_clusters ${num_clusters} --bp_or_ode ode --test_or_train test

Evaluations

# seen cases
CUDA_VISIBLE_DEVICES=${GPU_ID} python evaluate.py -e ${experiment_name} -c 2000 -s ${train_data_split} -d ${data_source} -t train -l fine -n 2500 --debug
# unseen cases
CUDA_VISIBLE_DEVICES=${GPU_ID} python evaluate.py -e ${experiment_name} -c 2000 -s ${test_data_split} -d ${data_source} -t test -l fine -n 2500 --debug

Acknowledge

This code repo is heavily based on the DeepSDF and DeepImplicitTemplates. Our evaluatin codes are developed based on NeuralMeshFlow. We thank the authors for their great job!

Citation

@InProceedings{Sun_2022_CVPR,
    author    = {Sun, Shanlin and Han, Kun and Kong, Deying and Tang, Hao and Yan, Xiangyi and Xie, Xiaohui},
    title     = {Topology-Preserving Shape Reconstruction and Registration via Neural Diffeomorphic Flow},
    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    month     = {June},
    year      = {2022},
    pages     = {20845-20855}
}

License

MIT License