This is the official PyTorch implementation for the TaG-Net method to handle the head and neck vessel labeling based on CTA image.
@ARTICLE{10032183,
author={Yao, Linlin and Shi, Feng and Wang, Sheng and Zhang, Xiao and Xue, Zhong and Cao, Xiaohuan and Zhan, Yiqiang and Chen, Lizhou and Chen, Yuntian and Song, Bin and Wang, Qian and Shen, Dinggang},
journal={IEEE Transactions on Medical Imaging},
title={TaG-Net: Topology-Aware Graph Network for Centerline-Based Vessel Labeling},
year={2023},
volume={42},
number={11},
pages={3155-3166},
doi={10.1109/TMI.2023.3240825}}
We propose a novel framework for centerline-based vessel labeling. The framework contains two separate models (SegNet and TaG-Net). SegNet is utilized to offer the initial vessel segmentation. TaG-Net is used for centerline labeling. Besides, a graph-based vessel completion method is proposed and utilized in test stage to alleviate the vessel interruption and adhesion resulted from the initial vessel segmentation. Experimental results show that our proposed method can significantly improve both head and neck vessel segmentation and labeling performance.
nnU-Net (3D U-Net cascade) is trained on our dataset to offer the initial vessel segmentation.
Hu range is set as [0, 800] (Window width/level = 800/400).
git clone https://github.com/PRESENT-Y/TaG-Net.git
cd TaG-Net
conda create -n TaG-Net python=3.7
conda activate TaG-Net
conda install pytorch==0.4.1 torchvision cuda90 -c pytorch
pip install -r requirements.txt
mkdir build && cd build
cmake .. && make
./SampleData
.python ./GraphConstruction/gen_cl_graph.py
For visualization of the centerline graph, you can run the following python files.
python ./GraphConstruction/vis_cl_graph.py
CUDA_VISIBLE_DEVICES=0 python ./train.py
You can modify ./cfgs/config_train.yaml
.
CUDA_VISIBLE_DEVICES=0 python ./test.py
We conduct the vessel completion based on the labeled vascular graph (output of the TaG-Net).
sh ./VesselCompletion/vessel_completion.sh
For visualization of the labeled centerline graph, you can run the following python files.
python ./VesselCompletion/vis_labeled_cl_graph.py
The code is released under GPL License (see LICENSE file for details).
If you have some ideas or questions about our research, please contact yaolinlin23@sjtu.edu.cn.