Road detection using segmentation models and albumentations libraries on Keras | by Insaf Ashrapov | Towards Data Science

In this article, I will show how to write own data generator and how to use albumentations as augmentation library. Along with segmentation_models library, which provides dozens of pretrained heads to Unet and other unet-like architectures. For the full code go to Github. Link to dataset.
Check out my Machine & Deep Learning blog https://diyago.github.io/
The task of semantic image segmentation is to label each pixel of an image with a corresponding class of what is being represented. For such a task, Unet architecture with different variety of improvements has shown the best result. The core idea behind it just few convolution blocks, which extracts deep and different type of image features, following by so-called deconvolution or upsample blocks, which restore the initial shape of the input image. Besides after each convolution layer, we have some skip-connections, which help the network to remember about initial image and help against fading gradients. For more detailed information you can read the arxiv article or another article.
We came for practice, let us go for it.
For segmentation we don’t need much data to start getting a decent result, even 100 annotated photos will be enough. For now, we will be using Massachusetts Roads Dataset from https://www.cs.toronto.edu/~vmnih/data/, there about 1100+ annotated train images, they even provide validation and test dataset. Unfortunately, there is no download button, so we have to use a script. This script will get the job done (it might take some time to complete).
Lets take a look at image examples:
Annotation and image quality seem to be pretty good, the network should be able to detect roads.
First of all, you need Keras with TensorFlow to be installed. For Unet construction, we will be using Pavel Yakubovskiy`s library called segmentation_models, for data augmentation albumentation library. I will write more details about them later. Both libraries get updated pretty frequently, so I prefer to update them directly from git.
As a data generator, we will be using our custom generator. It should inherit keras.utils.Sequence and should have defined such methods:
__init__ (class initializing)
__len__ (return lengths of dataset)
on_epoch_end (behavior at the end of epochs)
__getitem__ (generated batch for feeding into a network)
One main advantage of using a custom generator is that you can work with every format data you have and you can do whatever you want — just don’t forget about generating the desired output(batch) for keras.
Here we defining __init__ method. The main part of it is setting paths for images (self.image_filenames) and mask names (self.mask_names). Don’t forget to sort them, because for self.image_filenames[i] corresponding mask should be self.mask_names[i].
Next important thing __getitem__. Usually, we can not store all images in RAM, so every time we generate a new batch of data we should read corresponding images. Below we define the method for training. For that, we create an empty numpy array (np.empty), which will store images and mask. Then we read images by read_image_mask method, apply augmentation into each pair of image and mask. Eventually, we return batch (X, y), which is ready to be fitted into the network.
Next, we define generators, which will be fed into a network:
Data augmentation is a strategy that enables to significantly increase the diversity of data available for training models, without actually collecting new data. It helps to prevent over-fitting and make the model more robust.
There are plenty of libraries for such task: imaging, augmentor, solt, built-in methods to keras/pytorch, or you can write your custom augmentation with OpenCV library. But I highly recommend albumentations library. It’s super fast and convenient to use. For usage examples go to the official repository or take a look at example notebooks.
In our task, we will be using basic augmentations such as flips and contrast with non-trivial such ElasticTransform. Example of them you can in the image above.
After defining the desired augmentation you can easily get your output this:
We will be using common callbacks:
ModelCheckpoint — allows you to save weights of the model while training
ReduceLROnPlateau — reduces training if a validation metric stops to increase
EarlyStopping — stop training once metric on validation stops to increase several epochs
TensorBoard — a great way to monitor training progress. Link to the official documentation
As the model, we will be using Unet. The easiest way to use it just get from segmentation_models library.
backbone_name: name of classification model for using as an encoder. EfficientNet currently is state-of-the-art in the classification model, so let us try it. While it should give faster inference and has less training params, it consumes more GPU memory than well-known resnet models. There are many other options to try
encoder_weights — using imagenet weights speeds up training
encoder_freeze: if True set all layers of an encoder (backbone model) as non-trainable. It might be useful firstly to freeze and train model and then unfreeze
decoder_filters — you can specify numbers of decoder block. In some cases, a heavier encoder with a simplified decoder might be useful.
After initializing Unet model, you should compile it. Also, we set IOU ( intersection over union) as metric we will to monitor and bce_jaccard_loss (binary cross-entropy plus jaccard loss) as the loss we will optimize. I gave links, so won’t go here for further detail for them.
After starting training you can for watching tensorboard logs. As we can see model train pretty well, even after 50 epoch we didn’t reach global/local optima.
So we have 0.558 IOU on validation, but every pixel prediction higher than 0 we count as a mask. By picking the appropriate threshold we can further increase our result by 0.039 (7%).
Metrics are quite interesting for sure, but a much more insightful model prediction. From the images below we see that our network caught up the task pretty good, which is great. For the inference code and for calculating metrics you can read full code.
@phdthesis{MnihThesis,    author = {Volodymyr Mnih},    title = {Machine Learning for Aerial Image Labeling},    school = {University of Toronto},    year = {2013}}