Contextual Embedding Learning to Enhance 2D Networks for Volumetric Image Segmentation
This work addresses a domain-specific problem in medical imaging by providing a plug-and-play, memory-efficient solution for volumetric segmentation, though it is incremental as it builds on existing 2D and 3D CNN approaches.
The paper tackles the problem of volumetric medical image segmentation by enhancing 2D CNNs with contextual information to overcome memory and computation limitations of 3D CNNs, resulting in improved segmentation performance on prostate MRI and abdominal CT datasets.
The segmentation of organs in volumetric medical images plays an important role in computer-aided diagnosis and treatment/surgery planning. Conventional 2D convolutional neural networks (CNNs) can hardly exploit the spatial correlation of volumetric data. Current 3D CNNs have the advantage to extract more powerful volumetric representations but they usually suffer from occupying excessive memory and computation nevertheless. In this study we aim to enhance the 2D networks with contextual information for better volumetric image segmentation. Accordingly, we propose a contextual embedding learning approach to facilitate 2D CNNs capturing spatial information properly. Our approach leverages the learned embedding and the slice-wisely neighboring matching as a soft cue to guide the network. In such a way, the contextual information can be transferred slice-by-slice thus boosting the volumetric representation of the network. Experiments on challenging prostate MRI dataset (PROMISE12) and abdominal CT dataset (CHAOS) show that our contextual embedding learning can effectively leverage the inter-slice context and improve segmentation performance. The proposed approach is a plug-and-play, and memory-efficient solution to enhance the 2D networks for volumetric segmentation. Our code is publicly available at https://github.com/JuliusWang-7/CE_Block.