Magician's Corner: 4. Image Segmentation with U-Net

Abstract

A popular deep learning framework (Keras) is applied to the problem of image segmentation using a U-Net.

Figure 1a:

The Dice similarity coefficient (or Dice score) ranges from 0 to 1. (a) In the case where there is no overlap between the reference standard contour (such as provided by a human) and the predicted contour, the Dice score is 0. (b) In cases where there is a small overlap, the Dice score will be a small number, about 0.1 in this case. (c) In the case where they exactly match, the Dice score is 1.0.

Figure 1b:

The Dice similarity coefficient (or Dice score) ranges from 0 to 1. (a) In the case where there is no overlap between the reference standard contour (such as provided by a human) and the predicted contour, the Dice score is 0. (b) In cases where there is a small overlap, the Dice score will be a small number, about 0.1 in this case. (c) In the case where they exactly match, the Dice score is 1.0.

Figure 1c:

The Dice similarity coefficient (or Dice score) ranges from 0 to 1. (a) In the case where there is no overlap between the reference standard contour (such as provided by a human) and the predicted contour, the Dice score is 0. (b) In cases where there is a small overlap, the Dice score will be a small number, about 0.1 in this case. (c) In the case where they exactly match, the Dice score is 1.0.

Figure 2:

U-Net architecture. Each block on the left side consists of a convolution, layer-wise normalization, and maximum value pooling. The latter pools a 2 × 2 block into 1 pixel, thus reducing both X and Y dimensions by a factor of two. The blocks on the right side use transpose convolution with a stride of two to increase the X and Y dimensions by a factor of two. In addition, there is a skip connection where information from the down-resolution blocks is fed to these blocks. Conv = 2D convolution, Norm = layer-wise normalization, MaxPool = maximum value pooling, TransposeConv = transpose convolution, Concat = concatenation of convolution kernel from left block to the kernel in this block. Sigmoid is the function to map to a probability value from 0 to 1.

Figure 3:

Graph shows the Dice loss (which is 1 − Dice similarity coefficient) at each epoch out to 100 epochs. Note that the blue and orange lines track each other until about 10, and after that point, the performance on the training set continues to improve but the performance on validation set plateaus, indicating that overfitting is occurring. Blue line = training loss curve, orange line = validation loss curve.

Figure 4:

The performance on the test set of images. The images on the left are the ground truth (pancreas in cyan) and the right are the predicted. The overall Dice score for all images in the test set is 0.60, which is somewhat less than the validation performance but much less than the training set performance indicating that overfitting is occurring.