This repo holds code for BiUNet: Towards More Effective U-Net with Bi-Level Routing Attention.
conda create -n biunet python=3.7 -y
conda activate biunet
conda install pytorch==1.10.1 torchvision==0.11.2 cudatoolkit=11.3.1 torchaudio==0.10.1 -c pytorch -c conda-forge -y
pip install -r requirements.txt
The original data can be downloaded in following links:
-
QaTa-COV19 Dataset - Link (Original)
-
MoNuSeG Dataset - Link (Original)
Then prepare the datasets in the following format for easy use of the code:
├── datasets
├── QaTa-Covid19
│ ├── Test_Folder
│ │ ├── img
│ │ └── labelcol
│ ├── Train_Folder
│ │ ├── img
│ │ └── labelcol
│ └── Val_Folder
│ ├── img
│ └── labelcol
└── MoNuSeg
├── Test_Folder
│ ├── img
│ └── labelcol
├── Train_Folder
│ ├── img
│ └── labelcol
└── Val_Folder
├── img
└── labelcol
The first step is to change the settings in Config.py, all the configurations including learning rate, batch size and etc. are in it. Then run:
python train_model.py
Here, we provide pre-trained weights on Covid19 and MoNuSeg, if you do not want to train the models by yourself, you can download them in the following links:
-
Covid：https://pan.baidu.com/s/1hacTAlo2hkNIIUbNsyM26A
training_log: https://pan.baidu.com/s/1LSg_omWie5rC_1hcPhIGDA
-
MoNuSeg: https://pan.baidu.com/s/1j6o1Xz9j6Jpfc4q7rXWrKg
training_log: https://pan.baidu.com/s/1cCKRNZTIxDv3QM9N1atDyA
password: 1234
First, change the session name in Config.py as the training phase. Then run:
python test_model.py
You can get the Dice and IoU scores and the visualization results.
In our code, we carefully set the random seed and set cudnn as 'deterministic' mode to eliminate the randomness. However, there still exsist some factors which may cause different training results, e.g., the cuda version, GPU types, the number of GPUs and etc. See https://pytorch.org/docs/stable/notes/randomness.html for more details.
Open an issue or mail me directly in case of any queries or suggestions.