Skip to content

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.

License

Notifications You must be signed in to change notification settings

McGregorWwww/UDTransNet

Repository files navigation

[Neural Networks] UDTransNet

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.

framework

Requirements

Install from the requirements.txt using:

pip install -r requirements.txt

Usage

1. Data Preparation

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

2. Training

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'.

3. Testing

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.

3.1. Get Pre-trained Models

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.

3.2. Test the Model and Visualize the Segmentation Results

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.

Codes Used in Our Experiments

Contact

Haonan Wang (haonan1wang@gmail.com)

About

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.

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages