Multiclass semantic segmentation using U-Net

Multiclass semantic segmentation using U-Net in TensorFlow 2 / Keras.

Directory Structure

All the code files and folders follow the following structure in ipynb file:

├── images_256_256_256.tif
├── masks_256_256_256.tif
├── patches
│   ├── images
│   └── masks
└── model.hdf5

Introduction

The line of research is motivated by the need to accurately segment 4 regions from images about XRM (tomography) scan of a sandstone cylinder of size about 2 mm diameter. To solve this problem, we will use multiclass semantic segmentation using U-Net in TensorFlow 2 / Keras.

It is used U-Net model, which is trained on this Sandstone Dataset. It is used two tif files that include 256 images of 256 x 256: images_256_256_256.tif and masks_256_256_256.tif.

The 4 regions of the images consist of air/void (dark regions), quartz (light grey), clay (darker grey with texture), and pyrite (bright pixels). Clay and quartz are minority classes but since quartz appears in bright, it is easier to segment. Please read the Readme document for more information.

U-Net is a semantic segmentation technique originally proposed for medical imaging segmentation. U-Net was introduced in the paper, U-Net: Convolutional Networks for Biomedical Image Segmentation. The model architecture is fairly simple: an encoder (for downsampling) and a decoder (for upsampling) with skip connections.

The gray arrows indicate the skip connections that concatenate the encoder feature map with the decoder, which helps the backward flow of gradients for improved training.

Instructions

Please, follow the instruction in the provided notebook.

Results

With just 256 images of 256x256 and 50 epochs of training, it achieves a Mean Intersection over Union (IoU) of 0.88 in the unseen test partition.

Predicted segmentation examples