VesselFusion

English | 日本語

Official implementation of VesselFusion: Diffusion Models for Vessel Centerline Extraction from 3D CT Images. VesselFusion is a conditional diffusion model that predicts coronary centerline point sets from CT patches. It uses a coarse-to-fine (C2F) coordinate representation and voxel voting over predictions generated from different initial noise samples.

Paper settings

The default configuration follows the paper:

• ImageCAS split: train/validation/test = 7:1:2.
• CT patches: 64 x 64 x 64; training overlap 0.5 and inference overlap 0.
• C2F grid: 16 x 16 x 16; at most 320 centerline elements per patch.
• Denoiser: Transformer conditioned on 3D CNN features.
• Training: 300 epochs, batch size 48, AdamW, learning rate 1e-4, weight decay 1e-2.
• Inference: 100-step DDIM per seed, 100 generated seeds, and voting over 100 seeds.

Settings shared by multiple stages are defined in config.yml. Training and model hyperparameters are Typed Argument Parser (TAP) fields in main.py.

config.yml setting	Meaning	CLI override
inference.ddim_steps	DDIM steps used to generate one seed	--ddim-steps for generate/predict
inference.generation_seeds	number of seeds to generate	--num-seeds for generate; --generation-seeds for predict
inference.voting_seeds	number of generated seeds used for voting	--num-seeds for vote; --voting-seeds for predict

Installation

Python 3.12 and uv are required. Run all commands from the repository root.

uv sync --group dev

Log in when using Weights & Biases online:

uv run wandb login

Replace angle-bracket placeholders such as <training-run> with actual names or paths.

Placeholder	Meaning
<training-run>	training run name
<checkpoint>	checkpoint name
<inference-run>	test inference run name
<vote-name>	voting result name
<predictions-dir>	directory containing *.points.pt
<evaluation-dir>	evaluation output directory
<case-id>	case ID
<ct-path>	path to an input CT volume
<production-run>	production inference run name
<wandb-project>	W&B project name
<wandb-entity>	W&B entity name

Download and place ImageCAS

Download and extract ImageCAS from the ImageCAS Kaggle dataset page.

The extracted images and labels are each distributed across multiple directories. Merge all image files into one data/1-1000_img directory and all label files into one data/1-1000_label directory. Place the NIfTI files directly in these directories without retaining the downloaded subdirectory hierarchy.

data/
├── 1-1000_img/
│   ├── 1.img.nii.gz
│   ├── 2.img.nii.gz
│   └── ...
└── 1-1000_label/
    ├── 1.label.nii.gz
    ├── 2.label.nii.gz
    └── ...

File names must match data.image_pattern and data.label_pattern in config.yml.

Data preparation

Run the following stages in order:

1. Split paired images and labels into train, validation, and test sets.
2. Skeletonize labels in each split with Kimimaro/TEASAR.
3. Build training patches from the train and validation sets.

1. Split the dataset

uv run vesselfusion-split-data \
  --image-dir data/1-1000_img \
  --label-dir data/1-1000_label

The default mode creates symbolic links. Use --mode copy to copy files and --overwrite to replace an existing split. Split assignments are written to data/split_manifest.json.

2. Skeletonize labels

uv run vesselfusion-skeletonize --split all --workers 4

--split accepts train, val, test, or all. --workers sets the number of concurrently processed cases.

3. Build training data

Sequential version:

uv run vesselfusion-build-data --split all

Multiprocessing version:

uv run vesselfusion-build-data-multiprocessing --split all --workers 4

With --split all, training_data is created only for train and validation. Use --overwrite to rebuild existing artifacts.

The prepared data layout is:

data/
├── 1-1000_img/
├── 1-1000_label/
├── split_manifest.json
├── train/
│   ├── img/
│   ├── label/
│   ├── centerline/
│   └── training_data/
├── val/
│   ├── img/
│   ├── label/
│   ├── centerline/
│   └── training_data/
└── test/
    ├── img/
    ├── label/
    └── centerline/

Training

uv run vesselfusion-train --run_name <training-run>

The default W&B project is vesselfusion. Use --wandb_project <wandb-project> and --wandb_entity <wandb-entity> to change it, or --wandb_offline for local logging only.

outputs/training/<training-run>/
├── checkpoints/
│   ├── <checkpoint>.ckpt
│   └── last.ckpt
├── wandb/
├── config.yml
└── train_args.json

The checkpoint with the lowest validation loss and the final checkpoint are saved.

Test inference and voting

Separate generation and voting

Generate and save a point set for each seed:

uv run vesselfusion-generate \
  --checkpoint outputs/training/<training-run>/checkpoints/<checkpoint>.ckpt \
  --split test \
  --run-name <inference-run> \
  --num-seeds 100 \
  --paper-protocol

outputs/eval/<inference-run>/
├── config.yml
└── generation/
    ├── metadata.json
    └── samples/
        └── <case-id>/
            ├── case.pt
            ├── seed_000000.points.pt
            ├── seed_000001.points.pt
            └── ...

Create a voting result from saved seeds:

uv run vesselfusion-vote \
  --generation outputs/eval/<inference-run>/generation \
  --num-seeds 10 \
  --name <vote-name>

outputs/eval/<inference-run>/voting/<vote-name>/
├── config.yml
├── metadata.json
└── predictions/
    ├── <case-id>.points.pt
    └── <case-id>.centerline.nii.gz

Voting uses the lowest saved seed values up to --num-seeds. The selected values are stored in metadata.json and in the used_seeds field of each <case-id>.points.pt.

Integrated generation and voting

uv run vesselfusion-predict \
  --checkpoint outputs/training/<training-run>/checkpoints/<checkpoint>.ckpt \
  --split test \
  --run-name <inference-run> \
  --generation-seeds 100 \
  --voting-seeds 10 \
  --paper-protocol

Final predictions are written to outputs/eval/<inference-run>/predictions. Add --save-seeds to save per-seed point sets.

--paper-protocol restricts inference to patches containing at least one ground-truth centerline point. Use it for paper-protocol test evaluation and omit it for unlabeled images.

Generation displays progress for completed cases and for patch batch × seed × DDIM step within the current case.

Prediction visualization

One points.pt file:

uv run vesselfusion-visualize \
  --points <predictions-dir>/<case-id>.points.pt \
  --split test

The default output is <predictions-dir>/vis/<case-id>.png. Ground truth is loaded from data/test/centerline/<case-id>.pt. Use --ground-truth to specify another file.

Every *.points.pt below a directory:

uv run vesselfusion-visualize-directory \
  --directory <predictions-dir> \
  --split test

The default output directory is <predictions-dir>/vis. Predictions within 3 voxels of ground truth are dark green, more distant predictions are red, and the ground-truth centerline is gray. Change the cutoff with --distance-threshold.

Evaluation

Sequential version:

uv run vesselfusion-evaluate \
  --predictions <predictions-dir> \
  --split test \
  --output <evaluation-dir>

Multiprocessing version:

uv run vesselfusion-evaluate-multiprocessing \
  --predictions <predictions-dir> \
  --split test \
  --workers 4 \
  --output <evaluation-dir>

Select metric groups with --metrics:

uv run vesselfusion-evaluate-multiprocessing \
  --predictions <predictions-dir> \
  --split test \
  --workers 4 \
  --metrics coor pc \
  --output <evaluation-dir>

Metric name	Computation
coor	coordinate precision, recall, and F1 at R = 1, 2, 3
pc	squared symmetric Chamfer distance, Hausdorff distance, and HD95
topo	Betti-0 and Betti-1 with Gudhi
vesselness	Frangi vesselness on false positives farther than R = 3

Omitting --metrics runs all four groups. Results are written to <evaluation-dir>/per_case.csv, <evaluation-dir>/summary.json, and <evaluation-dir>/config.yml. Distances are computed in voxel coordinates.

Production inference

uv run vesselfusion-predict \
  --checkpoint outputs/training/<training-run>/checkpoints/<checkpoint>.ckpt \
  --image <ct-path> \
  --sample-id <case-id> \
  --run-name <production-run>

Artifacts are written to outputs/production/<production-run>. Add --save-seeds to save per-seed point sets. Use vesselfusion-generate and vesselfusion-vote for separate generation and voting.

Commands

Command	Operation
vesselfusion-split-data	split ImageCAS
vesselfusion-skeletonize	skeletonize labels with Kimimaro/TEASAR
vesselfusion-build-data	build training data
vesselfusion-build-data-multiprocessing	multiprocessing training-data builder
vesselfusion-train	train the model
vesselfusion-generate	generate per-seed predictions
vesselfusion-vote	vote over saved seeds
vesselfusion-predict	integrated generation and voting
vesselfusion-visualize	visualize one prediction
vesselfusion-visualize-directory	visualize predictions below a directory
vesselfusion-evaluate	sequential evaluation
vesselfusion-evaluate-multiprocessing	multiprocessing evaluation

Source structure

src/
├── data/                  # splitting, skeletonization, training data, and DataModule
├── model/                 # C2F, network, diffusion, and LightningModule
├── inference/             # seed generation, voting, and inference artifacts
├── eval/
│   ├── metric_components/ # coordinate, point-cloud, topology, and vesselness metrics
│   └── ...                # sequential and multiprocessing evaluation
├── vis/                   # single-file and directory visualization
└── utils/                 # shared configuration

Generated files under data and outputs are ignored by Git.

Citation

@inproceedings{mita2026vesselfusion,
  title     = {VesselFusion: Diffusion Models for Vessel Centerline Extraction from 3D CT Images},
  author    = {Mita, Soichi and Takezaki, Shumpei and Bise, Ryoma},
  booktitle = {IEEE International Symposium on Biomedical Imaging (ISBI)},
  year      = {2026}
}