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:
- 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.
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.
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
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".
We have included DiffusionNet as an additional baseline. The code is copy & pasted from this excellent repository.
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}
}