Biomedical Image Segmentation: UNet++

Source: Deep Learning on Medium

UNet++ have 3 additions to the original U-Net:

  1. redesigned skip pathways (shown in green)
  2. dense skip connections (shown in blue)
  3. deep supervision (shown in red)

Redesigned skip pathways

In UNet++, the redesigned skip pathways (shown in green) have been added to bridge the semantic gap between the encoder and decoder subpaths.

The purpose of these convolutions layers is aimed at reducing the semantic gap between the feature maps of the encoder and decoder subnetworks. As a result, it is possibly a more straightforward optimisation problem for the optimiser to solve.

Skip connections used in U-Net directly connects the feature maps between encoder and decoder, which results in fusing semantically dissimilar feature maps.

However, with UNet++, the output from the previous convolution layer of the same dense block is fused with the corresponding up-sampled output of the lower dense block. This brings the semantic level of the encoded feature closer to that of the feature maps waiting in the decoder; thus optimisation is easier when semantically similar feature maps are received.

All convolutional layers on the skip pathway use kernels of size 3×3.

Dense skip connections

In UNet++, Dense skip connections (shown in blue) has implemented skip pathways between the encoder and decoder. These Dense blocks are inspired by DenseNet with the purpose to improve segmentation accuracy and improves gradient flow.

Dense skip connections ensure that all prior feature maps are accumulated and arrive at the current node because of the dense convolution block along each skip pathway. This generates full resolution feature maps at multiple semantic levels.

Deep supervision

In UNet++, deep supervision (shown in red) are added to enable the model to operate in accurate mode or fast mode.

For accurate mode, the output from all segmentation branch is averaged, whereas, for fast mode, the final segmentation map is selected from one of the segmentation branches. UNet++ can be pruned to adjust the model complexity, to balance between speed (inference time) and performance.

Zhou et al. conducted experiments to determine the best segmentation performance with different levels of pruning. The metrics used are Intersection over Union and inference time. They experimented on four segmentation tasks: a) cell nuclei, b) colon polyp, c) liver, and d) lung nodule.

The results are as follows: