Comparing foundation models and nnU-Net for segmentation of primary brain lymphoma on clinical routine post-contrast T1-weighted MRI

This repository contains material associated to this paper.

It contains:

  • link to trained models for segmentation of lymphoma from post-constrast T1-weighted MRI (link)
  • link to trained models for segmentation of enhancing tumor in MSD-BraTS datasets (link)
  • code and material for reproducing the experiments on MSD-BraTS (link)

If you use this material, we would appreciate if you could cite the following reference.

Citation

  • Guanghui Fu, Lucia Nichelli, Dario Herran, Romain Valabregue, Agusti Alentorn, Khê Hoang-Xuan, Caroline Houillier, Didier Dormont, Stéphane Lehéricy, Olivier Colliot. Comparing foundation models and nnU-Net for segmentation of primary brain lymphoma on clinical routine post-contrast T1-weighted MRI. Preprint, https://hal.science/hal-04447318.

Download and install

Data split pre-processing

To be as similar as possible to lymphoma segmentation task, we selected the T1-weighted MRI after gadolinium injection and focused on the enhancing tumor as the target region. Therefore, we need some basic preprocessing to obtain T1-GD from MRI and select enhancing tumors from the label file of MSD-BraTs dataset.

Trained nnU-net models

Lymphoma segmentation in post-contrast T1-weighted MRI

Enhancing tumor segmentation in MSD-BraTs datasets

There is no preprocessing needed for using trained nnU-net model to inference on MSD-BraTS.

Contents for reproducing MSD-BraTS experiments

We provide the following contents for reproduction of MSD-BraTS experiments:

  • list of subjects of MSD-BraTS that were used (link)
  • manual box prompts for SAM and MedSAM models (link)
  • support sets for UniverSeg experiments (link)
  • code to train nnU-Net models (link)
  • code for inference of all models (SAM/MedSAM, UniverSeg, nnU-Net)
  • code for computation of metrics and statistical analysis (link)
  • Note that, here train means train/validation for nnU-Net.

Manual box prompt annotation

SAM and MedSAM require prompts. In this paper, we drew box prompts manually.

The figure below shows our process of using ITK-SNAP for drawing box prompts.

We also provide screen recording videos during annotation: https://owncloud.icm-institute.org/index.php/s/9LWatZ2xDB9SvE0

manual_box

In order to reproduce the experiments, you need the coordinates of the box prompts which are given here:

Generating boxes from grount truth

Also we provide the prompts generated from ground-truth.

The code to generate the boxes from ground truth is provided here:

Support sets for Universeg

The different support sets are given here: brats_support_set.zip

The code to build the support sets:

  • universeg_experiment/universeg_select_support_set.py: this code offers strategies for selecting the support set, a crucial step in preparing data for the UniverSeg [2] model. It includes options to select slices based on their size (largest, smallest, or medium).

Code to train nnU-net

Training nnU-Net only requires the following commands. DATASET_ID and FOLD are defined in the same way as nnU-net official codes. See the official code for details.

  • Plan and preprocessing:
nnUNetv2_plan_and_preprocess -d DATASET_ID --verify_dataset_integrity
  • 3D model training:
nnUNetv2_train DATASET_ID 3d_fullres FOLD
  • 2D model training:
nnUNetv2_train DATASET_ID 2d FOLD

Inference for SAM and MedSAM models

Inference for UniverSeg

Inference for nnU-Net

Inferencing nnU-Net only requires the following commands. DATASET_ID, INPUT_FOLDER and OUTPUT_FOLDER are defined in the same way as nnU-net official codes. See the official code for details.

  • 3D inference:
nnUNetv2_predict -d DATASET_ID -i INPUT_FOLDER -o OUTPUT_FOLDER -f  0 1 2 3 4 -tr nnUNetTrainer -c 3d_fullres -p nnUNetPlans
  • 2D inference:
nnUNetv2_predict -d DATASET_ID -i INPUT_FOLDER -o OUTPUT_FOLDER -f  0 1 2 3 4 -tr nnUNetTrainer -c 2d -p nnUNetPlans

Computation of metrics and statistical analysis

Related codes

  1. SAM [1]: https://github.com/bingogome/samm
  2. MedSAM [5]: https://github.com/bowang-lab/MedSAM
  3. UniverSeg [2]: https://github.com/JJGO/UniverSeg
  4. nnU-Net [4]: https://github.com/MIC-DKFZ/nnUNet
  5. TorchIO [6]: https://torchio.readthedocs.io/index.html

References

  1. Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C Berg, Wan-Yen Lo, et al. Segment anything. In Proc. ICCV 2023, pages 4015–4026, 2023
  2. Victor Ion Butoi, Jose Javier Gonzalez Ortiz, Tianyu Ma, Mert R Sabuncu, John Guttag, and Adrian V Dalca. UniverSeg: Universal medical image segmentation. In Proc. ICCV 2023, pages 21438–21451, 2023.
  3. Michela Antonelli, Annika Reinke, Spyridon Bakas, Keyvan Farahani, Annette Kopp-Schneider, Bennett A Landman, Geert Litjens, Bjoern Menze, Olaf Ronneberger, Ronald M Summers, et al. The medical segmentation decathlon. Nature communications, 13(1):4128, 2022.
  4. Fabian Isensee, Paul F Jaeger, Simon AA Kohl, Jens Petersen, and Klaus H Maier-Hein. nnU-Net: a self-configuring method for deep learning-based biomedical image segmentation. Nature methods, 18(2):203–211, 2021.
  5. Ma, Jun, et al. "Segment anything in medical images." Nature Communications 15.1 (2024): 654.
  6. Pérez-García, Fernando, Rachel Sparks, and Sébastien Ourselin. "TorchIO: a Python library for efficient loading, preprocessing, augmentation and patch-based sampling of medical images in deep learning." Computer Methods and Programs in Biomedicine 208 (2021): 106236.