SamDSK: Combining Segment Anything Model with Domain-Specific Knowledge for Semi-Supervised Learning in Medical Image Segmentation

Our proposed SamDSK uses probability maps generated by a medical image segmentation network for unlabeled images. Although probability maps are not exactly taken as labels or pseudo-labels, these maps are often referred to as “soft” pseudo-labels. Hence, below we give a review of the methods that use pseudo-labels for semi-supervised learning in medical image segmentation.

Wu et al. [27] proposed Mutual Consistency Learning (MCL) for utilizing unlabeled samples in segmentation model training. MCL creates multiple decoders which all share one encoder for generating segmentations. Because of the random initialization, these decoders generate possibly different segmentation predictions, and the differences between these predictions are used to compute prediction uncertainties. In addition, mutual learning [30] is performed during model training to enforce the output from each decoder to be close (consistent) to the pseudo-labels generated by the other decoders. At the end of the training, only one decoder is used for model deployment. Since pseudo-labels inherently contain errors, Thompson et al. [23] proposed to use super-pixels to refine pseudo-labels in semi-supervised learning. The proposed process largely follows the classical approach of using pseudo-labels for SSL [14], with additional steps which refine pseudo-labels according to super-pixels. The main motivation of this method is to impose a spatial structure induced from the super-pixels for regularizing and refining pseudo-labels so that pseudo-labels would yield better quality and benefit the subsequent segmentation model training. Seibold et al. [20] proposed to use reference images (with annotations) for generating pseudo-labels. This approach works for medical image segmentation tasks with well-defined and stable object structures (e.g., chest radiographic anatomy segmentation). For tasks under more dynamic scenes where objects can be at arbitrary locations and with large variations of shapes (e.g., polyp segmentation in endoscopic images), reference-based pseudo-label generation could become less reliable and effective. Recently, Basak et al. [4] proposed to utilize pseudo-labels in conjunction with contrastive learning.

Our proposed method utilizes probability maps generated by a task-specific medical image segmentation model. Similar to [23], we employ an external component (i.e., SAM in our case) to work with the medical image segmentation model. Instead of refining pseudo-labels using super-pixels, we aim to use DSK to match and select SAM-generated segmentation proposals and construct segmentations for unlabeled images for model re-training. Comparing to the super-pixel based methods, SAM is a much more advanced model and using SAM allows us to design a more controllable optimization process (e.g, multiple rounds of optimal matching) in the usage of domain-specific knowledge both at the pixel level and image level for generating segmentations of unlabeled images.

Since the introduction of the Segment Anything Model (SAM), there has been a wave of attempts to develop a more accurate and robust medical image segmentation system with the utilization of SAM [11, 32, 16, 26, 28, 18]. Recent work has shown that SAM alone, without further fine-tuning and/or adaptation, often delivers unsatisfactory results for medical image segmentation tasks [11, 32]. To utilize SAM more effectively, Ma et al. [16] proposed to fine-tune SAM using labeled images. Wu et al. [26] proposed to add additional layers to adapt SAM for a medical image segmentation task. Our work aims at using SAM for semi-supervised learning of a medical image segmentation model.

Recently, Chen et al. [7] proposed to use SAM to refine pseudo-labels in weakly-supervised learning. SAM is employed in three places: (1) refining CAM-generated attention maps, (2) refining segmentation generated by post-processing, and (3) refining segmentations generated by a segmentation model (DeepLab). In this paper, we utilize probability maps as pixel-level DSK for selecting segmentation proposals produced by SAM. To make the selection more reliable, we further employ an image-level DSK (e.g., the potential number of regions of interests) during the selection process to constrain the optimization process as a way of realizing the idea of combining SAM with domain-specific knowledge (at both the pixel level and image level).

Given a set $U$ of unlabeled medical images, we apply the state-of-the-art SAM to these images to obtain segmentation proposals in each image. A key consideration in this step is to use relatively lower threshold settings in SAM to ensure the inclusion of all potential Regions of Interest (RoIs) in the generated segmentation proposals. By default, we use the following parameter values when employing SAM: “crop_nms_threshold = 0.5, box_nms_thresh = 0.5, pred_iou_thresh = 0.5, stability_score_thresh = 0.5”. With the ongoing development of the segmentation foundation model, future SAM variants are expected to provide improved coverage of RoIs for various medical image segmentation tasks.

Suppose we already train a segmentation model ${M}$ using the currently available labeled data. Also, suppose for an unlabeled image $u$ with width $w$ and height $h$, SAM provides a set of $K$ segmentation proposals, $S=\{s_{1},s_{2},\dots,s_{K}\}$, where $s_{i}\in\{0,1\}^{h\times w}$ for $i=1,2,\dots,K$. The segmentation model ${M}$ gives $C$ distinct probability maps $\{p_{1},p_{2},\dots,p_{C}\}$, where $p_{c}\in\mathbb{R}^{h\times w}$ for $c=1,2,\dots,C$. $C$ is the number of classes in the segmentation task. The background class is taken always as the class $C$. For each class $c\in\{1,2,\dots,C\}$, we create a set of binary scalar indicators, $Z_{c}=\{z_{c,1},z_{c,2},\dots,z_{c,K}\}$, where $z_{c,k}\in\{0,1\}$ for $k=1,2,\dots,K$. Each indicator $z_{c,k}$ is multiplied with the corresponding segmentation proposal $s_{k}$ in constructing the segmentation map for class $c$. We aim to find an optimal configuration of $z_{1,1},z_{1,2},\dots,z_{C,K}$ which maximizes the total sum of IoU scores between the segmentation proposals and the segmentation probability maps, as:

$$z_{1,1}^{*},\dots,z_{C,K}^{*}=\operatorname*{argmax}_{z_{1,1},z_{1,2},\dots,z_{C,K}}\sum_{c=1}^{C-1}\textrm{IoU}(p_{c},\tau(\sum_{k=1}^{K}z_{c,k}s_{k})),$$

(1)

	$\displaystyle s.t.,$	$\displaystyle\sum_{c=1}^{C-1}z_{c,k}\leq 1,$		(2)
		$\displaystyle z_{c,k}\in\{0,1\},\ c=1,2,\dots,C,\ k=1,2,\dots,K,$

In Fig. 3, we provide a high-level illustration of the mechanism that we use for this optimal matching formulation. The optimization problem in Eq. (1) (with the constraints in Eq. (2)) is a maximum bipartite matching problem which can be efficiently solved in polynomial time [3].

In this section, we aim to improve the effectiveness and reliability of the above proposed segmentation annotation construction by adding image-level DSK as new constraints to the optimal matching formulation. Suppose the number of RoIs for an image of a given medical image segmentation task ranges from $v_{c}^{lower}$ to $v_{c}^{upper}$ for class $c$. We add these constraints to Eq. (2). The optimization objective remains unchanged, but the constraints are updated with additional constraints imposed by $v_{c}^{lower}$ and $v_{c}^{upper}$. Formally, we aim to optimize the objective in Eq. (1), with the following constraints:

	$\displaystyle s.t.,$	$\displaystyle\sum_{c=1}^{C-1}z_{c,k}\leq 1,$		(3)
		$\displaystyle v_{c}^{lower}\leq\sum_{k=1}^{K}z_{c,k}\leq v_{c}^{upper},c\in\{1,2,\dots,C-1\},$
		$\displaystyle z_{c,k}\in\{0,1\},\ c=1,2,\dots,C,\ k=1,2,\dots,K.$

In each round of the model training and labeled set expansion, we control the $v_{c}^{lower}$ and $v_{c}^{upper}$ values to ensure a stable way of gradually including more samples with more RoIs into the labeled set. For the first round, we can set $v^{lower}$ and $v^{upper}$ as both equal to 1. This indicates that the optimal matching is only a one-to-one matching between the segmentation proposals and the probability maps. In the second round, we can set $v_{c}^{lower}$ as 1 and $v_{c}^{upper}$ as 2 to allow a two-to-one matching. Depending on the specific segmentation task, we can use a larger increment of $v_{c}^{upper}$ to allow a more effective inclusion of unlabeled samples into the labeled set.

With the identified values $z_{1,1}^{*},z_{1,2}^{*},\dots,z_{C,K}^{*}$ through optimizing Eq. (1) and Eq. (3), we compute the overall score $\beta$ for the image, as follows:

$\displaystyle\beta=\frac{1}{C-1}\sum_{c=1}^{C-1}\textrm{IoU}(p_{c},\tau(\sum_{k=1}^{K}z_{c,k}^{*}s_{k})).$

(4)

A sample with a $\beta$ score higher than a pre-defined threshold (denoted as $\beta^{*}$; its default value is 0.9) is selected to advance to the next round of segmentation model training. For a selected sample, its annotation map is constructed as $q_{c}=\tau(\sum_{k=1}^{K}z_{c,k}^{*}s_{k})$, for $c=1,2,\dots,C-1$. Pixel areas not covered by any segmentation proposals selected by the optimal matching are considered as belonging to the background class (class $C$).

Suppose the unlabeled samples with their annotation maps generated by the above optimization procedure are denoted as ${D^{machine}}=\{(u_{j},q_{j,c})\}$, for $j=1,2,\dots,m^{\prime}$, $m^{\prime}\leq m$, and $c=1,2,\dots,C$. The original human-labeled set is denoted as $D=\{(x_{i},y_{i,c})\}$, for $i=1,2,\dots,n$, $c=1,2,\dots,C$. We aim to optimize the following objective function with respect to the parameters $\theta$ of a medical image segmentation model ${M}$, as:

$$\operatorname*{argmin}_{\theta}(\sum_{c=1}^{C}\sum_{i=1}^{n}{L}(M(x_{i},y_{i,c}))+\lambda\sum_{c=1}^{C}\sum_{j=1}^{m^{\prime}}{L}(M(u_{j},q_{j,c}))),$$

(5)

where $\lambda$ is set as 1 by default. Assuming that the medical image segmentation model ${M}$ generates probability maps for each class, the loss function $L$ can take the form of a region-based loss (e.g., Dice loss), a pixel-level entropy-based loss, or a combination of these two types of losses. A mini-batch based stochastic gradient descent method (e.g., Adam) can be applied to optimize the loss via updating the parameters $\theta$ of the model ${M}$.

The SamDSK method combines SAM-generated segmentation proposals with domain-specific knowledge both at the pixel-level and image-level in order to generate annotations for unlabeled images. Pixel-level DSK is provided by the segmentation model ${M}$ trained on the current labeled image set, and optimal matching is performed with constraints imposed by image-level DSK as well as the base constraints that only one segmentation proposal can be assigned to one class label. There exist clearly two cases after the optimal matching process for constructing annotations of the unlabeled images.

• •

  Case-1: Prediction maps produced by the current segmentation model ${M}$ match well with a subset of the segmentation proposals (generated by SAM). In this scenario, the optimal matching gives a high $\beta$ matching score (i.e., higher than $\beta^{*}$), and this image with the constructed annotations is added to the labeled image set for the subsequent rounds of segmentation model retraining.

• •

  Case-2: Prediction maps produced by the current segmentation model ${M}$ do not match well with any subsets of the segmentation proposals generated by SAM. In this situation, optimal matching gives a low $\beta$ matching score, and this image is not added to the labeled set for the next round of model retraining.

For simplicity, we use a binary segmentation task to illustrate, and focus on its segmentation class #1 (the foreground class) to provide the analyses below. The same logic applies to multi-class segmentation tasks with multiple classes of foreground objects.

Assumption 1: For an unlabeled image $u$, there exists a subset of segmentation proposals, $\Gamma_{1}$, generated by SAM, such that the union of its elements closely approximates the ground truth of segmentation class 1 in $u$.¹¹1SamDSK is not yet ready to be applied to those medical image segmentation tasks for which SAM fails to generate sufficiently good segmentation proposals. More formally, this assumption can be described as:

$$\textrm{IoU}(\bigcup_{s\in\Gamma_{1}}s,gt_{1})>1-\epsilon,$$

(6)

where $gt_{1}$ represents the ground truth annotation map of class 1 for the image $u$, and $\epsilon$ is a small positive value (e.g., $\epsilon$ = 0.02). It is important to note that we do not have access to the ground truth annotations of the unlabeled images.

Since the union of elements in $\Gamma_{1}$ closely approximates the ground truth and $p_{1}$ is not well aligned with the union of elements in $\Gamma_{1}$, it follows that $p_{1}$ is NOT closely aligned with the ground truth annotation map $gt_{1}$. Formally, an upper bound of the Intersection over Union (IoU) between the ground truth annotation map and the prediction map of class #1 is given by:

$$\textrm{IoU}(gt_{1},p_{1})\lesssim\beta^{*}+\epsilon.$$

(8)

With the above argument, we have demonstrated that the unlabeled samples which are not added to the labeled set (Case-2) do not yet contain sufficiently accurate annotations (e.g., $p_{1}$ is not sufficiently accurate with respect to the ground truth). On the other hand, for the unlabeled images with annotations that are added to the labeled set (Case-1), it is still possible that although their annotations match well with a subset of the segmentation proposals, denoted as $\Gamma_{1}$, the ground truth annotation map is actually closer to another subset $\Gamma_{1}^{\prime}$ of the segmentation proposals, and $\Gamma_{1}$ may or may not be equal to $\Gamma_{1}^{\prime}$. Consequently, we cannot provide an explicit guarantee on the correctness of $p_{1}$ and its corresponding $\bigcup_{s\in\Gamma_{1}}s$ for the samples in Case-1. Nevertheless, in our experiments, we present empirical evidence to demonstrate that the annotations of samples in Case-1 (those added to the labeled set) are closer to their corresponding ground truth annotations than those in Case-2.

Medical image segmentation encompasses a multitude of imaging modalities (e.g., MRI, CT, microscopy), covering a wide range of diverse objects of interest and clinical tasks. SamDSK relies on SAM to provide segmentation proposals. Therefore, SamDSK is especially applicable to those medical image segmentation tasks in which the regions of interest can be adequately addressed by segmentation proposals generated by SAM. It is important to note that our optimal matching formulation allows for the utilization of multiple segmentation proposals to cover one or more regions of interest. Consequently, SamDSK accommodates segmentation proposals that consist of over-segmented regions with respect to the regions of interest.

For illustration, we highlight several specific medical image segmentation tasks for which SamDSK is well-suited: (1) polyp segmentation in endoscopic images, (2) skin lesion segmentation in dermoscopic images, and (3) tumor region segmentation in ultrasound images. Fig. 4 presents visual examples of these tasks and the corresponding segmentation proposals generated by SAM. We conduct comprehensive experiments to demonstrate the effectiveness of SamDSK in tackling these segmentation tasks, as shown in the next section.

We first utilize the breast cancer segmentation task in ultrasound images [1] to validate the effectiveness of our proposed SamDSK. This task is a binary segmentation problem of segmenting breast cancer regions in ultrasound images. We randomly split the original dataset into two equally-sized sets, one set for model training and the other set for testing. In the training set, we randomly select 30% of samples and treat them as labeled images, and treat the remaining 70% of samples as unlabeled images.²²2All the data split information will be published together with the code release. Two state-of-the-art medical image segmentation models, TransUNet [6] and HSNet [29], are taken as the segmentation model $M$ for the experiments. We perform SamDSK for three rounds of processing (step-0 to step-4 as illustrated in Fig. 2 are one round of processing). For Eq. (3), both $v_{1}^{lower}$ and $v_{1}^{upper}$ are set as 1. In Table 1, the results show that SamDSK outperforms the state-of-the-art methods for SSL. SamDSK yields considerably better segmentation performance than the closely related BPC method, indicating that SAM plays a critical role in sample selection and annotation construction of unlabeled images. In addition, the results show that the PL method may degrade the segmentation performance (see the TransUNet case) when the initial segmentation model incurs too many errors in pseudo-labels.

Automatic polyp segmentation in endoscopic images can help improve the efficiency and accuracy of clinical screenings and tests for gastrointestinal diseases. Many deep learning (DL) based methods have been proposed for robust and automatic segmentation of polyps. Here, we utilize the SOTA polyp segmentation model HSNet [29] and the widely-used polyp segmentation model PraNet [10] for evaluating our proposed SamDSK method. HSNet uses the PVT backbone (Transformer-based) and PraNet uses the Res2Net backbone (CNN-based). We perform SamDSK for three rounds of processing. For Eq. (3), $v_{1}^{lower}$ is set as 1, 1, and 1 for the first, second, and third rounds, respectively, and $v_{1}^{upper}$ is set as 1, 2, and 3 for the first, second, and third rounds, respectively (the polyp class is set as class 1). That is, the constraints on the number of RoIs for the polyp class are gradually changed when we perform more rounds of model training and labeled set expansion, and we allow an increase in the number of segmentation proposals for the foreground-class segmentation. Following the data split settings in [10, 32], 900 images from Kvasir [12] and 550 images from CVC-ClinicDB [5] are randomly selected to form the training set. The remaining samples from these two datasets (i.e., Kvasir and CVC-ClinicDB) and the samples from CVC-ColonDB [22], ETIS [21], and CVC-300 [24] form the test set.

The results are reported in Table 2, which show that SamDSK improves the segmentation performances of the HSNet and PraNet models for the settings where 10% and 30% of the total labeled samples from the training set are available. Compared to the classical pseudo-label method [14] and the recently proposed methods, SamDSK achieves similar or better segmentation performances. Notably, we observe that SamDSK is stable and reliable in providing segmentation improvement, while some of the known methods can encounter situations where worse segmentation performances might be yielded in some cases.

The ISIC 2018 skin lesion segmentation dataset [8] contains 2594 training dermoscopic images and 1000 test images for melanoma segmentation. Once again, we employ HSNet and PraNet for the experiments. We are aware that the image-level DSK for skin lesion segmentation in most of these images is that they contain only one region of interest. Consequently, we set $v_{1}^{lower}$ to 1 and $v_{1}^{upper}$ to 1 in Eq. (3) (the melanoma class is designated as class 1) throughout all the rounds of processing. The results in Table 3 demonstrate that SamDSK enhances the segmentation performances of the state-of-the-art models when utilizing 10% of the total labeled samples. SamDSK achieves a higher performance gain in comparison to the competing methods.