U-Net: Convolutional Networks for Biomedical Image Segmentation

Olaf Ronneberger; Philipp Fischer; Thomas Brox

Papers · arXiv 1505.04597

U-Net: Convolutional Networks for Biomedical Image Segmentation

Olaf Ronneberger, Philipp Fischer, Thomas Brox

Published 2015-05-18cs.CV arXiv PDF

Read with translation and AI

Summary in three lines

🎯 Train high-precision pixel-wise biomedical segmentation networks end-to-end using very small annotated datasets.
🔧 A symmetric u-shaped convolutional architecture combining a contracting path for context, an expansive path with skip connections for precise localization, unpadded valid convolutions, elastic deformation augmentation, and a separation border weight map.
📈 Outperformed sliding-window methods on the ISBI EM segmentation challenge (warping error 0.000353) and won the ISBI Cell Tracking Challenge 2015 with 0.9203 IOU on PhC-U373 and 0.7756 IOU on DIC-HeLa.

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 .

Reading guide

Problem and motivation

Biomedical image segmentation requires dense, pixel-level classification (localization) rather than image-level categorisation. Two fundamental challenges dominate this domain:

Annotated training data is scarce, often consisting of only a few dozen images rather than thousands or millions.
Prior deep learning methods, such as patch-based sliding-window convolutional networks (e.g., Cireşan et al.), suffered from severe redundancy, slow inference, and an inherent trade-off between patch context and localization precision.

The authors sought an end-to-end fully convolutional network that can capture global context and fine local details simultaneously, process arbitrarily large images seamlessly, and train successfully from extremely limited data.

Key ideas

Symmetric U-Shaped Architecture: Extends fully convolutional networks by matching the contracting (downsampling) encoder with an equally expressive expansive (upsampling) decoder equipped with numerous feature channels to pass context upward.
Skip Connections: High-resolution feature maps from the contracting path are copied and concatenated with upsampled feature maps in the expansive path (after cropping), providing precise spatial localization details lost during pooling.
Overlap-Tile Strategy with Valid Convolutions: The network uses unpadded convolutions. Border regions are seamlessly predicted by tiling and extrapolating missing context through mirroring, bypassing GPU memory limits on large images.
Elastic Deformation Augmentation: Realistic, smooth elastic deformations generated via coarse random grids teach the network deformation invariance, drastically reducing the required number of training samples.
Separation Border Loss Weighting: A pre-computed pixel-wise weight map penalizes errors on background borders between touching cells of the same class, enabling distinct instance separation.

Method

Network Structure

The network has 23 convolutional layers organized into two paths:

Contracting Path: Repeated blocks of two $3 \times 3$ unpadded convolutions (each followed by a ReLU) and a $2 \times 2$ max pooling with stride 2. Feature channels double at each pooling step.
Expansive Path: Repeated blocks of an upsampling operation followed by a $2 \times 2$ convolution ("up-convolution") halving feature channels, concatenation with the cropped contracting feature map of matching resolution, and two $3 \times 3$ convolutions followed by ReLUs.
Final Layer: A $1 \times 1$ convolution maps the 64-channel feature representations directly to the desired $K$ class scores.

Loss Function and Border Weighting

The network is trained with pixel-wise softmax cross-entropy: $p_{k} (x) = \frac{exp ( a _{k} ( x ))}{\sum _{k^{'} = 1}^{K} exp ( a _{k^{'}} ( x ))}$ $E = x \in Ω \sum w (x) lo g (p_{ℓ (x)} (x))$ where $a_{k} (x)$ denotes the activation for feature channel $k$ at pixel $x$ , $ℓ (x)$ is the ground truth class label, and $w (x)$ is a spatial loss weight map.

To balance class frequencies and separate touching objects, $w (x)$ is defined as: $w (x) = w_{c} (x) + w_{0} \cdot exp (- \frac{( d _{1} ( x ) + d _{2} ( x ) ) ^{2}}{2 σ ^{2}})$ where $w_{c} (x)$ balances class frequencies, $d_{1} (x)$ is the distance to the border of the nearest cell, $d_{2} (x)$ is the distance to the border of the second nearest cell, and the authors set $w_{0} = 10$ and $σ \approx 5$ pixels.

Optimization

Trained using SGD in Caffe with batch size 1 (single large image tile) and high momentum (0.99).
Weights initialized from a Gaussian distribution with standard deviation $2/ N$ where $N$ is the number of incoming nodes per neuron.

Results

ISBI EM Segmentation Challenge (Neuronal Structures):
- Achieved the best rank with a warping error of $0.000353$ , Rand error of $0.0382$ , and pixel error of $0.0611$ (averaged over 7 rotated versions, no additional post-processing).
- Outperformed the previous state of the art by Cireşan et al. (warping error $0.000420$ , Rand error $0.0504$ ).
ISBI Cell Tracking Challenge 2015:
- PhC-U373 (Phase Contrast): Achieved an average Intersection over Union (IOU) of $0.9203$ , compared to the second-best algorithm's $0.83$ .
- DIC-HeLa (Differential Interference Contrast): Achieved an average IOU of $0.7756$ , compared to the second-best algorithm's $0.46$ .
Efficiency: Segmenting a $512 \times 512$ image requires less than a second on a modern GPU, with total training time taking approximately 10 hours on an NVIDIA Titan GPU (6 GB).

Limitations and open questions

Unpadded Convolutions and Cropping: Because convolutions are valid (unpadded), output tiles are strictly smaller than input tiles, requiring careful selection of input tile dimensions divisible by $2$ across all pooling steps and mirroring at image borders.
Instance Separation via Heuristic Weights: Separating touching instances relies on handcrafted morphological weight maps rather than end-to-end instance embeddings or detection branches.
2D-Specific Design: The paper focuses on 2D slices, leaving adaptation to volumetric 3D biomedical acquisitions for subsequent research.

Key terms

contracting path: The encoder half of U-Net that repeatedly applies unpadded 3x3 convolutions and max pooling to capture context and extract high-level feature representations.
expansive path: The symmetric decoder half of U-Net that upsamples feature maps, reduces channel dimensions, and reconstructs precise spatial segmentations.
up-convolution: A 2x2 convolution applied after upsampling that halves the number of feature channels before concatenation.
overlap-tile strategy: An inference and training approach that divides arbitrarily large images into overlapping tiles, mirroring border regions to supply required context for unpadded convolutions.
elastic deformations: A data augmentation technique generating smooth displacement vectors on a coarse grid via bicubic interpolation to simulate biological tissue variations.
weight map: A spatially varying loss weighting function combining class-frequency balancing and an exponential penalty based on distances to the two nearest cells to separate touching objects.
warping error: A segmentation evaluation metric used in the ISBI EM challenge that assesses topological correctness, such as unbroken neuronal membranes.
fully convolutional network: A neural network composed entirely of convolutional and pooling operations without fully connected layers, allowing dense prediction over variable input sizes.

Questions

Why does U-Net crop feature maps before concatenation in the skip connections?

Because U-Net uses valid (unpadded) convolutions, each convolution reduces the feature map size by border pixels. Cropping the contracting feature map ensures its spatial dimensions match the upsampled feature map for concatenation.

How does U-Net handle the borders of large images during inference?

U-Net applies an overlap-tile strategy: to predict the valid inner region of a tile, surrounding contextual pixels are provided as input, and any missing data outside image boundaries is extrapolated by mirroring.

How does U-Net separate cells that are touching each other?

It uses a pre-computed weight map in the loss function that adds a large exponential penalty to background pixels located between the nearest and second-nearest cell boundaries, forcing the network to maintain thin separation borders.

Why is batch size set to 1 during training?

The authors chose large image tiles over large batch sizes to maximize context within available GPU memory (6 GB NVIDIA Titan), compensating for the single-tile batch size by using a high SGD momentum of 0.99.

The overview on this page was written by AI from the paper and may contain mistakes. Read the paper for the details.

Read with translation and AI