Semantic-Aware Contrastive Learning for Multi-object Medical Image Segmentation
Ho Hin Lee, Yucheng Tang, Qi Yang, Xin Yu, Leon Y. Cai, Lucas W. Remedios, Shunxing Bao, Bennett A. Landman, Yuankai Huo
Paper: https://ieeexplore.ieee.org/document/10149329
Code: https://github.com/MASILab/DCC_CL
Abstract
Medical image segmentation, or computing voxel-wise semantic masks, is a fundamental yet challenging task in medical imaging domain. To increase the ability of encoder-decoder neural networks to perform this task across large clinical cohorts, contrastive learning provides an opportunity to stabilize model initialization and enhances downstream tasks performance without ground-truth voxel-wise labels. However, multiple target objects with different semantic meanings and contrast level may exist in a single image, which poses a problem for adapting traditional contrastive learning methods from prevalent “image-level classification” to “pixel-level segmentation”. In this paper, we propose a simple semantic-aware contrastive learning approach leveraging attention masks and image-wise labels to advance multi-object semantic segmentation. Briefly, we embed different semantic objects to different clusters rather than the traditional image-level embeddings. We evaluate our proposed method on a multi-organ medical image segmentation task with both in-house data and MICCAI Challenge 2015 BTCV datasets. Compared with current state-of-the-art training strategies, our proposed pipeline yields a substantial improvement of 5.53% and 6.09% on Dice score for both medical image segmentation cohorts respectively (p-value < 0.01). The performance of the proposed method is further assessed on external medical image cohort via MICCAI Challenge FLARE 2021 dataset, and achieves a substantial improvement from Dice 0.922 to 0.933 (p-value < 0.01). The code is available at: https://github.com/MASILab/DCC_CL
