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

Siwensun/Neural_Diffeomorphic_Flow--NDF