Segmentation

We have already learnt about Object Detection, which allows us to locate objects in the image by predicting their bounding boxes. However, for some tasks we do not only need bounding boxes, but also more precise object localization. This task is called segmentation.

Segmentation can be viewed as pixel classification, whereas for each pixel of image we must predict its class (background being one of the classes). There are two main segmentation algorithms:

• Semantic segmentation only tells pixel class, and does not make a distinction between different objects of the same class
• Instance segmentation divides classes into different instances.

For instance segmentation 10 sheep are different objects, for semantic segmentation all sheep are represented by one class.

Image from this blog post

There are different neural architectures for segmentation, but they all have the same structure:

• Encoder extracts features from input image
• Decoder transforms those features into the mask image, with the same size and number of channels corresponding to the number of classes.

Image from this publication

Prerequsites

To begin with, we will import required libraries, and check if there is GPU available for training.

The Dataset

We will use the PH² Database of dermoscopy images of human nevi. This dataset contains 200 images of three classes: typical nevus, atypical nevus, and melanoma. All images also contain corresponding mask that outline the nevus.

The code below downloads the dataset from the original location and decompresses it. You would need to have unrar utility installed in order for this code to work, you may install it using sudo apt-get install unrar on Linux, or by downloading command-line version for Windows here.

Now we will define the code to load the dataset. We will transform all images into 256x256 size, and split the dataset into train and test part. This function returns train and test datasets, each containing original images and masks outlining the nevus.

Let's now plot some of the images from the dataset to see how they look like:

We will also need dataloaders to feed the data into our neural network.

SegNet

The simplest encoder-decoder architecture is called SegNet. It uses standard CNN with convolutions and poolings in the encoder, and deconvolution CNN that includes convolutions and upsamplings in decoder. It also relies on batch normalization to train multi-layered network successfully.

Image from this paper: Badrinarayanan, V., Kendall, A., & Cipolla, R. (2015). SegNet: A deep convolutional encoder-decoder architecture for image segmentation

We should especially mention the loss function that is used for segmentation. In classical autoencoders we need to measure the similarity between two images, and we can use mean square error to do that. In segmentation, each pixel in the target mask image represents the class number (one-hot-encoded along the third dimension), so we need to use loss functions specific for classification - cross-entropy loss, averaged over all pixels. If the mask is binary (as in our example) - we will use binary cross-entropy loss (BCE).

Training loop is defined in the usual way:

To evaluate our model, we will just plot target masks and predicted masks for a number of images:

There are also some formal metrics to evaluate the performance, which you can read about here. The easiest one to understand is pixel accuracy - a percentage of pixels classified correctly.

U-Net

SegNet architecture is very natural, but it is not the most accurate. Indeed, we first apply pyramid CNN architecture to the original image, which reduces the spatial accuracy of image features. Then, when we reconstruct the image, we cannot correctly reconstruct the pixel positions.

This leads us to the idea of skip connections between convolution layers in encoder and decoder. This architecture is very common for semantic segmentation, and is called U-Net. Skip connections at each convolution level helps network not to lose information about features from original input at this level.

We will use quite simple CNN architecture here, but U-Net can also use more complex encoder for feature extraction, such as ResNet-50.

Image from paper: Ronneberger, Olaf, Philipp Fischer, and Thomas Brox. U-Net: Convolutional networks for biomedical image segmentation.