Supervised Contrastive Embedding for Medical Image Segmentation

Abstract

Deep segmentation networks generally consist of an encoder to extract features from an input image and a decoder to restore them to the original input size to produce segmentation results. In an ideal setting, the trained encoder should possess the semantic embedding capability, which maps a pair of features close to each other when they belong to the same class, and maps them distantly if they correspond to different classes. Recent deep segmentation networks do not directly deal with the embedding behavior of the encoder. Accordingly, we cannot expect that the features embedded by the encoder will have the semantic embedding property. If the model can be trained to have the embedding ability, it will further enhance the performance as restoring from those features is much easier for the decoder. To this end, we propose supervised contrastive embedding, which employs feature-wise contrastive loss for the feature map to enhance the segmentation performance on medical images. We also introduce a boundary-aware sampling strategy, which focuses on the features corresponding to image patches located at the boundary area of the ground-truth annotations. Through extensive experiments on lung segmentation in chest radiographs, liver segmentation in computed tomography, and brain tumor and spinal cord gray matter segmentation in magnetic resonance images, it is demonstrated that the proposed method helps to improve the segmentation performance of popular U-Net, U-Net++, and DeepLabV3+ architectures. Furthermore, it is confirmed that the robustness on domain shifts can be enhanced for segmentation models by the proposed contrastive embedding.