brats-transunet

BraTS Brain Tumor Segmentation — TransUNet and U-Net checkpoints

Weights for a controlled comparison of U-Net and TransUNet on BraTS-style multi-modal brain MRI, trained from scratch under a fixed 15-epoch budget with identical data and optimization settings across model families.

Code and full write-up: https://github.com/yhj3/brats-unet-vs-transunet

Which file to use

Use checkpoint_jobA_epoch19.pth. It is the best checkpoint and the one reported in the paper table and on the project page. Despite its name, checkpoint_improved_transunet_epoch20.pth is from a different run and is not the reported model; it is kept only for reference.

Context

Foreground Dice is averaged over labels 1–3 and computed only over slices where the class is present — including absent classes inflates the score, since a model then gets credit for predicting "nothing here" on the majority of slices. Numbers here are not comparable to Dice computed over all classes.

For reference, on the same split and metric:

TransUNet's advantage is not uniform: on the compact structures the tuned U-Net is competitive or better (Dice 3: 0.8351 vs. 0.8322), and essentially all of the gain sits in label 2 — the diffuse subregion — at 0.7652 vs. 0.6800.

Input format

4-channel 2D axial slices — FLAIR, T1, T1-CE, T2 — z-score normalized per slice, with a brain crop computed from non-zero tissue. Four output classes (BraTS label 4 is remapped to 3). The checkpoint stores the model configuration and the crop coordinates alongside the weights, so evaluation reconstructs the exact model:

git clone https://github.com/yhj3/brats-unet-vs-transunet && cd brats-unet-vs-transunet
pip install -r requirements.txt
hf download yihangj3/brats-transunet checkpoint_jobA_epoch19.pth --local-dir checkpoints/
python src/eval.py --checkpoint checkpoints/checkpoint_jobA_epoch19.pth

Training data

BraTS multi-modal MRI (FLAIR, T1, T1-CE, T2). The dataset is not redistributed here; obtain it from the BraTS challenge organizers.

Attribution

The TransUNet encoder/decoder modules are adapted from the official TransUNet implementation (Chen et al., 2021), MIT licensed. The 4→3 channel projection, patch-grid handling at non-default input sizes, and the decoder-side skip refinement are my additions.

Contact

Yihang Jiao — [email protected] · https://yhj3.github.io