This is a warehouse for SwinTransformerV2-UNet using pytorch framework, can be used to train your segmentation datasets.
链接：https://pan.baidu.com/s/13CZeZxIJlFo9NgtOSoqJpw?pwd=0615
提取码：0615
--来自百度网盘超级会员V3的分享
├── datasets: Load datasets
├── mydataset.py: Build the train dataloader and valid dataloader
├── ext_transforms.py: Additional data augmentation methods
├── models: SwiftFormer Model
├── cbam.py: Construct "CBAM: Convolutional Block Attention Module" module.
├── swintransformerv2.py: Construct "swintransformerv2" module.
├── swinv2UNet.py: Construct "SwinTransformerV2UNet" module.
├── util:
├── losses.py: Construct "DiceScore", "DiceBCELoss" and "Dice_th_pred" modules.
├── scheduler.py: Construct a lr_scheduler.
├── engine.py: Function code for a training/validation process.
├── model_configs.py: Define config parameters.
└── train_gpu.py: Training model startup file
Before you use the code to train your own data set, please first enter the model_configs.py file and modify the train_bs and self.num_classes parameters. Then, enter the mydataset.py file and modify the prefix parameter of your dataset. The parameter perfix just for joining the data root. If you have multi-class labels, enter the losses.py and modify the DiceBCELoss module, just replace the F.binary_cross_entropy_with_logits to the F.cross_entropy.
python train_gpu.py
@inproceedings{cao2022swin,
title={Swin-unet: Unet-like pure transformer for medical image segmentation},
author={Cao, Hu and Wang, Yueyue and Chen, Joy and Jiang, Dongsheng and Zhang, Xiaopeng and Tian, Qi and Wang, Manning},
booktitle={European conference on computer vision},
pages={205--218},
year={2022},
organization={Springer}
}
@article{liu2021swin,
title={Swin Transformer V2: Scaling Up Capacity and Resolution},
author={Liu, Ze and Hu, Han and Lin, Yutong and Yao, Zhuliang and Xie, Zhenda and Wei, Yixuan and Ning, Jia and Cao, Yue and Zhang, Zheng and Dong, Li and others},
journal={arXiv preprint arXiv:2111.09883},
year={2021}
}