vis.mp4
-
Apr 15, 2025: π We are excited to introduce the inference code for the PraNet series models (PraNet-V1, PVT-PraNet-V1, PraNet-V2, PVT-PraNet-V2) running on the Jittor framework! Performance evaluation and inference speed tests have shown that Jittor significantly boosts inference speed compared to the PyTorch framework. Give it a try and feel the difference! β¨π -
Apr 15, 2025: π Released training/testing code! PraNet-V2 now supports bothπΉ Binary Segmentation (e.g., polyp segmentation) π©Ί
πΉ Multi-Class Segmentation (e.g., multi-organ/tissue segmentation) π
Explore the training and testing details in their respective directories or follow the detailed setup steps below in this README! π
Accurate medical image segmentation is essential for effective diagnosis and treatment. PraNet-V1 enhanced polyp segmentation by introducing a parallel reverse attention (RA) module that leverages background information. However, it is limited in handling complex multi-organ segmentation within medical imaging datasets. To extend the applicability of RA to a broader range of medical imaging tasks, we propose a Dual-Supervised Reverse Attention (DSRA) module, which incorporates explicit background supervision, independent background modeling structure and semantically enriched attention fusion. Based on DSRA, we develop the PraNet-V2 framework, which shows strong performance across diverse polyp segmentation datasets. Additionally, leveraging DSRA as a plug-and-play module, we integrate DSRA into three state-of-the-art models for medical image semantic segmentation and achieved a maximum improvement of 1.36% in the Dice score compared to their original performances.
git clone git@github.com:ai4colonoscopy/PraNet-V2.git
cd PraNet-V2
# We use Python3.9, CUDA 12.2, PyTorch2.0.1.
conda env create -f pranet2.yaml
conda activate pranet2
cd PraNet-V2- Polyp datasets
- To download the training and testing datasets, use this Google Drive Link and place it in the directory
./binary_seg/data.
- To download the training and testing datasets, use this Google Drive Link and place it in the directory
- ACDC and Synapse datasets
- Please refer to This Google Drive Link to obtain preprocessed dataset.
We also require the pre-trained weights for PVTV2B2, Res2Net, and VIT, Please click here to download them π―. Then, make sure to place the files in the following locations:
-
ποΈ
pvt_v2_b2.pth(ckpt):./binary_seg/models./multiclass_seg/EMCAD/pretrained_pth/pvt
-
ποΈ
maxvit_rmlp_small_rw_224_sw-6ef0ae4f.pth(ckpt) andmaxxvit_rmlp_small_rw_256_sw-37e217ff.pth(ckpt):./multiclass_seg/MERIT/maxvit./multiclass_seg/MIST/maxvit
-
ποΈ
res2net50_v1b_26w_4s-3cf99910.pth(ckpt):./binary_seg/models
π The final file organization should look like this:
binary
βββ models
β βββ pvt_v2_b2.pth
β βββ res2net50_v1b_26w_4s-3cf99910.pth
βββ ...
multi-class
βββ EMCAD
β βββ pretrained_pth
β β βββ pvt
β β β βββ pvt_v2_b2.pth
β βββ ...
βββ MIST
β βββ pretrained_pth
β β βββ maxvit
β β β βββ maxvit_rmlp_small_rw_224_sw-6ef0ae4f.pth
β β β βββ maxxvit_rmlp_small_rw_256_sw-37e217ff.pth
β βββ ...
βββ MERIT
β βββ pretrained_pth
β β βββ maxvit
β β β βββ maxvit_rmlp_small_rw_224_sw-6ef0ae4f.pth
β β β βββ maxxvit_rmlp_small_rw_256_sw-37e217ff.pth
β βββ ...
π¦ Before running the training or inference scripts, make sure to grab the pre-trained model weights π.
We provide RES-V1.pth for PraNet-V1, RES-V2.pth for PraNet-V2, PVT-V1.pth for PVT-PraNet-V1 and PVT-V2.pth for PVT-PraNet-V2.
We recommend to organize the PraNet series models under ./binary_seg/snapshots following the structure below:
snapshots
βββ PraNet-V1
β βββ RES-V1.pth
βββ PraNet-V2
β βββ RES-V2.pth
βββ PVT-PraNet-V1
β βββ PVT-V1.pth
βββ PVT-PraNet-V2
βββ PVT-V2.pth
For the multi-class segmentation models, place our trained models in their respective model_pth folders. For example:
- EMCAD_Synapse.pth --->
./multiclass_seg/EMCAD/model_pth/Synapse/EMCAD_Synapse.pth - MIST_ACDC.pth --->
./multiclass_seg/MIST/model_pth/ACDC/MIST_ACDC.pth - MIST_Synapse.pth --->
./multiclass_seg/MIST/model_pth/Synapse/MIST_Synapse.pth
The data path is pre-configured. Follow the TODO markers in MyTrain_med.py to set the model save path and select the training model type. (PraNet-V2 with ResNet50 as the backbone; PVT-PraNet-V2 with PVTV2B2 as the backbone)
cd ./binary_seg
python -W ignore ./MyTrain_med.py --model_type PraNet-V2 # train PraNet-V2 with resnet50
python -W ignore ./MyTrain_med.py --model_type PVT-PraNet-V2 # train PVT-PraNet-V2 with PVTV2B2Youβll also need to provide the dataset paths for your environment by following the TODO markers in the training scripts. Once done, run the command:
export CUDA_VISIABLE_DEVICES=0 # using single GPU
### train MERIT (w/ DSTA)###
cd ./multiclass_seg/MERIT
# On ACDC dataset
python -W ignore train_ACDC.py --dual
### train MIST (w/ DSTA)###
cd ./multiclass_seg/MIST
# On ACDC dataset
python -W ignore ACDC_train_test.py --dual
# On Synapse dataset
python -W ignore Synapse_train.py --dual
### train EMCAD (w/ DSTA)###
cd ./multiclass_seg/EMCAD
# On Synapse dataset
python -W ignore train_synapse.py --dualPlease donβt forget to fill in the model path and segmentation result save path as guided by the TODO markers in ./binary_seg/MyTest.py. Once done, run the following command:
cd ./binary_seg
python -W ignore ./MyTest_med.pyThe polyp segmentation results will be saved in your predefined result save path. π
After that, follow the steps in the Evaluation section to obtain all the metrics! π
Next, update the dataset path and model path in the test script based on the TODO markers. We recommend changing the default values directly to avoid overly long command-line arguments. For example:
parser.add_argument('--volume_path', type=str, default='./data/synapse/test_vol_h5_new', help='root dir for validation volume data') # TODO: replace with actual pathFinally, run the corresponding inference scripts to get test results (πLogs will be saved in ./test_log) :
export CUDA_VISIABLE_DEVICES=0 # using single GPU
# MIST
cd ./multiclass_seg/MIST
python -W ignore Synapse_test.py --dual
python -W ignore test_ACDC.py --dual
# EMCAD
cd ./multiclass_seg/EMCAD --dual
python -W ignore test_synapse.py --dual
# γVisualizationγ Use the βis_savefig option to save visualization results, e.g., python -W ignore Synapse_test.py --is_savefig.For the PraNet series models, follow the inference steps to generate segmentation results, which will be saved in the results folder. Afterward, run the eval.py script to generate a performance evaluation table in the eval_results folder.
cd ./binary
python -W ignore ./eval.py
The eval.py script provides four eval_config options for evaluating the performance of the following models: PraNet-V1γPVT-PraNet-V1γPraNet-V2γPVT-PraNet-V2. You can try different configs in the script to check out the evaluation results for these models.
For the three multi-class segmentation models, the evaluation results are already logged in the test_log folder during the inference step.
To facilitate reproducibility and further analysis, we provide our segmentation predictions alongside their ground truth for both binary segmentation and multi-class segmentation tasks. π―
A huge thanks to the following excellent works: PraNet, Polyp-PVT, UACANet, CASCADE, MERIT, MIST, and EMCAD. πβ¨
@article{hu2025pranet2,
title={PraNet-V2: Dual-Supervised Reverse Attention for Medical Image Segmentation},
author={Hu, Bo-Cheng and Ji, Ge-Peng and Shao, Dian and Fan, Deng-Ping},
journal={arXiv preprint arXiv:2504.10986},
year={2025},
}