This repository contains code for a image segmentation model based on UNet++: A Nested U-Net Architecture for Medical Image Segmentation implemented in PyTorch.
[NEW] Add support for multi-class segmentation dataset.
[NEW] Add support for PyTorch 1.x.
- PyTorch 1.x or 0.41
- Create an anaconda environment.
conda create -n=<env_name> python=3.6 anaconda conda activate <env_name>- Install PyTorch.
conda install pytorch torchvision cudatoolkit=10.1 -c pytorch- Install pip packages.
pip install -r requirements.txtTraining on 2018 Data Science Bowl dataset
- Download dataset from here to inputs/ and unzip. The file structure is the following:
inputs └── data-science-bowl-2018 ├── stage1_train | ├── 00ae65... │ │ ├── images │ │ │ └── 00ae65... │ │ └── masks │ │ └── 00ae65... │ ├── ... | ... - Preprocess.
python preprocess_dsb2018.py- Train the model.
python train.py --dataset dsb2018_96 --arch NestedUNet- Evaluate.
python val.py --name dsb2018_96_NestedUNet_woDS- Clone LovaszSoftmax from bermanmaxim/LovaszSoftmax.
git clone https://github.com/bermanmaxim/LovaszSoftmax.git - Train the model with LovaszHingeLoss.
python train.py --dataset dsb2018_96 --arch NestedUNet --loss LovaszHingeLoss Make sure to put the files as the following structure (e.g. the number of classes is 2):
inputs └── <dataset name> ├── images | ├── 0a7e06.jpg │ ├── 0aab0a.jpg │ ├── 0b1761.jpg │ ├── ... | └── masks ├── 0 | ├── 0a7e06.png | ├── 0aab0a.png | ├── 0b1761.png | ├── ... | └── 1 ├── 0a7e06.png ├── 0aab0a.png ├── 0b1761.png ├── ... - Train the model.
python train.py --dataset <dataset name> --arch NestedUNet --img_ext .jpg --mask_ext .png - Evaluate.
python val.py --name <dataset name>_NestedUNet_woDS Here is the results on DSB2018 dataset (96x96) with LovaszHingeLoss.
| Model | IoU | Loss |
|---|---|---|
| U-Net | 0.839 | 0.365 |
| Nested U-Net | 0.842 | 0.354 |
| Nested U-Net w/ Deepsupervision | 0.843 | 0.362 |