Unet-pytorch

U-Net Architecture Full Pytorch Implementation

Title

U-Net: Convolutional Networks for Biomedical Image Segmentation

Discription

From Paper. The u-net is convolutional network architecture for fast and precise segmentation of images. Up to now it has outperformed the prior best method (a sliding-window convolutional network) on the ISBI challenge for segmentation of neuronal structures in electron microscopic stacks. It has won the Grand Challenge for Computer-Automated Detection of Caries in Bitewing Radiography at ISBI 2015, and it has won the Cell Tracking Challenge at ISBI 2015 on the two most challenging transmitted light microscopy categories (Phase contrast and DIC microscopy) by a large margin (See also our annoucement).

Model Architecture

Data

Dataset could be downloaded either from this website or here(you would need to create log-in for this one)
Please run dataset.py python file to create dataset that will be used in training.

Training

$ python train.py 
            --lr YOUR LEARNING RATE 
            --batch_size YOUR BATCH SIZE
            --num_epoch YOUR EPOCH 
            --data_dir [data directory] \ (i.e. './drive/My Drive/Python/Unet-pytorch/datasets' in GoogleColab)
            --ckpt_dir [checkpoint directory] \
            --log_dir [log directory] \
            --result_dir [result directory] \
            --mode train \
            --train_continue off

Change 'train_continue' varialbe to 'on' if you want to continue training from where you left off
Generated images from training are saved in the images subfolder along with [result directory] folder.

Test

$ python train.py 
            --lr YOUR LEARNING RATE 
            --batch_size YOUR BATCH SIZE
            --num_epoch YOUR EPOCH 
            --data_dir [data directory] \ (i.e. './drive/My Drive/Python/Unet-pytorch/datasets' in GoogleColab)
            --ckpt_dir [checkpoint directory] \
            --log_dir [log directory] \
            --result_dir [result directory] \
            --mode test \
            --train_continue off

Change testing code is embedded in the game train.py file
Change '--mode' to 'test' to begin testing
Generated images from training are saved in the images subfolder along with [result directory] folder.

Tensorboard

$ tensorboard --logdir ./log/[data name] \
                  --port 6006

%load_ext tensorboard --> Activate Tensorboard in Jupyter Notebook or Google Colab
Type URL http://localhost:6006 to view training logs recorded in tensorboard (Default)

Results

First row: input; 
Second row: label;
Thrid row: Prediction output;

The results were generated by a network trained with the dataset in 100 epochs.

hyunjoonbok/Unet-pytorch