This repo is the official implementation of 'Narrowing the semantic gaps in U-Net with learnable skip connections: The case of medical image segmentation' which is an improved journal version of UCTransNet.

🎉 The manuscript is accepted by Neural Networks, 2024.

Install from the requirements.txt using:

pip install -r requirements.txt

The GlaS dataset is in the 'datasets' folder. ISIC-2018 and Synapse datasets can be downloaded in following links:

• ISIC-2018 Dataset - Link
• Synapse Dataset - Link

Then prepare the datasets in the following format for easy use of the code:

├── datasets
│   ├── GlaS
│   │   ├── Test_Folder
│   │   │   ├── img
│   │   │   └── labelcol
│   │   └── Train_Folder
│   │       ├── img
│   │       └── labelcol
│   ├── ISIC18
│   │   └── Train_Folder
│   │       ├── img
│   │       └── labelcol
│   └── Synapse
│       ├── lists
│       │   └── lists_Synapse
│       ├── test_vol_h5
│       └── train_npz

We use five-fold cross validation strategy to train all the models on all the three datasets.

The first step is to change the settings in Config.py, all the configurations including learning rate, batch size and etc. are in it.

We optimize the convolution parameters in U-Net and the DAT parameters together with a single loss. Run:

python train_kfold.py

The results including log files, model weights, etc., are in '[TaskName]_kfold' folder, e.g., 'GlaS_kfold'.

For GlaS and Synapse, we test the models of five folds and take the average score on the test set.

For ISIC'18, since the annotation of test set is not publicly available, we test the model of each fold on each validation set.

Here, we provide pre-trained weights of five folds on the three datasets, if you do not want to train the models by yourself, you can download them in this Google Drive link.

First, change the session name in Config.py as the training phase.

Then, for GlaS and Synapse, run:

python test_kfold.py

For ISIC, run:

python test_each_fold.py

You can get the Dice and IoU scores and the visualization results.

• UNet++
• Attention U-Net
• MultiResUNet
• MedT
• TransUNet
• Swin-Unet
• MCTrans
• UCTransNet

Haonan Wang (haonan1wang@gmail.com)