U-Net: Convolutional Networks for Biomedical Image Segmentation

Olaf Ronneberger, Philipp Fischer, Thomas Brox

Seminal May 2015

U-Net instance masks (right) on a DIC-HeLa time-lapse from the ISBI cell tracking challenge, input on the left.

TL;DR A U-shaped fully convolutional network, with skip connections from a contracting to an expanding path, is trained end to end on only 20-35 annotated images using elastic-deformation augmentation. It took the best warping error on the ISBI EM segmentation challenge and won the 2015 ISBI cell tracking challenge on two light-microscopy datasets with 92% and 77.6% IoU, against 83% and 46% for the second-best entry.

What they do

Biomedical segmentation needs a label for every pixel, and thousands of annotated training images are usually out of reach. The ISBI 2012 EM winner classified each pixel from a sliding-window patch, which is slow and trades localization against context. U-Net extends the fully convolutional network into a symmetric encoder-decoder that segments a whole tile in one pass and trains from a few dozen images.

How it works

Architecture. A contracting path of unpadded 3x3 convs, ReLUs and 2x2 max pooling doubles the channels at each step. The expanding path up-convolves, halves the channels, concatenates the cropped features from the contracting path and convolves again. A final 1x1 conv maps to the classes, for 23 conv layers and no fully connected layers.
Overlap-tile inference. Only valid convolutions are used, so arbitrarily large images are segmented tile by tile, with missing context at the image border filled by mirroring.
Weighted loss. Pixel-wise softmax with cross-entropy, weighted by w = w_c + w0·exp(-(d1+d2)²/2σ²), where d1 and d2 are the distances to the nearest and second-nearest cell. This puts weight on the thin borders between touching cells (w0 = 10, σ ≈ 5 px).
Training. A batch of one large tile with momentum 0.99, plus random elastic deformations: 3x3 grid displacements with a 10 px standard deviation and bicubic interpolation.

Results

ISBI EM challenge (30 training images): warping error 0.000353 and Rand error 0.0382, rank 1 by warping error. The sliding-window CNN scored 0.000420 and 0.0504.
PhC-U373 (35 partially annotated images): 92.0% IoU vs. 83% for the second-best 2015 entry.
DIC-HeLa (20 partially annotated images): 77.6% IoU vs. 46%.
Training takes 10 h on an NVidia Titan GPU (6 GB); segmenting a 512x512 image takes under a second.

Limitations

The paper has no limitations section. By Rand error it is not the best entry on the EM challenge: methods with dataset-specific post-processing on the sliding-window CNN’s output do better (0.0189 for IDSIA-SCI). All errors remain well above the human values reported. Input tile sizes must keep every max-pooling input even-sized.

Videos

U-Net masks (right) on a phase-contrast U373 cell sequence, input on the left.

Figures

From Figure 1: The U-Net architecture, shown for 32x32 pixels at the lowest resolution. Blue boxes are multi-channel ... — From Figure 1: The U-Net architecture, shown for 32x32 pixels at the lowest resolution. Blue boxes are multi-channel feature maps (channel count on top, x-y size at the lower left), white boxes are copied feature maps, and the arrows mark the operations.

From Figure 2: Overlap-tile strategy for segmenting arbitrarily large images, here neuronal structures in EM stacks. ... — From Figure 2: Overlap-tile strategy for segmenting arbitrarily large images, here neuronal structures in EM stacks. Predicting the yellow area needs input from the blue area; missing input at the image border is extrapolated by mirroring.

From Figure 3: Pixel-wise loss weight map for DIC HeLa cells, which forces the network to learn the border pixels bet... — From Figure 3: Pixel-wise loss weight map for DIC HeLa cells, which forces the network to learn the border pixels between touching cells.

From Figure 4: DIC-HeLa segmentation result (random colored masks) with the manual ground truth as yellow borders.

Abstract

There is large consent that successful training of deep networks requires many thousand annotated training samples. In this paper, we present a network and training strategy that relies on the strong use of data augmentation to use the available annotated samples more efficiently. The architecture consists of a contracting path to capture context and a symmetric expanding path that enables precise localization. We show that such a network can be trained end-to-end from very few images and outperforms the prior best method (a sliding-window convolutional network) on the ISBI challenge for segmentation of neuronal structures in electron microscopic stacks. Using the same network trained on transmitted light microscopy images (phase contrast and DIC) we won the ISBI cell tracking challenge 2015 in these categories by a large margin. Moreover, the network is fast. Segmentation of a 512x512 image takes less than a second on a recent GPU. The full implementation (based on Caffe) and the trained networks are available at http://lmb.informatik.uni-freiburg.de/people/ronneber/u-net .

Summaries are written automatically from each paper's text. Figures and videos come from the paper and its project page; all credit goes to the authors.

Matt Strong