Class-aware Patch Based Contrastive Learning for Medical Image Segmentation
Shaozhi Wu, Han Li, Xingang Liu, Dan Tian, Han Su · 2024
For medical image segmentation, contrastive learning recently becomes the dominant method to get the class-separated visual representations by pulling augmented views of the same samples closer in a representation space, and pushing apart augmented views of different samples. However, most current contrastive learning methods may suffer from the issue of class collision because those images (pixels) that from the same class are forcefully contrasted due to the pretext task simply takes another augmentation of the same image as positive pairs and all other images are view as negative samples. Moreover, the problem looks more serious in pixel-level contrastive learning in that massive pixels from the same class are viewed as negative pairs to differ is intolerable. Therefore, to alleviate this, we propose a novel semi-supervised learning framework based on hierarchical data augmentation in contrastive learning for medical image segmentation. Specifically, We first propose a novel contrastive loss function based on class-aware patch by using the pseudo labels and source image effectively. With the novel loss, we learn better intra-class compactness and inter-class separability in the feature space. We then introduce a patch-based hierarchical data augmentation module to learn the hierarchical invariance by applying different strength of augment to different probability of patches, thus making the patches with a high probability of belong to a certain class more alike and avoiding the issue of class collision in contrastive learning. Experiments on several public accessible datasets from multiple domains reveals the superiority of our proposed method as compared to the state-of-the-art semi-supervised and contrastive learning methods.