This repository contains the PyTorch implementation of the paper:
PialNN: A Fast Deep Learning Framework for Cortical Pial Surface Reconstruction [Paper]
Qiang Ma, Emma C. Robinson, Bernhard Kainz, Daniel Rueckert, Amir Alansary
The 4th Workshop on Machine Learning in Clinical Neuroimaging (MLCN2021)
Abstract. Pial Neural Network (PialNN) is a 3D deep learning framework for pial surface reconstruction. It is trained end-to-end to deform a given initial white matter surface to a target pial surface by a sequence of learned deformation blocks. PialNN is fast and memory-efficient, which allows reconstructing a pial surface mesh with 150k vertices within 1s runtime.
The following packages are required for running PialNN:
- Python >= 3.6
- PyTorch >= 1.5.0
- PyTorch3D = 0.4.0
- PyTorch Geometric (PyG)
- Trimesh
- NiBabel
- SciPy, NumPy
The following files are required for training and testing. The data should be generated by FreeSurfer or in the same format.
./data/train/SUBJECT_NAME/mri/orig.mgz
./data/train/SUBJECT_NAME/surf/lh.white
./data/train/SUBJECT_NAME/surf/lh.pial
./data/train/SUBJECT_NAME/surf/rh.white
./data/train/SUBJECT_NAME/surf/rh.pial
orig.mgz
is the brain MRI volume. lh.white/rh.white
are the meshes of initial white matter cortical surfaces of left/right cerebrum hemispheres. lh.pial/rh.pial
are ground truth pial surface meshes. An example data is provided in ./data/test/example/
.
We use HCP Young Adult Dataset for training and testing in this paper.
You can run the following codes to predict a pial surface based on the provided example data.
git clone https://github.com/m-qiang/PialNN.git
python eval.py --data_path=./data/test/ --hemisphere=lh --save_mesh_eval=True
The output pial surface mesh will be saved in ./ckpts/eval/
as .obj format.
Download your training data to ./data/train/
and run
python train.py --data_path=./data/train/ --hemisphere=lh
For more information about the training configuration, please see
python train.py --help
The pretrained models of PialNN for both left and right hemispheres are provided:
./ckpts/model/pialnn_pretrained_lh.pt
./ckpts/model/pialnn_pretrained_rh.pt
Download your testing data to ./data/test/
and run
python eval.py --data_path=./data/test/ --hemisphere=lh
Three distance-based matrics will be computed: Chamfer Distance (CD), Average Absolute Distance (AD), and Hausdorff Distance (HD).
For the purpose of debugging or validation, the evaluation result of the provided example data in ./data/test/example/
should be
- CD=0.39, AD=0.20, HD=0.40