A segmentation pipeline using Keras and Keras-Transform

Here’s a basic pipeline which handles data augmentation and allows you to quickly start training. In this exemple, create_model and get_paths should be created by you.


You will need Keras, Keras-Transform and OpenCV.

The quickest possible way to install OpenCV is through conda :

conda install -c menpo opencv3

Building a Sequence

Let’s begin by importing Keras, Keras-Transform and OpenCV. We’ll design our Sequence object in the most simple way. It will be provided by a list of paths to the input image and its segmentation mask. The Sequence will load the image and resize it.

import cv2
import numpy as np
from keras.utils import Sequence
from model import create_model  # Function to create a model which does segmentation (U-Net, Tiramisu, etc)
from transform.sequences import (RandomRotationTransformer, SequentialTransformer,
                                 RandomZoomTransformer, RandomHorizontalFlipTransformer)
from your_dataset import get_paths  # Method which will return the path to both the image and its segmentation mask

INPUT_SHAPE = (300, 300)

class YourDatasetSequence(Sequence):
    def __init__(self, paths: [(str, str)], input_shape: (int, int), batch_size: int):
        self.paths = paths
        self.input_shape = input_shape  # We assume that the input_shape is the same as the output
        self.batch_size = batch_size

    def __len__(self):
        return len(self.paths)

    def __getitem__(self, idx: int):
        data = self.paths[idx * self.batch_size:(idx + 1) * self.batch_size]
        X, mask = zip(*data)  # Split the data between input and groundtruth

        # Load and resize the images, we also apply some preprocessing.
        X = [cv2.resize(cv2.imread(fi), self.input_shape) for fi in X]
        X = [self.apply_preprocessing(x) for x in X]

        # Load the masks
        mask = np.array([cv2.resize(cv2.imread(fi), self.input_shape) for fi in mask])

        return X, mask

    def apply_preprocessing(self,x: np.array):
      """Dummy function that puts `x` between 0,1."""
      return x / 255.

Using the Sequence

With this simple Sequence object, you are now able to data augmentation and multiprocessing loading. Now we need to instanciate those Sequences and apply the data augmentation to the training Sequence. Using Keras-Transform, we will be able to apply random transformations on both the input and the mask.

def get_sequences(paths, batch_size=1):
    """Will create the sequences from `paths`"""
    # We split `paths` in 3 and then create 3 sequences.
    train, val, test = np.split(paths, [0.6 * len(paths), 0.8 * len(paths)])
    return YourDatasetSequence(train, INPUT_SHAPE, batch_size), \
           YourDatasetSequence(val, INPUT_SHAPE, batch_size), \
           YourDatasetSequence(test, INPUT_SHAPE, batch_size)

# We create our three Sequences
paths = get_paths()  # [(str,str)]
train_seq, val_seq, test_seq = get_sequences(paths, BATCH_SIZE)

# We create a sequence of transformer to do DA
sequence = SequentialTransformer([RandomZoomTransformer(zoom_range=(0.8, 1.2)),

# By supplying a mask, we can perform transformation on the groundtruth as well.
train_seq = sequence(train_seq, mask=[True, True])

Now that we have everything in place, we only need to instanciate the model and train it!

# Create the model
model = create_model()  # Model

# We can now train and evaluate!
model.fit_generator(train_seq, len(train_seq), validation_data=val_seq, validation_steps=len(val_seq),
                    use_multiprocessing=True, workers=6, epochs=15)

model.evaluate_generator(test_seq, len(test_seq), use_multiprocessing=True, workers=6)


In this post, we’ve created a pipeline for segmentation using Keras and Keras-Transform. With Sequences, we can safely train our model using multiprocessing. If you have any question or comment, feel free to e-mail me!

Written on September 11, 2017