The best approach to semantic segmentation of biomedical images
By: Sergey Kolchenko
Tags: Machine Learning, Segmentation, Deep Learning, Pytorch

Photo by Arnold Francisca on Unsplash

Intro.

Semantic segmentation is a problem of computer vision in which our task is to assign a class to each pixel in the image using that image as an input. In the case of semantic segmentation, we don’t care if we have multiple instances (objects) of the same class, we simply label them all with their class. There are multiple courses\reviews on different computer vision problems (for example https://nanonets.com/blog/semantic-image-segmentation-2020/), but one picture can summarize different computer vision problems better than a description

Semantic segmentation has lots of applications for biomedical image analysis: X-ray\MRI scans, digital pathology, microscopy, endoscopy, etc. The list of challenges at https://grand-challenge.org/challenges is quite impressive, there are lots of different interesting and important questions to explore.

From the technical point of view, if we think about the semantic segmentation problem, for the N*M*3 (let’s assume we have an RGB image) image we want to produce the corresponding map N*M*k (where k is the number of classes). There are many architectures that can solve this problem, but here I want to talk about two particular architectures, Unet and Unet++.

TL;DR: If you have a semantic segmentation problem start with Unet++ with resnest200e encoder (available at https://github.com/qubvel/segmentation_models.pytorch)

There are many reviews on Unet, how it changed the field forever (the original publication now has 20+k citations, quite impressive). It is unitively very clear architecture, which consists of an encoder, which generates a representation of the image and a decoder, which uses that representation to build the segmentation. Two maps at each spatial resolution are concatenated (the grey arrows), so you combine two different representations of the image together. And it worked!

The next bog thing was using a pre-trained encoder. Think about the image classification problem. Effectively, we try to build a feature representation of the image, such that different classes in that feature space can be separated. We can take (almost) any CNN and use it as an encoder, take features from that encoder and feed them to our decoder. To my best knowledge, one of the first published papers of that kind were from Iglovikov & Shvets and Buslaev et al., where the group of authors used VGG11 and resnet34 respectively to generate better features for Unet decoder and boost its performance.

Unet++ is a recent improvement of Unet architecture (a nice review is also posted on medium), which has multiple skip connections

According to the paper, Unet++ seems to outperform original Unet. Just like in Unet, multiple encoders (backbones) can be used here to generate strong features for the input image.

Which encoder should I take?

Here I want to focus on Unet and Unet++ and compare their performance with different pre-trained encoders. For that, I chose to use a chest X-ray dataset (a combination of two datasets: first, second), with quite a simple task to segment lungs. It is a binary segmentation, so we should assign each pixel a probability of a class “1”, which we can later binarize to make a mask. First, let’s look at the data

Examples of labeled images from chest X-ray dataset (image by Author)

So these are quite large images, typically 2000*2000 pixels, with large masks and visually it seems like finding a lung is not a problem. Using a great segmentation_models_pytorch library we have access to 100+ different pre-trained encoders for Unet and Unet++. Let’s make a quick pipeline to train a model using Catalyst (another great library for pytorch, which helps you to train your models without having to write lots of boring code) and Albumentations (which helps you to apply different transformations to your images)

After all, it is going to looks something like this

1. Define dataset and augmentations. We will resize images to 256*256 and apply some heavy augmentations to the train dataset

2. Define model and loss function. Here we take Unet++ with regnety_004 encoder and train with a sum of DICE + BCE losses using RAdam + Lookahed optimizers

3. Define callbacks and train!

If we do this for Unet and Unet++ with different encoders we can look at the validation quality of each trained model and summarize it like this: