This repository contains code for a image segmentation model based on UNet++: A Nested U-Net Architecture for Medical Image Segmentation implemented in PyTorch.
- 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.txt
Make sure to put the files as the following structure (e.g. the number of classes is 2):
inputs
└── <ica>
├── images
| ├── nj_100_1_LCA_RAO.png
│ ├── nj_100_2_LCA_RAO.png
│ ├── nj_100_1_LCA_RAO.png
│ ├── ...
|
└── masks/0
├── nj_100_1_LCA_RAO.png
├── nj_100_2_LCA_RAO.png
├── nj_100_1_LCA_RAO.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 ica dataset (96x96) .
Model | IoU | Loss |
---|---|---|
ica_NestedUNet_1000 | 0.839 | 0.365 |
ica_NestedUNet_woDS | 0.842 | 0.354 |