Awesome

Medical Image Segmentation Foundation Model

This repository provides the official implementation of pretrained medical image segmentation models based on Volume Fusion (VolF), accordingn to the following papers:

• [1] Guotai Wang, Jianghao Wu, Xiangde Luo, Xinglong Liu, Kang Li, Shaoting Zhang, MIS-FM: 3D Medical Image Segmentation using Foundation Models Pretrained on a Large-Scale Unannotated Dataset. arXiv:2306.16925, 2023.

• [2] Guotai Wang, Jia Fu,Jianghao Wu, Xiangde Luo, Yubo Zhou, Xinglong Liu, Kang Li, Jingsheng Lin, Baiyong Sheng, Shaoting Zhang, Volume Fusion-based Self-Supervised Pretraining for 3D Medical Image Segmentation. Submitted to MedIA, 2024.

Details

The following figure shows an overview of our proposed method for pretraining with unannotated 3D medical images. We introduce a pretext task based on pseudo-segmentation, where Volume Fusion is used to generate paired images and segmentation labels to pretrain the 3D segmentation model, which can better match the downstream task of segmentation than existing Self-Supervised Learning (SSL) methods.

VolF for Pre-training

1, The implementation of VolF (and some other self-supervised learning methods) is provided in PyMIC. Please see the PyMIC repository for details. Especially, the following three classes are key components of VolF:

• Crop4VolumeFusion (see pymic/transform/crop4vf.py) A data transform operation that crops two sub-volumes from a large volume with data augmentation.

• VolumeFusion (see pymic/transform/crop4vf.py) Fuse two sub-volumes based on a discretized fusion coefficient map, and obtain the ground truth for pseudo-segmenation task.

• SelfSupVolumeFusion (see pymic/net_run/self_sup/self_volf.py) A class for self-supervised training with VolF, which inherits from a supervised segmentation class.

2, For examples of using VolF for pretraining, please see a demo implemented in PyMIC_examples. In that demo, we train a 3D UNet with the LUNA dataset using VolF, and show applying the pretrained model to a downstream dataset LCTSC2017 for chest organ segmentation.

List of Pretrained Models

Here we provide a list of models pretrained by VoLF using different network structures and datasets. The pretrained weights can be downloaded from Google Drive.

where FMUNet is a modified version of 3D UNet as described in [2], and PCT-Net is a hybrid model using CNN and Transformer proposed in [1].

Demo for Using the Pretrained Model

Package and data requirement

In this demo, we show the use of pretrained FMUNet and PCT-Net for left atrial segmentation. Please install PyMIC and its requirements first. PyMIC can be installed by:

pip install PYMIC==0.5.0

The downstream dataset for this demo can be downloaded from PYMIC_data.

The settings for dataset, network and training/testing can be found in configuration files:

• demo/pctnet_scratch.cfg and demo/fmunet_scratch.cfg for training from scratch.

• demo/pctnet_pretrain.cfg and demo/fmunet_pretrain.cfg for using the pretrained models for training.

After downloading the data, edit the value of train_dir in the configuration files, and make sure the path to the images is correct.

Training

python net_run.py train demo/{NET_NAME}_scratch.cfg

where NET_NAME could be fmunet or pctnet.

Inference

python net_run.py test demo/{NET_NAME}_scratch.cfg

where NET_NAME could be fmunet or pctnet.

Evaluation

pymic_eval_seg -cfg demo/evaluation.cfg

You may need to edit demo/evaluation.cfg to specify the path of segmentation results before evaluating the performance.

In this simple demo, for PCT-Net, the segmentation Dice was 90.71% and 91.80% for training from scratch and from the pretrained weights, respectively. Note that the maximal iteration was set to 6000 in this demo for a quick training. You may use a larger iteration number for an even better performance.

🛡️ License

This project is under the Apache license. See LICENSE for details.