U-Net: Convolutional Networks for Biomedical Image Segmentation

The problem: a label for every pixel

Image segmentation is the task of assigning a class to every pixel of an image instead of giving one label to the whole picture. In biomedical research the images usually come from microscopy, such as stacks of electron micrographs in which neural tissue has to be traced, or light microscopy recordings of cells that have to be followed over time. The bottleneck is annotation: a human expert has to mark the boundaries of every structure by hand, so annotated datasets stay small and expensive. At the time, the working assumption was that training a deep network well requires many thousand annotated samples^[1]. U-Net was built to work against that assumption.

How the network is built

The architecture has two halves. A contracting path captures context: it processes the image at steadily lower spatial resolution, so that each position in the feature maps summarizes an ever wider neighbourhood of the input. A symmetric expanding path then brings the resolution back up, which the paper describes as enabling precise localization^[1]. Because the output is restored to the size of the input, the network can be trained end to end, from image to segmentation map, without a separate post-processing stage^[1].

Training on very few images

The network and its training strategy rely heavily on data augmentation, used to exploit the available annotated samples more efficiently^[1]. In practice this means the same annotated image is fed to the network many times, each time slightly deformed, shifted, rotated, flipped or changed in intensity, so that a small set of hand-labelled images yields a large and continually changing stream of training examples. The authors present this as a challenge to the consensus that successful deep network training requires many thousand annotated samples^[1], and report that the network can be trained end to end from very few images^[1].

Results

The paper reports two evaluations. On the ISBI challenge for segmentation of neuronal structures in electron microscopic stacks, the network outperformed the prior best method, a sliding-window convolutional network^[1]. The same network, trained on transmitted light microscopy images, won the ISBI cell tracking challenge 2015 in its categories by a large margin^[1]. Speed was part of the result as well: a 512x512 image is segmented in less than a second on a recent GPU^[1], which matters when segmentation is meant to be part of an interactive workflow rather than an overnight batch job.

Why it matters

The full implementation, based on Caffe, and the trained networks were made publicly available^[1]. Together with the small-data training strategy, that release made the method something other researchers could reproduce, retrain on their own microscopy images and compare against. The paper's contribution is therefore twofold: an architecture that pairs context with precise localization, and a training recipe that turns a modest annotated dataset into enough supervision for a deep network^[1].