[Paper Review] 3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation

Özgün Çiçek, Ahmed Abdulkadir+3 more|arXiv (Cornell University)|Jun 21, 2016

Advanced Neural Network Applications12 references491 citations

TL;DR

The paper extends U-Net to 3D to densely segment volumetric biomedical images using sparsely annotated slices, enabling semi-automated densification and fully automated segmentation with end-to-end training from scratch.

ABSTRACT

This paper introduces a network for volumetric segmentation that learns from sparsely annotated volumetric images. We outline two attractive use cases of this method: (1) In a semi-automated setup, the user annotates some slices in the volume to be segmented. The network learns from these sparse annotations and provides a dense 3D segmentation. (2) In a fully-automated setup, we assume that a representative, sparsely annotated training set exists. Trained on this data set, the network densely segments new volumetric images. The proposed network extends the previous u-net architecture from Ronneberger et al. by replacing all 2D operations with their 3D counterparts. The implementation performs on-the-fly elastic deformations for efficient data augmentation during training. It is trained end-to-end from scratch, i.e., no pre-trained network is required. We test the performance of the proposed method on a complex, highly variable 3D structure, the Xenopus kidney, and achieve good results for both use cases.

Motivation & Objective

• Enable dense 3D segmentation from sparsely annotated volumetric images.
• Provide a semi-automated workflow where user-specified slices guide full-volume segmentation.
• Enable fully automated segmentation with a representative sparsely labeled training set.
• Demonstrate the approach on a complex Xenopus kidney structure and quantify gains over 2D approaches.

Proposed method

• Extend the 2D U-Net architecture to 3D by replacing all operations with 3D equivalents (3D convolutions, 3D max-pooling, 3D up-convolutions).
• Use no bottleneck architecture and apply batch normalization to accelerate training.
• Train end-to-end from scratch with a weighted softmax loss that ignores unlabeled voxels to enable learning from sparse annotations.
• Apply on-the-fly elastic deformations for data augmentation to increase generalization from few labeled slices.
• Employ a tiling strategy to handle large volumes and enable training with 3D data sizes.
• Provide an OpenSource Caffe-based implementation.

Experimental results

Research questions

• RQ1 Can a 3D fully convolutional network learn dense volumetric segmentations from sparsely annotated slices?
• RQ2 How does 3D U-Net compare to a 2D approach when trained on sparse annotations for semi-automated densification?
• RQ3 What is the impact of the number of annotated slices on segmentation performance in semi-automated and fully automated setups?
• RQ4 Does batch normalization improve performance in this low-sample, high-dimensional setting?
• RQ5 How does the method perform on a highly variable 3D biological structure (Xenopus kidney) and generalize to new volumes?

Key findings

• In semi-automated segmentation, the 3D U-Net with BN achieves IoU up to 0.863 (mean across folds) outperforming a 2D baseline.
• Increasing the number of annotated slices improves IoU nonlinearly, with substantial gains even from very few slices.
• In fully automated segmentation, the 3D model generally outperforms the 2D counterpart when trained with BN, though BN can be ineffective in some highly heterogeneous data settings.
• The 3D architecture provides a quantitative performance gain over a 2D implementation in semi-automated and fully automated scenarios.
• The approach can be trained from scratch and does not require pre-trained networks.