Semantic-Aware Contrastive Learning for Multi-object Medical Image Segmentation

Contrastive Learning:

Self-supervised representation learning approaches have recently been proposed to learn useful representation from unlabeled data. Some approaches propose learning embeddings directly in lower-dimensional representation spaces instead of computing a pixel-wise predictive loss [13]. Self-distillation with a teacher-student network is further proposed to enrich the semantic correspondence using pseudo-label predictions [14], while the masked autoencoder provides an alternative to learn the spatial feature correspondence with the image reconstruction task [15]. Contrastive learning is one of such state-of-the-art methods for self-supervised learning to model the semantic-wise relationships in the latent space [1], [16]. It employs a loss function to pull latent representations closer together for positive pairs, while pushing them apart for negative pairs. Maximizing mutual information between embeddings has also been proposed as an alternative to extract the similar information between targets [17|. Adapting with memory bank and momentum contrastive approaches have been proposed to increase the batch sizes and generate more dissimilar pairs in a minibatch for contrastive learning [18]. Additionally, to constrain and stabilize the embedding spaces, class label information has been added to provide additional supervision to stabilize contrastive learning process [2].

However, most of the prior works in contrastive learning focused on improving the “image-wise classification”, while relatively fewer methods have been proposed for the “pixel-wise segmentation”. Pixel-wise contrastive loss is proposed to adapt the representation from the ground-truth label information [9], while dense contrastive loss is also proposed to minimize the discrepancy of image-level prediction and pixel-level prediction [19]. Furthermore, one-stage contrastive learning framework is proposed to enforce the pixel embeddings belonging to a same semantic class to be more similar than embeddings from different classes [20]–[22]. In the medical domain, the contrastive learning framework is extended to leverage the structural similarity and learn the distinctive representations of local regions without using pixel-wise ground truth labels [7] and image-wise labels [23]. Similarly, limited ground-truth labels are adapted into the pretraining step for contrastive learning and enhanced the segmentation performance [24], [25]. However, such methods typically need pixel-wise segmentation labels. Therefore, our proposed method identifies the semantic-aware regions with one-hot attentional guidance and leverages multi-class image-level labels to define region-bounded representations as an arbitrary number of positive pairs for contrastive learning, without using pixel-wise labels.

Medical Image Segmentation & Multi-Organ Segmentation:

Fully-supervised deep learning methods have been developed to enhance both the segmentation performance and the generalizability across different datasets [26]–[29]. However, the supervised learning strategies are limited to the quality of pixel/voxel-wise labels and the resolution of volumes [30]. Thus, hierarchical approaches and patch-wise approaches have been proposed to perform segmentation across scales and resolutions [31]-[33]. Another study further enhances the segmentation accuracy with the statistical fusion from multi-view predictions [34]. Apart from the multi-view attention, shape-aware network is proposed to consistently smoothen the label prediction by learning the signed distance function as additional constraints [35]. Furthermore, RAP-Net is proposed to leverage one-hot shape-aware mappings to provide additional localization context as additional input channel and refine the segmentation mapping hierarchically [36]. nn-UNet further enhances the generalizability with self-configuring structure to diversely predict segmentation for multi-modality imaging [37]. In terms of generic network backbone, vision transformer is introduced as the encoder network to extract attention features with large receptive field for robust segmentation [38]. On the other hand, partially-supervised, semi-supervised and self-supervised learning have also been explored to adapt unlabeled data in the medical imaging domain. Multiple single organ-labeled datasets are used to provide structural prior knowledge during the training process with multiple organ-labeled dataset to enhance multi-organ segmentation performance [39]. A quality assurance module have been proposed to adapt the segmentation quality as the supervision from unlabeled data [40]. Pretext tasks such as colorization, deformation and image rotation, have been used as pre-training features to initialize the segmentation networks [41]. Self-supervised context has also been explored by predicting the relative patch location and the degree of rotation [42], [43]. Contrastive learning has been used to extract global and local representations for domain-specific MRI images [7]. A contrastive predictive network has been used to summarize the latent vectors in a minibatch and predicts the latent representation of adjacent patches [6].

A. Hierarchical Coarse Segmentation

The input data of the entire pipeline is a multi-contrast 3D image volume $V_i={\left\{X_i,Y_i\right\}}_{i=1,\dots ,L}$, where $L$ is the number of all imaging samples, $X$ is the volumetric image and $Y$ is the corresponding multi-organ label. The corresponding outcomes of the preprocessing stage are coarse segmentation masks (attention maps) $A_i=RAP\left(X_i\right)$ from a hierarchical segmentation network $RAP\left(\cdot \right)$ [36]. We define $A\in R^{H\times W\times D\times C}$, where $H$ and $W$ denote as the axial dimension of the image, $D$ denote as the number of slices and $C$ denotes as the number of label classes. The coarse segmentation network RAP-Net consists of two hierarchical stages: 1) low-resolution whole volume segmentation and 2) organ-specific patch segmentation refinement. The low-resolution model generates a rough segmentation map and provide anatomical context as additional channel input to refine the segmentation in patch-wise setting as the second step. Both low-resolution model and patch-wise model are trained in supervised setting with 5-fold cross-validations.

B. Data Preprocessing

The goal of the data preprocessing step is to randomly sample 2D training patches for downstream contrastive learning. In our design, we first slice all the volumetric scans (image, ground truth labels and coarse segmentations) and utilize the slice-wise attention maps (1) as spatial restrictions of the organ-specific sampling process, and (2) highlight the current organ of interest to define semantic-wise embeddings for segmentation refinement. Briefly, organ-specific patches $p_i={\left\{x_{C,i},y_{C,i},s_{C,i}\right\}}_{i=1,\dots ,N}$ are randomly sampled within each organ class $C$ in attention maps. The center point is randomly sampled from attention maps to crop the region of interest (ROI), $N$ denotes as the total number of query patches, $x_{C,i}$ is the organ-corresponding image patch, $y_{C,i}$ is the binary ground-truth label patch, and $s_{C,i}$ is the coarse organ-specific attention map from $A_i$ slice in binary setting. As the significant difference between $y_{C,i}$ and $s_{C,i}$ is the variation of segmentation quality, the trained model aims to refine the segmentation with the prior knowledge of $s_{C,i}$ as an additional input channel. As a standard process in data augmentation, random cropping, rotation (−30 to 30 degrees), scaling has been applied to augment the size of training samples.

C. Contrastive learning with Organ-Specific Attention

After generating augmented image pairs, pairwise images are then used as the inputs for contrastive learning. Specifically, a convolutional encoder network $E\left(\cdot \right)$ is used to extract high dimensional features. We further project each high-level feature mapping into 1D vector ${\tilde{z}}_i$ using a multi-layer perceptron network $P\left(\cdot \right)$, ${\tilde{z}}_i=P\left(E\left({\tilde{a}}_i\right)\right)$, ${\tilde{z}}_i\in R^{O_E}$ (pink box in Fig. 2), where $O_E$ is the size of the output vector. Then, the standard self-supervised contrastive loss (SSCL) [1] can be defined as the following:

where $T$ is a hyperparameter indicating temperature scaling to control the radius weighting on the positive pair/negative pairs. Both $k$ and $p\left(k\right)$ represent the index of the anchor sample and the corresponding positive sample, respectively. $J\left(k\right)$ represents the number of remaining negative samples. To incorporate the attention with modality and organ semantic meanings, we extend SSCL to adapt an arbitrary number of positive pairs by introducing multi-class image-level labels into contrastive loss. Here, modalities indicate the different contrast types in CT (we have utilized both contrast-enhanced and non-contrast CT scans in the training set). As the organ-specific attention only provides one-hot voxel-wise context to preserve organ regions, the multi-class labels represent different modalities and organs for the learned representations under the attention regions. It provides flexibility to further constrain the representations into semantic-aware clusters, which is conditional to multiple image-level labels. In each batch, pairwise patches with the same organ and modality label are defined as positive pairs, while the remaining pairs are specified as negative pairs. With such positive-negative pairs definition, we further extend the contrastive loss with conditional constraints as follows:

where $L\left(k\right)\equiv \left\{l\in J\left(k\right):m_k=m_l,o_p=o_l\right\}$, $m$ and $o$ denote as the corresponding modality and organ label, respectively. $l$ and ${\widehat{z}}_l$ are the index number and the projected feature representation of the corresponding positive sample with same organ and modality label. The feature vector output with 256 channels is directly used to compute the contrastive loss of modality and organ class respectively. By classifying the learned representations into multi-classes embeddings, the model is initially learned the attention-bounded representations with semantic meanings, which are hypothesized to be beneficial for downstream segmentation tasks.

D. Co-training with Multi-Organ Segmentation

The ultimate goal of our framework is to achieve a robust patch-wise contrastive learning without using pixel-wise labels, which benefits for downstream segmentation tasks. The native two-stage strategy is to train both contrastive loss and downstream segmentation loss independently. Here, we attempt to have a co-training strategy, by training the contrastive loss and segmentation refinement tasks simultaneously. The encoder network is followed with an atrous spatial pyramid pooling (ASPP) module as the decoder network to resample the bottleneck feature with multiple effect Field Of Views (FOVs) [44]. The DeepLabV3+ is employed as the segmentation part with the shared encoder structure for contrastive learning [44]. The distinctiveness of adapting ASPP is to obtain multi-scale features during upsampling. One $1\times 1$ convolution and three $3\times 3$ convolution layers with different dilation rate (e.g., 6, 12,18) are leveraged. With the increased number of dilation rate, kernel stride is constrained while a larger FOV is accomplished without increasing the number of model parameters. Furthermore, image pooling is also performed in parallel to extract the global features. Features from different FOVs are finally concatenated. The channel-wise features are mixed using a $1\times 1$ convolution layer before passing through the final layer for high-resolution prediction. The rationale of such design is to adapt the multi-view behavior and search the optimal tradeoff between the localized features (small FOV) and the global-assimilated features (large FOV). Dice loss is used to compute the predicted output with the ground truth label in binary setting for the co-training segmentation task. After computing all organ-specific patches predictions, we fuse the organ-wise patches according to the center point recorded in the data preprocessing stage. We employ majority voting [45] as the label fusion module to fuse predictions into multi-organ labels.

Datasets:

To evaluate our proposed learning approach, one in-house research cohort and two publicly available cohorts in medical imaging are used with multi-organ segmentation as the downstream task.

MICCAI 2015 Challenge Beyond The Cranial Vault (BTCV) dataset is comprised of 100 de-identified unpaired 3D contrast-enhanced CT scans with 7,968 axial slices in total. 20 scans are publicly available for the testing phase in the MICCAI 2015 BTCV challenge. All CT scans are in portal venous phase. Peak enhancement of contrast is observed in several organs, such as liver, kidney, spleen, and portal splenic vein. For each scan, 12 organ anatomical structures are well-annotated, including spleen, right kidney, left kidney, gallbladder, esophagus, liver, stomach, aorta, inferior vena cava (IVC), portal splenic vein (PSV), pancreas and right adrenal gland. Each volume consists of $47\sim 133$ slices of $512\times 512$ pixels, with the resolution of $\left(\left[0.54\sim 0.98\right]\times \left[0.54\sim 0.98\right]\times \left[2.5\sim 7.0\right]\right)m{m}^3$.

Non-contrast clinical cohort is retrieved in de-identified form from ImageVU database of Vanderbilt University Medical Center. It consists of 56 unpaired 3D CT scans with 3,687 axial slices and expert-refined annotations for the same 12 organs as the MICCAI 2015 BTCV challenge dataset. All volumetric scans are generated without contrast enhancement procedures. Each volume consists of $49\sim 174$ slices of $512\times 512$ pixels, with the resolution of $\left(\left[0.64\sim 0.98\right]\times \left[0.64\sim 0.98\right]\times \left[1.5\sim 5.0\right]\right)m{m}^3$.

MICCAI 2021 Challenge Fast and Low GPU memory Abdominal Organ Segmentation (FLARE) dataset leverage large scales of abdominal contrast-enhanced CT with 511 unpaired cases from 11 medical centers in multi-contrast phases (including both portal venous phase and non-contrast phase CTs). It consists of 361 3D CT scans with four organ-specific labels including spleen, kidney, liver and pancreas. Each volume consists of $43\sim 384$ slices of $512\times 512$ pixels, with the resolution of $\left(\left[0.64\sim 0.98\right]\times \left[0.64\sim 0.98\right]\times \left[1.0\sim 5.0\right]\right)m{m}^3$.

Preprocessing:

We apply the preprocessing steps as follows: (i) applying soft tissue windowing within the range of −175 to 250 Hu and performing intensity normalization of each 3D volume, $v$ with min-max normalization: $\left(v-v_1\right)/\left(v_{99}-v_1\right)$, where $v_p$ denote as the $p$^th intensity percentile in $v$, and (ii) applying volume-wise cropping in z-axis with body part regression algorithm to extract the abdominal region only for segmentation and ensure the similar field of view between scans [48].

Network Training:

Our proposed framework AGCL is trained with unpaired samples in our scenario, while it also allows to train with paired samples. 5 -fold cross-validation is performed for both contrast-enhanced phase and non-contrast phase CT (Training: 60 volumes (contrast-enhanced) and 44 volumes (non-contrast), validation: 20 volumes (contrast-enhanced) and 6 volumes (non-contrast), and testing: 20 volumes (contrast-enhanced) and 6 volumes (non-contrast)). For training the coarse segmentation network RAP-Net, we downsample all training volumes to a resolution of $2\times 2\times 6$ with the dimension of $168\times 168\times 64$. The low-resolution volumes are leveraged to train a low-resolution segmentation model with Adam optimizer using a batch size of 1 and a learning rate of $1e-4$. We then use the coarse segmentation output to guide and extract organ-specific patches with the dimension of $128\times 128\times 48$. The patch-wise segmentation refinement model is trained with Adam optimizer using a batch size of 2 and a learning rate of $1e-4$. For contrastive learning, we perform patients-level sampling and extract 30 2D query patches of each anatomical target in each axial slices of a subject scan. Such sampling strategy ensures that all patches are fully covered the organ-specific ROIs with significant variation of anatomical morphology. More than 400k patches with dimensions $128\times 128$ are used and shuffle to train with stochastic gradient descent (SGD) optimizer for 5 epochs with a batch size of 4 and a learning rate of $5\times {10}^{-4}$. We have evaluated the variation of the temperature parameter towards the segmentation performance and $T=0.1$ achieves the best performances across all other temperature values. For segmentation task, the encoder’s weight is frozen and only the decoder with ASPP module is trained for 10 epochs with Adam optimizer using a batch size of 4 and a learning rate of ${10}^{-4}$. We use the validation set to choose the model with the highest mean Dice score for all semantic targets segmentation and perform inference as the quantitative representation on the testing set.

Experimental Setup:

We evaluate the segmentation performance with Dice similarity coefficient on current state-of-the-art approach in contrastive learning and segmentation task for medical imaging domain, including the testing phase of the BTCV dataset, testing cohort of the non-contrast clinical cohort and the random sampled cohort from FLARE dataset. We further perform different pre-training strategies with the multi-class image-level label using different scenarios. Apart from learning image-wise embeddings with self-supervised setting, inspired by Khosla et al. [2], we introduce patch-wise multi-label (modality & organ) classification as the pretext task via the canonical cross-entropy (CE) loss in a fully-supervised setting. The learned representations are classified into label-corresponding clusters, as in AGCL. The CE loss is defined as following:

where $y_i$ is the ground-truth multi-class label and $p_i$ is the Softmax probability for the $i$^th, $i\in 1,\dots ,C$ classes. Apart from different pretext task strategies, we also perform ablation studies with the variation of hyperparameters such as temperature and the number of label used for pretraining, to investigate the optimal effect on fine-tuning segmentation task. For the encoder network, we evaluate with two common backbone architectures for segmentation in medical imaging domain: DeeplabV3+ with ResNet-50 and ResNet-101 encoder. The normalized activation of the final pooling layer with $D_E=2048$ are used as the distinctive feature representation vector. More details of both network training and data preprocessing are demonstrated in the supplementary material.