EAUWSeg: Eliminating annotation uncertainty in weakly-supervised medical image segmentation

Without the requirement of large densely annotated data, weakly-supervised learning has gained significant attention in medical image segmentation [13, 2]. As the most efficient weak annotation method, image-level annotations only require classification labels and generates class activation maps [14] for training. Although image-level annotations method is convenient, it has limited performance due to the extremely weak supervision [15]. Box-level annotation is usually defined by two corner coordinates, which provides localization-awareness compared to the image-level annotation [16]. However, boxes for different objects may tend to overlap with each other, making it difficult to accurately approximate the target boundary, especially for complex shapes [10]. Point annotations provide a small number of pixels for different classes and can better handle complex shapes, which may be more preferable to medical segmentation compared to box-level annotations. Despite its efficiency, the segmentation model trained with point annotations tends to overfit the small number of annotated pixels when comparing with the large number of unannotated pixels.

Scribble-based annotations provide labels for a sparse set of pixels of each class for training, and are usually more obtainable in medical image segmentation by considering its annotation efficiency, performance effectiveness as well as the friendliness to the annotation of nested structure [17]. Only the scribbles of the background and each object are given, while the groundtruth of other pixels remains unknown, which is harmful to the segmentation performance. An intuitive resolution is to expand the scribble annotations by considering the prior assumptions [18] or using the learned foreground features through the deep neural networks [19]. However, due to the lack of supervisory signals, the constructed models usually fail to capture the object structure and confuse on the object boundary. To address this issue, a series of studies have concentrated on learning adversarial shape priors at the expense of requiring additional fully-annotated masks [17]. However, acquiring such fully annotated datasets may present challenges in many clinical practices, rendering these existing methods both costly and lacking in scalability. Our work aims to explore new weak annotation method that can prompt the performance of automated medical image segmentation without auxiliary datasets.

Contrastive learning argues that similar samples should have similar representations, and the representations of different samples should be different [20]. Based on this, contrastive loss is usually designed to enforce representations to be similar for similar pairs and dissimilar for dissimilar pairs [21]. Considering its powerful self-supervised feature extracting ability from the unlabeled data, contrastive learning has been widely used in many image-level tasks. Among all these methods, the key is the selection mechanism designing of contrastive sample pairs, i.e., positive and negative pairs.

Recently, contrastive learning has been extended from image-level task to pixel-level ones to mine informative information from unlabeled data [22, 23]. As mentioned earlier, constructing contrastive sample pairs is crucial for discriminative feature learning. In the context of pixel-level tasks, sample pairs are usually constructed through pseudo labels or spatial structure, which may introduce noisy sampling. To alleviate this problem, prediction uncertainty has been injected into the sampling to reduce the number of noisy samples [24].

In the scenario of classical weakly-supervised segmentation, the input pixels $x$ are usually divided into the labeled pixels $x_{l}$ and unlabeled pixels $x_{ul}$. In this way, the corresponding labels $y_{l}$ for the labeled pixels $x_{l}$ will be directly used for supervision by employing the partial cross-entropy loss, which can be formulated as follows:

where $p$ is the segmentation prediction. For the unlabeled pixels, there is no off-the-shelf label for supervision, and a lot of work focus on assigning pseudo labels to unlabeled pixels for supervision [28, 29]. The overall objective function can be formulated as follow:

However, assigning pseudo labels to unlabeled pixels not only requires a time-consuming multi-stage training process, but also results in misleading or biases [10].

To address this problem, this work introduces the bounded polygon annotation method that simply requires labeling two polygons that are similar to the inscribed and outer envelope-like delineations of lesion (as shown in Fig. 1). To further eliminate the uncertainty included in the bounded polygon annotation, we explicitly treat bounded polygons as two separated annotation, i.e., inscribed-like annotation $y^{in}$ and envelope-like annotation $y^{en}$. Different from the classical weakly annotation methods, the input pixels $x$ are divided into the certain labeled pixels and the uncertain pixels in our bounded polygon annotations. This work aims at providing more reliable supervision signal for pixels in uncertain region during model training.

For convenience, we define $\Omega_{I}$, $\Omega_{\Delta}$, $\Omega_{O}$ as the spatial domain inside the inscribed-like annotation, between the inscribed-like and envelope-like delineations, and outside the envelope-like annotation, respectively. Here, the certain labeled pixels and uncertain pixels can be depicted as $x_{i}\in\Omega_{I}\cup\Omega_{O}$ and $x_{uc}\in\Omega_{\Delta}$, the corresponding labels $y_{i}=1$ if $x_{i}\in\Omega_{I}$ and $y_{i}=0$ if $x_{i}\in\Omega_{O}$, otherwise, $y_{i}$ is uncertain. The spatial domain of input image $x$ and the envelope-like annotation can be depicted as $\Omega=\Omega_{I}\cup\Omega_{O}\cup\Omega_{\Delta}$ and $\Omega_{E}=\Omega_{I}\cup\Omega_{\Delta}$, respectively. In this way, our proposed EAUWSeg tries to learn from the “certain/uncertain” pixels instead of “labeled/unlabeled” pixels in the classical weakly-supervised segmentation. The objective function in our EAUWSeg can be re-formulated as:

The feature learning of certain pixels have been well solved. Hence, this work focuses on eliminating annotation uncertainty and thus providing reliable supervision signals for pixels in the uncertain regions.

EAUWSeg is tailored for the proposed bounded polygon annotation, and mainly focuses on eliminating annotation uncertainty for pixels belong to $\Omega_{\Delta}$. As illustrated in Fig. 2, EAUWSeg consists of 1) a mainstream segmentation network supervised by two bounded polygons segmentation labels to implicitly define the certain region and uncertain region during network training, 2) a classification-guided confidence generator to provide the category-level prediction confidence for pixels $x_{i}\in\Omega_{\Delta}$ by leveraging a tailored multi-class classification task, 3) a confidence-auxiliary consistency learner to distinguish reliable pixels in uncertain region can assign the corresponding “certain” labels.

Let $\mathcal{S}_{e}$, $\mathcal{S}_{d}$, and $\mathcal{S}_{h}$ denote the encoder, the decoder, and segmentation head used in our proposed framework that are parameterized by $\Theta_{e}$, $\Theta_{d}$ and $\Theta_{h}$, respectively. In the proposed EAUWSeg, the bounded polygon annotation is treated as two separate masks, i.e., inscribed-like and envelope-like masks, and the basic segmentation loss function in EAUWSeg can be formulated as:

where $p$ is the predicted probability maps for input image $x$. In this work, the following dice loss is employed for both $\mathcal{L}_{in}$ and $\mathcal{L}_{en}$:

where $H\times W\times D$ denotes the input image size, $p_{i}$ and $y_{i}$ denotes the prediction probability and label for pixel $i$, respectively. However, training with $\mathcal{L}_{c}$ will introduce inconsistency supervision signals, since $y^{in}$ and $y^{en}$ have the following characteristics: $y^{in}_{i}=1$ for $x_{i}\in\Omega_{I}$, $y^{en}_{i}=1$ for $x_{i}\in\Omega_{E}$, and others are $0$. In this way, pixels in uncertain region will have different labels during training, i.e., $y^{in}_{i}=0$ while $y^{en}_{i}=1$ for $x_{i}\in\Omega_{\Delta}$.

To mitigate the influence and leverage this adversarial supervision signal to learn invariant feature during model training, this work focuses on assign more reliable labels for pixels in uncertain region by utilizing feature representation of certain pixels. For uncertain pixels, we want to utilize the potential similarity between pixels in the same category, i.e., $x_{i}\in\Omega_{I}\cup\Omega_{O}$ and $x_{j}\in\Omega_{\Delta}$ to mine informative information. With these definitions, the loss function of BPAnno-supervised segmentation can be formulated as:

Here, $\mathcal{L}_{uc}(x,y^{in},y^{en})$ denotes the loss function for uncertain pixels.

The key point for accurate BPAnno-supervised segmentation is reliable labels assigning for pixels in uncertain region. Different from existing methods that focus on iteratively assigning pseudo label for uncertain pixels, we propose to utilize the intra-class similarity and inter-class discriminative from both the feature and category perspective.

An intuitive idea to approximate the confidence for uncertain pixels $x_{i}$ is the predictive entropy that is calculated according to the following equation:

where $\Theta_{s}=\{\Theta_{e},\Theta_{d},\Theta_{h}\}$ are the parameters of standard segmentation network, $\epsilon$ is a constant to avoid overflow. Similar as previous works, prediction with large entropy is considered as the solid uncertain pixels, which will be dropped during the subsequent learning. For clarity, we define the solid uncertain pixels in uncertain region with category of $-1$:

where $\mu$ is a predefined threshold to mask the uncertain labels, and $\mathcal{U}^{e}\in\mathbb{R}^{C\times H\times W}$ is the estimated uncertainty map with the same size as input image.

To assign more reliable labels for uncertain pixels, we propose to explicitly leverage the potential similarity between certain and uncertain pixels by employing a tailored classification task, which aims at removing as much uncertainty as possible. Let $f_{\mathcal{S}}(x)$ denote the feature representation generated through the encoder and decoder network, $\mathcal{S}_{c}$ and $\Theta_{c}$ denote the classification head and its corresponding parameters respectively. Previous work [10] has shown that similar pixels in the feature space preferable to generate consistent category prediction. Based on this, the constructed classification head is used to model a multi-class classification task with the objective function of:

where $y^{c}$ is the classification labels that $y_{i}^{c}=0$ for $x_{i}\in\Omega_{O}$, $y_{i}^{c}=1$ for $x_{i}\in\Omega_{\Delta}$, and $y_{i}^{c}=2$ for $x_{i}\in\Omega_{I}$, and $N=H\times W\times D$.

During model training, we assume that “certain” pixels in uncertain region would prefer to generate prediction of $y_{i}^{c}=0$ for background and $y_{i}^{c}=2$ for foreground. In this way, the uncertain map generated by the confidence map can be formulated as:

where $\odot$ refers to the element-wise multiplication, and $\mathcal{M}_{u}$ is a mask with $x_{i}=1$ for $x_{i}\in\Omega_{\Delta}$, and $x_{i}=0$ otherwise. The final confidence for the uncertain pixels can be formulated as:

Here, $\mathcal{U}$ means that both pixels with large predictive entropy, i.e., $\mathcal{E}_{i}\geq\mu$ and with uncertain classification prediction, i.e., pixels with prediction of $1$ for the multi-class classification task, will be considered as solid uncertain and be assigned with label of $-1$.

Confidence-auxiliary consistency learner aims at generating “certain” information from uncertain region to facilitate stable training. An intuitive idea is utilizing contrastive learning to reduce the distance between pixels within same category while enlarging the distance between pixels in different categories. This strategy allows us to conduct the pixel-wise contrastive learning. However, the crucial question is the selection of positive and negative samples, especially for pixels in uncertain region. To reduce the influence of uncertain information, we propose to utilize the generated confidence for the uncertain pixels and only the solid certain pixels will be considered during the pixel-wise contrastive learning. In this way, the determined pseudo labels can be obtained as follows:

To provide more reliable supervision signal by using the pixel-wise contrastive learning, we follow two guidelines during sample selection: 1) only feature embedding of pixels in the certain region are stored in this study and further be sampled during the computation of contrastive loss; 2) the anchor sampling in this study focuses on hard samples with error prediction for $x_{i}\in\Omega_{I}\cup\Omega_{O}$, and samples with higher certainty for $x_{i}\in\Omega_{\Delta}$. The pixel-wise contrastive loss in this work can be defined as:

where $\mathcal{P}$ contains the indexes of all “certain” pixels in the uncertain region; $P$ and $N$ contains the indexes of positive pixels, i.e., pixels has same class with pixel $i$, and negative pixels, i.e., pixels with different labels to pixel $i$, in the certain region, respectively; $\tau$ is a temperature hyper-parameter.

Considering the semantic representation for deep layers and the effective information for uncertain pixels, feature representation before the segmentation head is embedded into a specific feature space and is employed as the prototype vector in contrastive learning. That is to say, $f_{i}$ denotes the feature embedding of pixels $x_{i}$ that is calculated according to the following equation:

To summarize, the overall objective function includes two parts: 1) losses for “certain” pixels using fully-supervised segmentation setting, 2) confidence-guided contrastive loss for uncertain pixels. At the early stage of training, segmentation model need to learn the feature representation of lesions with the guidance of supervised loss for “certain” pixels, i.e., $\mathcal{L}_{c}$. When the segmentation performance gradually improves, the contrastive loss $\mathcal{L}_{PCL}$ combined with a multi-class classification cross-entropy loss $\mathcal{L}_{ce}$ are added to apply constraints on uncertain pixels in same class to preserve consistency feature representation. Therefore, the overall objective function in this work is formulated as:

where $\lambda_{1}$ and $\lambda_{2}$ are the parameters to control the contribution of confidence-auxiliary consistency learner and classification-guided confidence generator, respectively.

To demonstrate the comprehensive segmentation performance of our method, we compare EAUWSeg with different state-of-the-art approaches:

• •

  Scribble-supervised methods: 1) different learning strategies on UNet, including partially Cross-Entropy loss [33], Total Variation loss [34], Gated Conditional Random Field loss [35], Mumford-Shah Loss [36], as well as Uncertainty-aware Self-ensembling and Transformation-consistent Mean Teacher techniques (USTM) [37]. 2) different scribble-supervised frameworks, including ScribbleVC [8], DMSPS [1], and TriMix [38].

• •

  Box-supervised methods. For fair comparison, we also present the results of classical segmentation networks, i.e., UNet and TransUNet, supervised with bounding box.

• •

  Fully-supervised segmentation methods: 1) CNN-based methods, including UNet [25], UNet++ [30], and DeepLabV3+ [26]. 2) Transformer-based methods, including TransUNet [27], TransFuseS [31], HiFormer [32]. Implementation of these networks follow the corresponding github repositories. During training, ResNet50 [42] is employed as the encoder for UNet and DeepLabV3+, ResNet34 is utilized in the UNet++ and TranFuseS, the default “R50-ViT-B_16” and “Hiformer-S” configurations are employed for TransUNet and HiFormer.

Table 1 presents the quantitative evaluation results of the aforementioned methods. For fair comparison with scribble-supervised methods with different learning strategies, we present the results of EAUWSeg with UNet as backbones. To demonstrate the effectiveness of EAUWSeg, we also give the results with TransUNet as backbone. The results illustrate that our method outperforms other weakly-supervised methods on both ISIC2017 and Kvasir-Seg datasets, including the scribble-supervised as well as the box-supervised methods. When compared with the fully-supervised methods, our proposed EAUWSeg can also deliver superior performance, yielding an average Dice score of 86.60% and 89.88% on ISIC2017 and Kvasir-Seg, respectively. This underscores the superiority and effectiveness of the proposed BPAnno-supervised strategy and its corresponding learning framework EAUWSeg. Fig. 3 shows some qualitative evaluation results, it can be seen that our proposed method achieves better segmentation performance.