TF-Unet

General purpose U-Network implemented in Keras for image segmentation

Getting started • Training • Evaluation

Getting started

Looking for Jupyter notebooks? checkout the training, evaulation and prediction notebooks or run make jupyter to serve them locally. Looking for pre-trained weights? download them here.

Dependencies

To quickly get started make sure you have the following dependencies installed:

• Docker (19 or newer) 📦
• Make [Optional macOS / Linux] 🛠
• Nvidia Container Toolkit [Optional for GPUs] 🗜

Setup

Clone (or download) the repository and cd into it

git clone https://github.com/juniorxsound/TF-Unet.git && cd TF-Unet

Next build the Docker image by simply running make build

The build process will pick either Dockerfile.cpu or Dockerfile.gpu based on your system

Training

This repository uses the ShapeDataset synthetic data generator written by Matterport (in Mask R-CNN). No download is needed, as all data is generated during runtime, here is a sample of the dataset

To start training, simply call make train which will start the training process using the parameters defined in train.py. A model will be saved at the end of the training process into the weights folder in SavedModel format.

If you are interested in following the training process, you can use make log during training to start a Tensorboard server with accuracy and loss metrics being updated every batch.

Tensorboard image here

If you want to train in a Jupyter notebook follow the Training notebook

Evaluation

To quickly evaluate download the pre-trained weights and unzip the contents into the weights folder. To run evaluation simply use make evaluate or the Jupyter Evaluation notebook.

The weights provided were trained for 50 epochs on 8000 samples with batch size of 18. Training takes 5 hours using 2 GTX 2080ti's and reaches 96.56% accuracy.

Prediction

See the Jupyter Prediction notebook.

Architecture

The implementation was inspired by U-Net: Convolutional Networks for Biomedical Image Segmentation

Thanks to

The original paper authors, this Keras UNet implementation, this Tensorflow UNet implementation and Mask R-CNN authors.