This repository provides the official implementation of the Trexplorer: Recurrent DETR for Topologically Correct Tree Centerline Tracking paper by Roman Naeem, David Hagerman, Lennart Svensson and Fredrik Kahl. The codebase builds upon DETR and Trackformer.
Tubular structures with tree topology such as blood vessels, lung airways, and more are abundant in human anatomy. Tracking these structures with correct topology is crucial for many downstream tasks that help in early detection of conditions such as vascular and pulmonary diseases. Current methods for centerline tracking suffer from predicting topologically incorrect centerlines and complex model pipelines. To mitigate these issues we propose Trexplorer, a recurrent DETR based model that tracks topologically correct centerlines of tubular tree objects in 3D volumes using a simple model pipeline. We demonstrate the model's performance on a publicly available synthetic vessel centerline dataset and show that our model outperforms the state-of-the-art on centerline topology and graph-related metrics, and performs well on detection metrics.
- Clone this repository:
git clone https://github.com/RomStriker/Trexplorer.git - Install requirements:
pip install -r requirements.txt - Install PyTorch 2.2 with CUDA 11.8:
pip install torch==2.2.0 torchvision==0.17.0 torchaudio==2.2.0 --index-url https://download.pytorch.org/whl/cu118
- Download the Synthetic Vessel Dataset.
- Convert the graph vtk files to the required annotation format [example].
- Create the following directory structure:
The '_train', '_val', and '_test' directories contain the training, validation, and test images respectively. The '_val_sub_vol' directory contains the validation images that are used for patch-level evaluation. The 'masks' directories contain the binary masks of the vessel trees. The 'annots' directories contain the annotation files in the required format. The 'annots_val_sub_vol.pickle' [example] file contains the annotations for the validation sub-volume images.
Trexplorer ├── data │ ├── synthetic │ │ ├── annots_train │ │ ├── annots_val │ │ ├── annots_val_sub_vol │ │ ├── annots_test │ │ ├── images_train │ │ ├── images_val │ │ ├── images_val_sub_vol │ │ ├── images_test │ │ ├── masks_train │ │ ├── masks_val │ │ ├── masks_val_sub_vol │ │ ├── masks_test │ │ ├── annots_val_sub_vol.pickle
The training script uses the configuration file ./configs/train.yaml to set the hyperparameters. To train the model, run the following command from the root directory:
python ./src/train_trx.py
For distributed training, use the following command:
python -m torch.distributed.launch --nproc_per_node=NUM_GPUS ./src/train_trx.py
For evaluation, in addition to ./configs/train.yaml, we use ./configs/eval.yaml. To evaluate the model, run the following command from the root directory:
python ./src/train_trx.py with eval
For distributed training, use the following command:
python -m torch.distributed.launch --nproc_per_node=NUM_GPUS ./src/train_trx.py with eval
If you use this software in your research, please cite our publication:
@InProceedings{10.1007/978-3-031-72120-5_69,
author="Naeem, Roman
and Hagerman, David
and Svensson, Lennart
and Kahl, Fredrik",
title="Trexplorer: Recurrent DETR for Topologically Correct Tree Centerline Tracking",
booktitle="Medical Image Computing and Computer Assisted Intervention -- MICCAI 2024",
year="2024",
publisher="Springer Nature Switzerland",
address="Cham",
pages="744--754",
abstract="Tubular structures with tree topology such as blood vessels, lung airways, and more are abundant in human anatomy. Tracking these structures with correct topology is crucial for many downstream tasks that help in early detection of conditions such as vascular and pulmonary diseases. Current methods for centerline tracking suffer from predicting topologically incorrect centerlines and complex model pipelines. To mitigate these issues we propose Trexplorer, a recurrent DETR based model that tracks topologically correct centerlines of tubular tree objects in 3D volumes using a simple model pipeline. We demonstrate the model's performance on a publicly available synthetic vessel centerline dataset and show that our model outperforms the state-of-the-art on centerline topology and graph-related metrics, and performs well on detection metrics. The code is available at https://github.com/RomStriker/Trexplorer.",
isbn="978-3-031-72120-5"
}