chinmay5 / vesselformer

Vesselformer: Towards Complete 3D Vessel Graph Generation from Images

Requirements

• CUDA>=11.0
• PyTorch>=2.0

For other system requirements please follow

pip install -r requirements.txt

Compiling CUDA operators

cd ./models/ops
python setup.py install

Code Usage

1. Dataset preparation

Please download Synthetic Vessel dataset and organize them as following:

code_root/
└── data/
    └── vessel_data/
        ├── raw/
        ├── seg/
        └── vtk/

After downloading the dataset run the following script to preprocess and prepare the data for training

python generate_data.py

2. Training

2.1 Prepare config file

The config file can be found at .configs/synth_3D.yaml. Make custom changes if necessary.

2.2.a Training on multiple-GPU (e.g. 3 GPUs)

For example, the command for training VesselFormer is following:

python train.py --config configs/synth_3D.yaml --cuda_visible_device 0 1 2 --nproc_per_node 3

3. Evaluation

Once you have the config file and trained model, run following command to evaluate it on test set:

python vesselformer_inference.py --config configs/synth_3D.yaml --model ./trained_weights/last_checkpoint.pt --eval

4. Interactive notebook

Please find the vesselmerger_inference.ipynb for interactive visualization for vesselmerger algorithm