U-Net: Convolutional Networks for Biomedical Image Segmentation

For a long time the bane of deep networks has been data availability and this is especially prevalent in the biomedical imaging field as high quality clinical samples are rare. To address this problem the authors propose a clever architecture,the U-Net, which excels at providing precise segmentation with limited training images.

Architecture

The architecture involves a contracting path followed by an expansive path. The contracting path is traditional in the sense that we're compressing the feature maps and expanding the channels until we reach our lowest resolution. Now, in the expansive path we employ many feature channels allowing contextual information to flow upwards in the model. Combining this with the localization information from the contracting path we end up with a more precise segmentation. In the final layer a 1x1 convolution is used to map the 64-dimensional feature vectors to the desired number of classes.

Weight map

The authors propose an interesting approach to solve segmentation of touching objects. As a pre-processing step a weight map is calculated from the ground truth segmentation, which assigns larger weights to bordering pixels. This is then used to create a weighted loss function where separating background labels between touching objects obtain a larger loss - forcing the model to learn the small separation borders.

Data augmentation

The final interesting takeaway from this paper is their heavy usage of data augmentation to reduce the need for excessive unique training samples. Their augmentation seems to depend on the kind of task at hand, meaning that an understanding of typical deformations in your domain is important to generate good results. They note that "especially random elastic deformations of the training samples seems to be the key".

Final thoughts

Per usual I don't care to touch on the results. The u-net architecture seems to achieve strong performance across a variation of biomedical segmentation applications and future work seems promising.