Supervised Contrastive Embedding for Medical Image Segmentation

Sangwoo Lee,Geongyu Lee,Sangheum Hwang,Yejin Lee

doi:10.1109/access.2021.3118694

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.

Highlights

R ECENTLY, learning methods based on deep neural networks are rapidly changing the field of medical imaging analysis and are being applied for clinical purposes, e.g., early detection or classification of lesions to improve current practices
If a model is ideally trained to meet the original purpose of the encoder and decoder, the feature vectors extracted from the encoder should have the semantic embedding property, i.e., they should be distinguishable according to their classes in the corresponding receptive fields
We demonstrate that our proposed method is effective in improving the domain robustness of segmentation models when trained on multi-source domains

Summary

Introduction

R ECENTLY, learning methods based on deep neural networks are rapidly changing the field of medical imaging analysis and are being applied for clinical purposes, e.g., early detection or classification of lesions to improve current practices. Deep neural networks learned from a large amount of data support precise diagnoses and accelerate time-consuming processes requiring medical expertise [1]– [3]. Semantic segmentation in medical imaging analysis is a important area and is essential for diagnosis, monitoring, and treatment [4]. Despite such significant progress, better and reliable performance is required for segmentation models to be widely used in clinical settings. The feature vectors representing the same class of regions should be located close to each other and the feature vectors representing different classes should be

Objectives

Methods

Results

Conclusion