U-Net: Convolutional Networks for Biomedical Image Segmentation

In this blog I will talk about one of the most famous network for biomedical image segmentation that is U-Net. U-Net takes the idea from the Fully Connected Network. This modify and extend the FCN network such that it works with few training images and yields more precise segmentation.

The main idea of FCN is use of Upsampling layer to get the full segmentated mask of input images and combining the contracting network output with the upsampling layer output for precise output. Read more about FCN network in this blog.

The important points of U-Net architecture are

• In Upsampling part large number of feature channels are used that allows the network to propagate context information to higher resolution layers. As a consequence, the expansive path is more or less symmetric to the contracting path, and yields a U-shaped architecture.

• Only valid convolutional is used in the proposed network (without any padding). The output segmented map contains the output for pixels for which full context is available. This strategy allows the seamless segmentation of arbitrarily large images by an overlap-tile strategy. To predict the pixels in the border region of the image extrapolation is used. Convolution with padding can be used for fixed dimension input images. 

  

  Overlap-tile strategy for seamless segmentation of arbitrary large image. Prediction of the segmentation in the yellow area, requires image data within the blue area as input.

  
  

• As in case of biomedical imaging very little training data is available, data augmentation by applying elastic deformations to the available training images were used. This allows the network to learn invariance to such deformations. This is particularly important in biomedical segmentation, since deformation used to be the most common variation in tissue and realistic deformations can be simulated efficiently. 

  From the below image, if we remove grid, then it will be hard to differentiate between real and deformed image.

   

Elastic Augmentation

• Another challenge in many cell segmentation task is the separation of touching object of the same class. For this author proposed the use of a weighted loss, where the separating background labels between touching cells obtain a large weight in the loss function. 

  
  

a) Raw image of HeLa cells. b) Overlay with ground truth segmentation. Different color indicate different instances of the HeLa cells. c) Generated segmented mask (white: foreground, black: background). d) Map with a pixel-wise loss weight to force the network to learn the border pixels.

 

U-Net Architecture

 

Each Blue box corresponds to a multi-channel feature map. White boxes represent copied feature maps. The arrows denote the different operations. The number of channels is denoted in top of the box. The x-y size is provided at the lower left edge of the box.

• It consists of a contracting path (left side / encoder) and an expansive path ( right side / decoder).

• The contracting path follows the typical architecture of a convolutional network. It consists of the repeated block of two 3x3 convolution (unpadded convolution), each followed by ReLU and 2x2 max pooling operation with stride 2 for downsampling. At each downsampling number of feature channels get doubled as usual.
  

• Every step in the expansive path consists of an upsampling of the feature map by a 2x2 up-convolution that halves the number of feature channels, a concatenation with the correspondingly cropped feature map from the contracting path, and two 3x3 convolutions each followed by a ReLU. The cropping is necessary due to the loss of border pixels in every convolution.

   

  

  2x2 up-convolution
  

  
   

• At the final layer a 1x1 convolution is used to map each 64-component feature vector to the desired number of classes. In total, network has 23 convolution layers.

One may look this contraction and expansion path as getting the information of WHAT and WHERE from the image.

 

Contraction and Expansion Path

 

As we move along the contraction network we loose WHERE information and gain WHAT information of the image and the expansion path recovers the WHERE information by gradually applying upsampling.

Weighted Loss for precise boundary 

For segmentation one can use cross-entropy loss or Dice loss.

As we can see in the above images of HeLa cells that cells are almost connected and with normal loss it's little difficult to separate the instances of cell during prediction.

In the U-Net paper, author talked about using weighted cross entropy loss for precise prediction of overlapping boundary.

Weighted map for each ground truth segmentation in computed to compensate the different frequency of pixels from a certain class in the training data set and to force the network to learn the small separation border which is introduced in between touching cell.  

The separation border is computed using morphological operations. The weight map is then calculated as 

Weighted map

where, wc is the weight map to balance the class frequency (foreground, background), d1(x) is the distance to the nearest cell border at position x, d2(x) is the distance to the second nearest cell border. w0=10 and σ=5 were choosen.

Segmentation Mask

 

We can see the weight map in the above image.

Weighted Cross Entropy Loss

Pixel wise soft-max over the final feature map is calculated and combined with the cross entropy loss. The cross entropy loss is penalized at each position by the weight map w(x), which help the network to learn the separation boundary between touching cells.

Results

EM Segmentation Challenge

• Achieved new best score, in term of warping error which was much better than sliding window CNN network. 

• Network is fast, training time is 10 hours and inference time is 1s per image on 6 GB NVidia Titan GPU.  

   

  

  EM Segmentation Challenge 

  
  

  
  

  
  

 

ISBI Cell Tracking Challenge 

U-Net achieves 92% IoU while second best method algorithm achieves only 83% IOU .

 

Sample Image and Result of ISBI Cell

 

 

 ISBI Cell tracking Challenge

ISBI Cell Tracking Challenge: DIC HeLa

 

U-Net achieved 77.6% IoU while 2nd best method achieved only 46%

That's all for this blog, thanks for reading !!

 

Reference

• https://arxiv.org/pdf/1505.04597.pdf

• https://lmb.informatik.uni-freiburg.de/people/ronneber/u-net/
  

 

Fully Convolutional Networks for Semantic Segmentation

In this blog I will talk about the Fully Convolution Network for Semantic Segmentation Paper. 

Let's first understand what is semantic segmentation problem. Semantic segmentation classify each pixel into one of the class including background without differentiating instance of object.

Semantic Segmentation

In the above image we can see two cow but semantic segmentation does not differentiate instance of connected similar object. One can use Instance segmentation if want to segment two cow separately.

Fully Convolutional Networks for Semantic Segmentation Paper majorly talk about two things

• First is using fully convolutional network pixel-to-pixel end-to-end training, that takes arbitrary size input images and produces corresponding size of output.
  
• Second is using novel "skip" architecture that combine the semantic information from a deep, coarse layer with local appearance information from a shallow, fine layer to produce accurate and detailed segmentation.
  

 

Classification Network to FCN Network for Semantic Segmentation

Lets see what we mean by fully convolution network. Usually, in the classification network like AlexNet, VGG and GoogLeNet input image goes through multiple convolution block followed by fully connected layer and then output layer.

Because of fully connected layer these network requires fixed input dimension images because fully connected layer takes fixed dimension of input.

In FCN, author, replaces this fully connected layer by 1x1 convolutional layer hence making the network fully convolutional.  

In classification network we usually have N number of output node for N classes in output layer, but for semantic segmentation we want pixel mask as an output having spatial dimension same as input. For this author removes the final output layer of classification network and uses deconvolution (upsampling) layer to produce output mask.

Don't get confused with the name of deconvolution, it is not the reverse process of convolution, in deep learning deconvolution or some time also called back convolution or up convolution or transposed convolution is only used for upsampling.   

Note that the deconvolution filter in such layer need not to be fixed (e.g to bilinear upsampling), but can be learned. A stack of deconvolution layers and activation functions can even learn a nonlinear upsampling.

So till now we have seen that, two changes requires in fully connected classification network in order to make it fully convolutional network for semantic segmentation. Let's see how to do these changes.

FCN-32

First recall the architecture of VGG network.

VGG original architecture

After fifth max pooling layer output feature map dimension is 7x7x512 and after that we are using two fully connected layer.

In order to use only convolutional layer, we change original fc6 layer with convolutional layer of filter size 7x7 (7x7x512x4096) which produces 7x7x4096 size feature map and replace fc7 with convolution layer of filter size 1x1 (1x1x4096x4096) which produces 7x7x4096 size of feature map. 

After that this output feature map (7x7x4096) of new fc7 layer passes through one more 1x1 (1x1x4096x21) convolution layer to give individual prediction mask for each of the PASCAL classes (including background), at last deconvolution layer is used to upsample the output 32 times, from 7x7 to 224x224.

 

 

Note for any arbitrary input dimension (512x512),  feature map after last pooling layer will be 1/32 times (16x16) and after deconvolution layer (upsample by 32 times) output mask will have same dimension (512x512) as input.  

This Network is what we call FCN-32. 

FCN-32

 

FCN-32 network outputs the segmented mask of input, but the output maps are coarse (rough output map) because of the 32 pixel stride at the final prediction layer which limit the scale of details in the upsampled output.

 

To overcome this, author uses skip connection which combine the semantic information from deep layers with spatial location information from the shallow layers to produce accurate and detailed segmentation.

This gives us different variant of FCN i.e. FCN-16 and FCN-8.

FCN-16

For FCN-16, output of pool4 is convolved with 1x1 convolution to get class specific prediction for 21 channels. This predicted output is fused with the 2x upsampled output of conv7 and 16x upsampling is performed on the fused output to get the final segmentation mask.

FCN-32, FCN-16 and FCN-8

FCN-8

Similarly for FCN-8, pool3 is convolved with 1x1 convolution to get class specific prediction for 21 classes and this output is fused with 2x upsampled feature map (last) of FCN-16 i.e 2x pool4 and 4x conv7 and after that 8x upsampling is performed on the fused output to get the final segmentation mask.

 

As FCN-8 is having both spatial information and deep semantic information, so it formed better than FCN-16 and FCN-32.

Refining fully convolutional nets by fusing information from layers with different strides improves segmentation detail

Result 

Pixel accuracy, mean accuracy, mean intersection over union (IU) and frequency weighted IU is reported for these 3 network.

 

Comparison of skip FCN's on a subset of PASCAL VOC 2011 validation

That's all for this blog, Thanks for reading.

Reference

• https://arxiv.org/pdf/1411.4038.pdf
• http://cs231n.stanford.edu/slides/2016/winter1516_lecture13.pdf