Coronary mesh convolution

This repository contains code accompanying our (MICCAI 2021) "Workshop on Statistical Atlases and Computational Modelling of the Heart" (STACOM) paper "Mesh convolutional neural networks for wall shear stress estimation in 3D artery models". For questions, feel free to contact me.

Dependencies & packages

Dependencies:

• Python (tested on 3.9.13)
• PyTorch (tested on 1.12.1)
• PyTorch Geometric "PyG" (tested on 2.0.3) with
  • torch-cluster (tested on 1.6.0)
  • torch-scatter (tested on 2.0.9)
  • torch-sparse (tested on 0.6.15)

Packages:

pip install prettytable vtk trimesh potpourri3d tensorboard h5py robust_laplacian

You can install all of these with the provided conda environment file (CUDA 11.6):

conda env create -f environment.yml -n cmc
conda activate cmc

Additionally, we need gauge-equivariant mesh convolution:

git clone https://github.com/Qualcomm-AI-research/gauge-equivariant-mesh-cnn.git
cd gauge-equivariant-mesh-cnn
pip install .

If you get an error regarding OpenMesh, try

conda install -c conda-forge openmesh-python

and then try to install again.

Data

You can download the dataset(s) from here. The physical units for wall shear stress are [dyn/cm^2] = 0.1 [Pa]. We additionally provide pre-trained model weights.

We adapt the dataset-directory structure used by PyTorch Geometric ("PyG"). The directory with the dataset should contain a folder raw with the unprocessed data. Pre-processing creates a folder processed with the transformed data.

vessel-datasets
└── stead
    ├── bifurcating
    │   └── raw
    │       └── database.hdf5
    └── single
        └── raw
            └── database.hdf5

The pre-trained model weights should be placed in a folder model-weights and are loaded automatically if present.

Usage

Experiments are run by executing e.g. (options listed in main.py)

python main.py --model gem_gcn --artery_type single

and produce visualised output in the vis directory which can be viewed with e.g. ParaView. If you get an error Unable to open file try downloading the HDF5 files directly instead of the whole directory and placing them in their respective folders manually. If everything works, first thing you will see is the pre-processing of the training data.

Hyperparameters for neural network training are set in an experiment file, e.g. exp/gem_gcn/stead.py. Training curves can be viewed with TensorBoard for PyTorch via

tensorboard --logdir=runs

This codebase supports parallelisation over multiple GPUs. Just use the command line option with a space-separated list.

python main.py --model gem_gcn --artery_type single --num_epochs 100 --gpu 0 1

Network layout

This repository implements a three-scale mesh-based graph convolutional residual neural network with gauge-equivariant convolution. For details refer to our paper "Mesh convolutional neural networks for wall shear stress estimation in 3D artery models".

DiffusionNet

We have included DiffusionNet as an additional baseline. The code is copy & pasted from this excellent repository.

Publication

If you found this repository useful, please consider citing our paper:

@inproceedings{Suk/et/al/2021,
  author = {Julian Suk and Pim de Haan and Phillip Lippe and Christoph Brune and Jelmer M. Wolterink},
  title = {Mesh convolutional neural networks for wall shear stress estimation in 3D artery models},
  booktitle = {Statistical Atlases and Computational Models of the Heart},
  year = {2021}
}