U-Net: Convolutional Network for Biomedical Image Segmentation

Implemantation of U-Net architecture using Keras framework

Brief info

U-Net is a convolutional network architecture for fast and precise segmentation of images. Input is a grey scale 512x512 image in jpeg format, output - a 512x512 mask in png format. Below is the implemented model's architecture

Project scheme

.
├── model
│   ├── unet.py             # defines U-Net class
│   └── utils.py            # layers for U-Net class
├── tools                     
│   ├── data.py             # generates data 
│   └── image.py            # image-related functions  
├── images                     
│   ├── img                 # image examples for readme
│   └── mask                # generated by nn masks
├── main.py                 # main script to run
└── ...

Requirements

• Python 3.7
• Anaconda 3

Installation

After installing the Python 3.7 version of Anaconda, clone the project

git clone https://github.com/sevakon/unet-keras.git

Create and activate a tensorflow virtual environment (no GPU support):

conda create -n tensorflow_env tensorflow
conda avtivate tensorflow_env

Alternatively, create a tensorflow virtual environment with GPU support:

conda create --name tf_gpu tensorflow-gpu 
conda avtivate tf_gpu

Install other packages to the virtual environment:

conda install -c conda-forge keras   # installs Keras

conda install matplotlib               # installs matpoltlib

conda install scikit-image             # installs skimage

Instructions

All the data preparation happens is data.py file. Data augmentation can be added to ImageDataGenerator parameters

    image_datagen = ImageDataGenerator(rescale=1. / 255)      # add data augmentation here
    mask_datagen = ImageDataGenerator(rescale=1. / 255)       # as function parameter

Note that your data should be structured like this:

.
├── train_path
│   ├── img                 # training images
│   │   ├── 1.jpg                       
│   │   ├── 2.jpg                 
│   │   └── ...    
│   └── mask                # training masks
│       ├── 1.png    
│       ├── 2.png     
│       └── ...    
├── test_path               # folder with images fot test
│       ├── 1.jpg      
│       ├── 2.jpg
│       └── ...    
├── save_path               # where results of tests will be saved
└── ...

Define your config in main.py file

img_height =          # set your image resolution (512x512 or 256X256 recommended)
img_width =           # every image from your dataset will be resacled
img_size = (img_height, img_width)
train_path = ''       # full path to the folder with training data
test_path = ''        # full path to the folder with testing data
save_path = ''        # full path to the folder to save results on testing data
model_name = 'unet_model.hdf5'     # how to name your model save
model_weights_name = 'unet_weight_model.hdf5'

Model Training

To start training your model, run main.py script. All the samples in the dataset will be first made square with black paddings added at random. Then the samples are pushed into training/testing generators, the pretrained weights (if defined) are loaded into the model, and the fitting process begins. At the end of the training, testing samples are predicted to see how well your model behaves on real data. If further training is required, run main.py script again, commenting prepare_dataset function since the dataset was already prepared during the first training.

Model Predict

To predict results with the already trained model, run predict.py script:

python predict.py n

where n is the number of samples that await prediction. Please note that testing samples must be structured like so:

.
├── data
│   ├── test               # testing images
│   │   ├── 1.jpg          # must have names of             
│   │   ├── 2.jpg          # ascending numbers from 1-to-n     
│   │   └── ...    
│   └── results            # predicted masks 
├── predict.py             # script
├── weight.hdf5            # file with the weights     
└── ...

About Keras

Keras is a minimalist, highly modular neural networks library, written in Python and capable of running on top of either TensorFlow or Theano. It was developed with a focus on enabling fast experimentation. Being able to go from idea to result with the least possible delay is key to doing good research.

Use Keras if you need a deep learning library that:

allows for easy and fast prototyping (through total modularity, minimalism, and extensibility). supports both convolutional networks and recurrent networks, as well as combinations of the two. supports arbitrary connectivity schemes (including multi-input and multi-output training). runs seamlessly on CPU and GPU. Read the documentation Keras.io

Keras is compatible with: Python 2.7-3.5.