SEGNN for hemodynamics

This repository contains the code for the paper "SE(3) symmetry lets graph neural networks learn arterial velocity estimation from small datasets" [arXiv] [Springer] (FIMH 2023).

Getting started

Set up the Conda environment (CUDA 11.7) by running

conda env create -f environment.yml -n segnn-hemo
conda activate segnn-hemo

and download (a tiny subset of) our dataset from here and pre-trained SEGNN weights from here. Place lumen_tiny.hdf5 into the directory ./lumen-dataset/raw/ and leave neural_network_weights.pt at ./. If everything is set up correctly, after running

python main.py --num_epochs 0

the first thing you will see is pre-processing of the dataset. I recommend ParaView to view the visuals.

Using your own data

You can train SEGNN on your own volume meshes. The easiest way to do so is to process your meshes to match our framework. Save your data in an HDF5 file, e.g. using h5py, matching our structure:

database.hdf5
├── sample-0000
│   ├── velocity           # N x 3
│   ├── pos_tets           # N x 3
│   ├── tets               # M x 4
│   ├── inlet_idcs         # indexing pos_tets (values in [0, N))
│   ├── lumen_wall_idcs    # indexing pos_tets (values in [0, N))
│   └── outlets_idcs       # indexing pos_tets (values in [0, N))
│
├── ...
│
└── sample-1999
    ├── velocity
    ├── pos_tets
    ├── tets
    ├── inlet_idcs
    ├── lumen_wall_idcs
    └── outlets_idcs5

If you have questions, feel free to open an issue or contact me. Depending on the size of your meshes, parallelised training across multiple GPUs might be neccessary. Distributed parallel training is available via the option --num_gpus. If you are not sure how to set the radii in PoolingClusterSampling you can try using scripts/estimate_neighbourhoods_radii.py.

Credits

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

@inproceedings{Suk2023SE3Symmetry,
  author = "Suk, Julian and Brune, Christoph and Wolterink, Jelmer M.",
  title = "{SE(3)} Symmetry Lets Graph Neural Networks Learn Arterial Velocity Estimation from Small Datasets",
  booktitle = "Functional Imaging and Modeling of the Heart",
  year = "2023",
  publisher = "Springer Nature Switzerland",
  address = "Cham",
  pages = "445--454",
  isbn = "978-3-031-35302-4"
}