heron117/Nested-UNet11

Code for an image segmentation model based on Nested UNet (Inculdes UNet).

Pytorch Implementation of Nested UNet

This repository contains code for an image segmentation model based on Nested UNet (Inculdes UNet).

Requirements

• PyTorch 1.x
• Albumentations 0.1.12
• Pandas 1.x.x

Training on custom dataset

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
            ├── ...

The file format doesn't matter, it can be modified with the training command.

1. Train the model

python train.py --dataset <dataset name> --arch NestedUNet --img_ext .jpg --mask_ext .png

usage: train.py [-h] [--name NAME] [--epochs N] [-b N] [--arch ARCH]
                [--deep_supervision DEEP_SUPERVISION]
                [--input_channels INPUT_CHANNELS] [--num_classes NUM_CLASSES]
                [--input_w INPUT_W] [--input_h INPUT_H]
                [--loss {BCEDiceLoss,LovaszHingeLoss,BCEWithLogitsLoss}]
                [--dataset DATASET] [--img_ext IMG_EXT] [--mask_ext MASK_EXT]
                [--optimizer {Adam,SGD}] [--lr LR] [--momentum MOMENTUM]
                [--weight_decay WEIGHT_DECAY] [--nesterov NESTEROV]
                [--scheduler {CosineAnnealingLR,ReduceLROnPlateau,MultiStepLR,ConstantLR}]
                [--min_lr MIN_LR] [--factor FACTOR] [--patience PATIENCE]
                [--milestones MILESTONES] [--gamma GAMMA] [--early_stopping N]
                [--num_workers NUM_WORKERS]

optional arguments:
  -h, --help            show this help message and exit
  --name NAME           model name: (default: arch+timestamp)
  --epochs N            number of total epochs to run
  -b N, --batch_size N  mini-batch size (default: 16)
  --arch ARCH, -a ARCH  model architecture: NestedUNet/UNet (default: UNet)
  --deep_supervision DEEP_SUPERVISION
  --input_channels INPUT_CHANNELS
                        input channels
  --num_classes NUM_CLASSES
                        number of classes
  --input_w INPUT_W     image width
  --input_h INPUT_H     image height
  --loss {BCEDiceLoss,LovaszHingeLoss,BCEWithLogitsLoss}
                        loss: BCEDiceLoss | LovaszHingeLoss |
                        BCEWithLogitsLoss (default: BCEDiceLoss)
  --dataset DATASET     dataset name
  --img_ext IMG_EXT     image file extension
  --mask_ext MASK_EXT   mask file extension
  --optimizer {Adam,SGD}
                        loss: Adam | SGD (default: Adam)
  --lr LR, --learning_rate LR
                        initial learning rate
  --momentum MOMENTUM   momentum
  --weight_decay WEIGHT_DECAY
                        weight decay
  --nesterov NESTEROV   nesterov
  --scheduler {CosineAnnealingLR,ReduceLROnPlateau,MultiStepLR,ConstantLR}
  --min_lr MIN_LR       minimum learning rate
  --factor FACTOR
  --patience PATIENCE
  --milestones MILESTONES
  --gamma GAMMA
  --early_stopping N    early stopping (default: -1)
  --num_workers NUM_WORKERS

2. Evaluate

python val.py --name <model name>