Annotation by Clicks: A Point-Supervised Contrastive Variance Method for Medical Semantic Segmentation

Weakly-supervised segmentation methods have proposed to use different types of weak supervision information, such as point-level (Bearman et al. 2016; Qian et al. 2019; Tang et al. 2018), scribble-level (Lin et al. 2016; Tang et al. 2018), box-level (Dai, He, and Sun 2015; Khoreva et al. 2017), and image-level (Ahn and Kwak 2018) supervisions. In particular, Lin et al. (2016) proposed to propagate information from scribble regions to unlabeled regions by generating proposals. Tang et al. (2018) combined the partial cross entropy for labeled pixels and the normalized cut to unlabeled pixels to obtain a more generalised model. Kim and Ye (2019) proposed a novel loss function based on Mumford-Shah functional that can alleviate the burden of manual annotation. Qian et al. (2019) considered both the intra- and inter-category relations among different images for distance metric learning. In order to obtain reasonable pseudo-labels, some methods stack temporal level predictions (Lee and Jeong 2020) or constrain the consistency of different outputs from the model to increase the performance (Ouali, Hudelot, and Tami 2020; Chen et al. 2021c). The performances of the above methods nevertheless still have considerable gaps from fully supervised methods and their study has focused on natural scene images.

Medical image semantic segmentation with weak supervision poses different challenges from the segmentation task for natural scene images (Luo et al. 2022; Dorent et al. 2021; Roth et al. 2021; Qu et al. 2020; Tajbakhsh et al. 2020). Some works have explored using extreme points as supervisory signals for medical image segmentation (Dorent et al. 2021; Roth et al. 2021). Roth et al. (2021) used the Random Walker according to extreme points to generate initial pseudo-labels, which were then used to train the segmentation model. Subsequently, Dorent et al. (2021) generated more annotated voxels by deep geodesics connecting and used a CRF regularised loss to obtain better performance. Kervadec et al. (2019) proposed a loss function that constrains the prediction by considering prior knowledge of size. Laradji et al. (2021) constrained the output predictions to be consistent through geometric transformations. Can et al. (2018) introduced an iterative two-step framework that uses conditional random field (CRF) to relabel the training set and then trains the model recursively. Recently, Luo et al. (2022) proposed a dual-branch network to generate pseudo-labels dynamically, which learns from scribble annotations in an end-to-end way. Although the above work has made some progress, they have not yet satisfactorily addressed the problem of weakly supervised semantic segmentation of medical organs, and more research effort is needed on this task.

Following the framework in Figure 2, the prediction output $\hat{Y}^{n}\in[0,1]^{K*H*W}$ for a training image $I^{n}$ can be generated through the pipeline of an encoder $f_{encoder}$ and a decoder $f_{decoder}$, as follows:

$$\hat{Y}^{n}=\mbox{softmax}(f_{decoder}\circ f_{encoder}(I^{n})),$$

(1)

where $\mbox{softmax}(\cdot)$ denotes the class-wise softmax function that predicts each pixel into a probability vector of $K$ classes. Let $\hat{Y}^{n}_{k}(r)$ denote the predicted probability of a pixel at location $r$ belonging to the $k$-th class. We then have $\sum_{k=1}^{K}\hat{Y}^{n}_{k}(r)=1$.

In general point-supervised learning, especially given a fraction of labeled points, the partial cross-entropy loss $L_{pCE}$ is typically minimized to train the segmentation network (Tang et al. 2018). Here given the few annotated points indicated by $\Omega_{y}^{n}$ for each training image $I^{n}$, the cross-entropy loss $L_{pCE}$ is computed by using the point-supervised ground truth $Y^{n}$ and the prediction output/mask $\hat{Y}^{n}$ as follows:

$$L_{pCE}=-\sum_{n=1}^{N}\sum_{r\in\Omega_{y}^{n}}\sum_{k=1}^{K}Y^{n}_{k}(r)\log\hat{Y}^{n}_{k}(r)$$

(2)

In this subsection, we provide an overview of the Mumford-Shah loss function (Kim and Ye 2019), which treats image segmentation as a Mumford-Shah energy minimization problem (Mumford and Shah 1989). This loss function aims to segment images into $K$ classes of piece-wise constant regions by enforcing each region to have similar pixel values with the regularization of the contour length.

Specifically, by using the prediction function output $\hat{Y}_{k}$ as the characteristic function for the $k$-th class (Kim and Ye 2019), the Mumford-Shah loss functional $L_{MS}$ is defined as below:

$$\begin{split}L_{MS}=&\lambda_{ms}\sum_{n=1}^{N}\sum_{k=1}^{K}\int_{\Omega}|I^{n}(r)-c_{k}^{n}|^{2}\hat{Y}^{n}_{k}(r)dr\\ &+\mu\sum_{n=1}^{N}\sum_{k=1}^{K}\int_{\Omega}|\nabla\hat{Y}^{n}_{k}(r)|dr,\end{split}$$

(3)

where $c_{k}^{n}\in\mathbb{R}^{1*1*1}$ is the average pixel value of the $k$-th class in the $n$-th image; the gradient $\nabla\hat{Y}^{n}_{k}(r)=\frac{\partial\hat{Y}^{n}_{k}(r)}{\partial r}$ computes the horizontal and vertical changes in the predictions, and is used in the second term to enforce the predictions to be smooth within each of the organs.

It is very challenging to identify boundaries between different classes with just a single point annotation per class in an image. The partial cross-entropy loss on the labeled pixels can only provide limited guidance. It is therefore essential to exploit the unannotated pixels. Although the Mumford-Shah loss above exploits the unlabeled pixels in a predictive K-means clustering manner, it still lacks sufficient discriminative capacity for separating different categories and identifying the segmentation boundaries. Meanwhile, medical images have some unique properties for the spatial distributions of different categories of organs—even with operations such as random rotation and flipping, the spatial locations of the same organ, especially the smaller ones, in different images will still be within some limited regions in contrast to the whole image background. Inspired by this observation, we develop a novel contrastive variance loss function to exploit the statistical spatial category distribution information in a contrastive manner, aiming to enhance the discrepancy between each organ and its background.

Specifically, we first calculate the pixel-level variance distribution map $z_{k}^{n}\in\mathbb{R}^{1*H*W}$ for each $k$-th organ category on the $n$-th image, such as:

$$z_{k}^{n}(r)=|I^{n}(r)-c_{k}^{n}|^{2}\hat{Y}_{k}^{n}(r),\forall r\in\Omega,$$

(4)

where $c_{k}^{n}\in\mathbb{R}^{1*1*1}$ is the mean pixel value for the $k$-th category region on the $n$-th image, which is defined as:

$$c_{k}^{n}=\frac{\int_{\Omega}I^{n}(r)\hat{Y}^{n}_{k}(r)dr}{\int_{\Omega}\hat{Y}_{k}^{n}(r)dr}.$$

(5)

We then use this variance distribution map $z_{k}^{n}$ as the appearance representation of the $k$-th class in the $n$-th image and define the following contrastive variance loss to promote the inter-image similarity between the same class of organs in contrast to the similarities between different classes:

$$\begin{split}L_{CV}=&\lambda_{cv}\sum_{n=1}^{N}\sum_{k=1}^{K}-\log\frac{pos}{pos+neg}\\ &+\mu\sum_{n=1}^{N}\sum_{k=1}^{K}\int_{\Omega}|\nabla\hat{Y}^{n}_{k}(r)|dr,\end{split}$$

(6)

with

	$\displaystyle pos$	$\displaystyle=\exp(\cos(z_{k}^{n},z_{k}^{m_{n}})/\tau),$		(7)
	$\displaystyle neg$	$\displaystyle=\sum_{i\neq n}\sum_{j\neq k}\exp(\cos(z_{k}^{n},z_{j}^{i})/\tau),$		(8)

where $\cos(\cdot,\cdot)$ denotes the cosine similarity function, and $\tau$ is the temperature hyper-parameter. In this loss function, given the pixel variance distribution map $z_{k}^{n}$ for the $k$-th category on the $n$-th training image, we build the contrastive relationship by randomly selecting another $m_{n}$-th image that contains the same $k$-th category and using $z_{k}^{m_{n}}\in\{z_{k}^{i},i\not=n|z_{k}^{n}\}$ as a positive sample for $z_{k}^{n}$, while using all pixel variance distribution maps for all the other categories {$j:j\not=k$} on all other training images $\{i:i\not=n\}$ as negative samples for $z_{k}^{n}$.

Different from the typical contrastive losses popularly deployed in the literature, which are computed over intermediate features (Khosla et al. 2020; Hu et al. 2021), our proposed contrastive variance loss above is specifically designed for the point-supervised medical image segmentation task and is computed over the pixel-level variance distribution map. This contrastive variance (CV) loss integrates both the variance and contrastive learning paradigms in a way that has not been attempted, and is designed to possess the following advantages. First, under this CV loss, by maximizing the similarities between the positive pairs in contrast to the negative pairs, we can enforce each prediction $\hat{Y}^{n}_{k}$ into the statistically possible spatial regions for the $k$-th category across all the training images. This can help eliminate the complex background formed by all the other categories and improve the discriminative capacity of the segmentation model. Second, by using the pixel-level variance distribution maps as the appearance representation of the organs for similarity computation, one can effectively eliminate the irrelevant inter-image pixel variations caused by different collection equipment or conditions. Third, using the cosine similarity function to compute the similarity between a pair of $(z_{k}^{n},z_{k}^{m_{n}})$ is equivalent to computing the dot product similarity between two normalized probability distribution vectors. Hence, the smaller effective length of the vectors will lead to larger similarities. This property consequently will push each organ prediction $\hat{Y}^{n}_{k}$ to cover more compact regions for more accurate predictions of the unlabeled pixels, while eliminating the boundary noise. Overall, we expect this novel loss function can make effective use of all the unlabeled pixels to support few-point-annotated segmentation model training.

The proposed PSCV is straightforward and can be trained in an end-to-end manner. It does not require any additional layers or learnable parameters on top of the base segmentation network. During the training phase, the learnable parameters of the segmentation network are learned to minimize the following joint loss function:

$$L_{total}=L_{pCE}+L_{CV},$$

(9)

where $L_{pCE}$ encodes the supervised information in the labeled pixels and $L_{CV}$ exploits all pixels in the training images in a self-supervised manner. Note for batch-wise training procedures, these losses are computed within each batch.

During the testing stage, each test image $I$ can simply be given as an input to the trained segmentation network (encoder and decoder) to obtain the prediction results $\hat{Y}$.

Following (Wang et al. 2022; Chen et al. 2021b), we benchmark the proposed PSCV on two medical image datasets: the Synapse multi-organ CT dataset and the Automated cardiac diagnosis challenge (ACDC) dataset. The Synapse dataset consists of 30 abdominal clinical CT cases with 2211 images and contains eight organ categories. Following (Wang et al. 2022; Chen et al. 2021b), 18 cases are used for training and 12 cases are used for testing. The ACDC dataset contains 100 cardiac MRI examples from MRI scanners with 1300 images and includes three organ categories: left ventricle (LV), right ventricle (RV), and myocardium (Myo). We evaluate the proposed PSCV on the ACDC via five-fold cross-validation (Luo et al. 2022). Following previous works (Ronneberger, Fischer, and Brox 2015; Wang et al. 2022; Chen et al. 2021b), the Dice Similarity Coefficient (DSC) and the 95% Hausdorff Distance (HD95) are used as evaluation metrics.

We randomly select one pixel from the ground truth mask of each category as labeled data to generate point annotations for each training image. The most widely used medical image segmentation network UNet (Ronneberger, Fischer, and Brox 2015) is used as the base architecture, and the proposed method is trained in an end-to-end manner. In the training stage, we first normalized the pixel values of input images to [0,1]. Then the input images are resized to 256$\times$256 with random flipping and random rotation (Luo et al. 2022). The weights of the proposed model are randomly initialized. The model is trained by stochastic gradient descent (SGD) with a momentum of 0.9, a weight decay of 1e-5, a batch size of 12, and a poly learning rate policy with a power of 0.9. We set the hyperparameters $\mu$=1e-5 and $\tau$=0.07. We use a learning rate of 0.01 for 60K total iterations and $\lambda_{cv}$=1e-3 on the ACDC dataset and a learning rate of 0.001 for 100K total iterations and $\lambda_{cv}$=3e-1 on the Synapse dataset. Inspired by the central bias theory (Levy et al. 2001), in the testing stage we filtered fifty pixels that are close to the left and right edges of the prediction as the background to further refine the results on Synapse for all the methods.

The proposed PSCV is compared with several state-of-the-art methods on the Synapse and the ACDC datasets: pCE (Lin et al. 2016) (lower bound), EntMini (Grandvalet and Bengio 2004), WSL4MIS (Luo et al. 2022), USTM (Liu et al. 2022) and GatedCRF (Obukhov et al. 2019). For fair comparisons, we use UNet as the backbone for all the comparison methods, while all of them are trained with the same point-annotated data. In addition, we also tested the fully supervised model trained with the cross-entropy loss, which uses the dense annotations of all pixels in the training images, and its performance can be treated as a soft upper bound for the point-supervised methods.

To demonstrate the effectiveness of the proposed CV loss, we conducted an ablation study by comparing the full PSCV method with two variant methods: (1) “pCE” denotes the variant that performs training only using the partial cross-entropy loss $L_{pCE}$; (2) “vanilla MS” denotes the variant that replaces the CV loss, $L_{CV}$, in the full PSCV method with the standard Mumford-Shah loss, $L_{MS}$. The comparison results on the Synapse and the ACDC datasets are reported in Table 3 and Figure 3 respectively. From Table 3, we can see that using $L_{pCE}$ alone produces poor segmentation results. With the additional $L_{MS}$ loss, “vanilla MS” improves the segmentation results by 16.5% in terms of class-average DSC and by 71.25 in terms of class-average HD95. By using the proposed contrastive variance loss $L_{CV}$, PSCV further improves the performance by another 8.3% in terms of DSC and by 4.2 in terms of HD95. Similar observations can be made from the results reported in Figure 3 as well. Overall, the proposed CV loss is more effective than the regular Mumford-Shah loss.

To provide a qualitative analysis of the proposed PSCV, we present a visualized comparison between our results and the results produced by several state-of-the-art methods (i.e. pCE, EntMini, GatedCRF, USTM, vanilla MS and WSL4MIS) in Figure 5. We can see that the various organs differ significantly in size and shape in medical organ images, and the background regions are frequently mistaken for organs by many methods. Additionally, different image samples have different intensities of the same organs, which seriously undermines the effectiveness of the existing segmentation methods. From the qualitative results in Figure 5, we can see that benefiting from using the contrastive paradigm and the variance distribution map representations in designing the CV loss, the proposed PSCV can better discriminate the organs from the background, and produce more compact and accurate segmentations than the other comparison methods. Its ability to accurately segment the smaller organs is especially notable.