Concentrate on Weakness: Mining Hard Prototypes for Few-Shot Medical Image Segmentation

1 Introduction

Automatic segmentation of organs, tissues and lesions appearing in medical images is broadly applied in clinical research, including disease diagnosis [Tsochatzidis et al., 2021], treatment planning [El Naqa et al., 2007; Zaidi and El Naqa, 2010; Chen et al., 2021] and surgical guidance [Wang et al., 2018; Hesamian et al., 2019]. Recently, fully supervised deep learning segmentation frameworks [Fang and Yan, 2020; Shi et al., 2021; Isensee et al., 2021] exhibit superior performance, but they heavily depend on large amounts of densely annotated training data, which is impractical for medical scenarios for the reason that labeling medical images requires considerable clinical expertise, cost and time [Luo et al., 2021].

∗Corresponding author.

To solve this challenge, Few-Shot Learning (FSL) [Snell et al., 2017; Sung et al., 2018; Garcia and Bruna, 2017; Finn et al., 2017] has been proposed as a potential learning paradigm to train models that can efficiently adapt to novel classes when given extremely scarce annotated data [Shaban et al., 2017], based on which corresponding segmentation method is called Few-Shot Medical Image Segmentation (FSMIS). In particular, FSMIS typically trains models with base classes and segment unseen classes in the inference stage. Consequently, models trained in this manner are capable of segmenting novel organ or lesion leveraging solely a few annotated data, evading the need for retraining, which holds significant importance for medical image segmentation.

As one of dense image prediction tasks, FSMIS has an indispensable demand for features with strong category information and precise spatial boundary details. In order to leverage the annotated samples, the majority of current FSMIS models distill prior knowledge based on prototype network [Ouyang et al., 2020; Shen et al., 2023; Wu et al., 2022], among which multiple prototypes-based network is the predominant method. However, these methods generate multiple prototypes from support images either by a random selecting strategy or simple local averaging, which incline to normal features in the central area of a specific organ or involve background information in foreground prototypes as illustrated in Fig. 1(a). Furthermore, those weaker features mainly distributed along the boundary of the organ are neglected to a considerable extent, which can lead to a poorly segmentation boundary. As a result, some foreground and background points will be misclassified due to the excessive focus on normal features.

In this paper, we propose a Concentrate on Weakness (CoW) approach that learns to mine hard prototypes for each category to represent its feature distribution more comprehensively as illustrated in Fig. 1(b), to tackle the aforementioned problem. Specifically, we devise a Support Self-Prediction (SSP) module to provide knowledge for weak features carrying critical boundary information, which are subsequently utilized by the Hard Prototypes Generation (HPG) module to generate hard prototypes to improve segmentation accuracy and robustness. Right after is the Multiple Similarity Maps Fusion (MSMF) module that fuses multiple similarity maps to incorporate comprehensive information included in them to get the final predicted mask. To acquire a more precise segmenting boundary, we also devise a boundary loss, which effectively suppress over-segmentation whether for the foreground or the background. In conclusion, our main contributions are summarized as follows:

• We propose a novel method to identify weak features by comparing predicted support mask and its true label. Besides, multiple hard prototypes are generated by further processing those weak features, ensuring more comprehensive representation of overall distribution.
• To make full use of the information provided by generated multiple prototypes, we design an MSMF module to fuse multiple similarity maps predicted by these prototypes to get the final mask.
• Extensive experiments conducted on three popular medical image datasets reveal that our proposed model outperforms other existing state-of-the-art (SOTA) models and exhibits significant performance improvements.

2 RelatedWork

2.1 Medical Image Segmentation

With the development of Convolutional Neural Networks (CNNs), deep learning methods have become the dominated approach to MIS on various tumor [Peiris et al., 2022], anatomical structures [Dolz et al., 2018] and lesions [Li et al., 2018] recently. Fully Convolutional Networks (FCNs), which replace the fully connected layers of standard CNNs with fully convolutional layers, prove to be a powerful net architecture for semantic segmentation [Long et al., 2015]. Afterwards, the encoder-decoder networks gradually become the prevailing architecture for semantic segmentation, among which U-Net is a widely recognized method for MIS with remarkable results [Ronneberger et al., 2015]. Besides, most recent MIS networks are designed based on visual transformer [Cao et al., 2022], multi-scale fusion [Wu et al., 2019], feature pyramid [Chen et al., 2018] and self-attention mechanism [Gu et al., 2020]. Unfortunately, all these models require abundant densely annotated data.

2.2 Few-Shot Medical Image Segmentation

FSMIS is a promising approach since it can work with a few labeled samples. The prototypical network [Wang et al., 2019; Ouyang et al., 2020], which extracts representative prototypes of semantic classes from the support, is one of prevalent FSMIS methods featured by its strong interpretability. PANet [Wang et al., 2019] argues that parametric segmentation generalizes poorly and proposes a non-parametric prototype matching method with a prototype alignment regularization. ADNet [Hansen et al., 2022] builds further on the branch initiated by PANet [Wang et al., 2019] and does not explicitly model the complex background class namely relying solely on one foreground prototype. Q-Net [Shen et al., 2023] replaces fixed thresholds used in ADNet with adaptive ones learned from the query to mitigate distribution shifts between the query image and the support set, and simultaneously predicts the query segmentation masks with dual-scale features. Most of these methods solely employ masked average pooling to extract a single prototype, which can not well represent the distribution of a specific category.

Unlike these single prototype-based methods, SSLALPNet [Ouyang et al., 2020] devises an adaptive local prototype pooling module where additional local prototypes are computed on a regular grid to preserve local information lost in global average operation. AAS-DCL [Wu et al., 2022] constructs prototypical contrastive learning and contextual contrastive learning among the labeled support guidance information and unlabeled tissue knowledge, to strengthen the discriminability of the prototypes. GMRD [Cheng et al., 2024] randomly selects multiple prototypes from both foreground and background to comprehensively represent the commonality within the corresponding class distribution. However, these methods pay no particular attention to weak boundary features even though they have constructed relatively comprehensive representation of each class with multiple prototypes.

3 Methodology

3.1 ProblemDefinition

The FSMIS task aims at training a model on a dataset Dtrain that contains labeled samples belonging to seen classes Cbase, allowing a quick adaptation to segmentation of previously unseen novel classes Cnovel when exposed to just a few annotated samples without retraining, where Cbase and Cnovel are disjoint, namely Cbase∩Cnovel = ∅. In general, the episodic [Vinyals et al., 2016; Dong and Xing, 2018] meta-learning setting is commonly utilized in FSS tasks. Specifically, by means of random sampling, the training set and the testing set are both divided into a support set and a query set (e.g., Dtrain = {(Si,Qi)}Ntrain i=1 and Dtest = {(Si,Qi)}Ntest i=1 ), where Ntrain and Ntest are the number of episodes for training and testing respectively. When considering a Nway K-shot task, the support set Si = {(I k s ,M k s (cj))}Kk=1 contains K image-mask pairs of an image I ∈ RH×W and its corresponding binary mask M ∈ {0, 1}H×W for class cj ∈ Cbase, j = 1, 2, ...,N and the query set Qi = {(I k q ,M k q (cj))}Nqry k=1 contains Nqry image-mask pairs from the same class as the support set. Notably, the background class denoted as c0 does not count toward Cbase or Cnovel. Given a support image-mask pair and a query image as input, the model for FSS task outputs predicted segmentation mask ∼Mq for the query image.

3.2 Network Overview

The proposed CoW is a dual-path architecture where similar operations are performed, one for segmenting foreground regions and the other for background. The framework comprises following primary components: (1) a generic feature encoder shared by the support and query images for extracting features; (2) an SSP module for acquiring a coarse predicted mask of the support image using global support prototype; (3) an HPG module for generating a specified number of hard prototypes for both foreground and background; (4) an MSMF module for fusing multiple similarity maps to get final segmenting mask of the query image. Figure 2 shows a more detailed pipeline of our method.

3.3 Support Self-Prediction

As shown in Figure 2, our proposed CoW starts by a parameter-sharing feature encoder denoted as fθ, which extracts features from both support and query images. We denote the support and the query features as Fs = fθ(Is) and Fq = fθ(Iq), Fs, Fq ∈ RC′×H′×W ′ , where C ′ is the channel depth of the feature, and H′ and W ′ denote the height and width of the feature respectively. Following the practice in Q-Net [Shen et al., 2023], we adopt ResNet-101 pre-trained on MS-COCO [Lin et al., 2014] as the backbone, and utilize the output feature of the third residual block.

In order to find out the weak feature points easily misclassified, which lead to over-segmentation and boundaryblurring, we devise a Support Self-Prediction (SSP) module to provide clues for them. Specifically, we firstly employ masked average pooling to generate a global support prototype ps ∈ R1×C′ from the support feature map Fs:

where Fs is resized to the size of corresponding mask Ms. ⊙ denotes the Hadamard product, and (h, w) is the pixel position. Then we duplicate and expand the support prototype ps to the same size as Fs, represented as Fp, which is subsequently concatenated with Fs to generate a new feature map:

where Concat(·) denotes the concatenation operator.

Next, the similarity map for the support image is generated after passing through the feature processing module and the decoder that adopts single-scale residual layers following an

where Ssupport ∈ R1×H′×W ′ is the predicted support similarity map. FPM(·) is the feature processing operation and D(·; ψ) means the support decoder parameterized with ψ. After that, we generate the predicted support mask ˆMs through a concatenation operation followed by an argmax operation:

where Ssupport is upsampled to (H,W) in advance and 1 - Ssupport represents the background similarity map.

For the purpose of finding weaker feature points in the support feature map, we adopt a cross-entropy loss on SSP module, which is written as:

where ˆM bg s = 1 - ˆM fg s , and the superscript fg and bg is the indicator for the foreground and background respectively.

Finally, we calculate the position of weak feature points and normal ones borrowing information from both predicted support mask ˆMs and support ground truthMs. This process can be mathematically expressed as:

where the superscript hf, hb, nf, nb refers to hard foreground, hard background, normal foreground and normal background respectively.

3.4 Hard Prototypes Generation

As opposed to previous works that employ multiple normal prototypes, which inevitably loses part of the detailed information and inadequately represents the overall feature distribution within a class, we devise a novel Hard Prototypes Generation (HPG) module to generate multiple hard prototypes for both foreground and background to guarantee full use of the support information. We will take the example of generating hard foreground prototypes as described below.

Building on the weak feature points instruction provided by the SSP module, we firstly remove all support feature points except for those belonging to hard foreground, which can be described as:

where ⊙ denotes element-wise multiplication.

Next, we reshape the features of hard foreground region, denoted as Dhf = [d1hf , d2hf , ..., dNhf hf ] ∈ RNhf×C′ . Then, a random sampling operation is performed on Dhf to extract some hard foreground points D′hf . This process can be expressed as:

where F(·) is a random sampling function and N is the number of D′hf .

Subsequently, we populate the hard foreground feature map F hf s withD′hf whereM hf s = 0 and retain original weak foreground feature points, which can be summarized mathematically in the following formulation:

where (h, w) denotes the spatial position of the pixel.

After that, we acquire a reconstructed hard foreground feature map where all feature points exclusively pertain to the hard foreground. In order to generate multiple hard foreground prototypes that fully capture the feature information, we pass the reconstructed F hf s through a lightweight MLP which is composed of a fully connected layer, a ReLU activation layer, and another fully connected layer. This process can be denoted as:

where MLP (·; ϕ) represents the MLP module parameterized with ϕ. The hard foreground feature map is reshaped to the dimension of HW × C ′ before being fed into the MLP and Phf ∈ RNhf×C′ denotes generated Nhf hard foreground prototypes.

Notably, the aforementioned procedure for generating hard prototypes also applies to hard background, normal foreground and background, namely Phb, Pnf , Pnb. Eventually, we concatenate the generated hard prototypes with normal ones and corresponding global support prototype for foreground and background respectively, which can be formu-

where ⊕ denotes the concatenation operation. pfg s and pbg s are support foreground and background prototypes generated via MAP. Pfg ∈ RNfg×C′ and Pbg ∈ RNbg×C′ are final representative prototypes used for predicting the foreground and background similarity maps of the query image respectively.

3.5 Multiple SimilarityMaps Fusion

In previous FSMIS tasks, researchers commonly leverage a single prototype or multiple normal prototypes to make predictions for the query image directly based on cosinesimilarity. However, there exist extreme imbalance between foreground and background in medical image and lost information in the masked feature, which easily lead to misclassified foreground and background points. To tackle this challenge, we design a Multiple Similarity Maps Fusion (MSMF) module, which fuses similarity maps generated by multiple prototypes obtained from HPG module, thus making more accurate segmentation.

First of all, we calculate the similarity maps between the multiple foreground and background prototypes, and the query image respectively:

where Sfg ∈ RNfg×H′×W ′ and Sbg ∈ RNbg×H′×W ′ denote the generated similarity maps for foreground and background of the query respectively.

Then, both Sfg and Sbg will be fed into a lightweight decoder to fuse multiple similarity maps, thereby generating the foreground prediction map and background prediction map. This process can be denoted as:

where D1(·; ψ1) and D2(·; ψ2) denote a straightforward decoder based on convolutional neural network with parameter ψ1 and ψ2. ˆM fg q ∈ R1×H′×W ′ and ˆM bg q ∈ R1×H′×W ′ are the final single predicted similarity map for query foreground and background respectively.

Finally, we upsample ˆM fg q and ˆM bg q to the size of (H,W) and integrate them by utilizing the concatenation operation and a softmax layer to get the final prediction, as follows:

where sof tmax(·) executes along the channel dimension and ˆMq ∈ R2×H×W will eventually be processed by an argmax operation to get the final segmentation mask of the query image.

3.6 Loss Function

In order to quantify the dissimilarity between the predicted and ground truth, we employ the binary cross-entropy loss as the segmentation loss for each training episode:

where Lseg works as the primary loss of our proposed CoW.

Following common practice [Wang et al., 2019; Ouyang et al., 2020], we inversely predict labels of the support images by using query images as the support set, then build a prototype alignment regularization loss, which is calculated as follows:

Additionally, we employ the intra-class loss to help diminish intra-class variation of the support and query and the interclass loss to improve discrimination among different classes, both enhancing the representativeness of our generated hard and normal prototypes for foreground and background, which is formulated as:

where the superscripts s and q denote the support and query image, cos is the cosine similarity, and max(·) means taking the maximum.

Furthermore, to get a more precise predicted mask, we devise a boundary loss to refine the segmentation boundary. Specifically, we generate the boundary map of ground truth and predicted mask through a maxpool operation, denoted as M b q and ˆM b q . Then, we calculate f1-score between them and get final boundary loss Lbound as follows:

where BF1(·, ·) computes the f1-score between input pair.

Overall, the loss function for each training episode is defined to be L = Lseg + Lalign + λ0Lbound + λ1(Lintra + Linter + Lssp), where λ0, λ1 are the balancing coefficients.

4 Experiment

4.1 Datasets, Settings and EvaluationMetrics

Datasets. To demonstrate the superiority of our method under different scenarios, we perform evaluations on three representative publicly available datasets: abdominal organs segmentation for MRI and CT (Abd-MRI [Kavur et al., 2021] and Abd-CT [Landman et al., 2015]) and cardiac segmentation for MRI (CMR [Zhuang, 2018]). The detailed introduction and split strategy is depicted in [Zhu et al., 2023].

Settings. Themajority of FSMISmodels follow the settings in the work [Ouyang et al., 2022] to evaluate their performance. In order to make a fair comparison, we also adopt the two experimental settings outlined in SSL-ALPNet [Ouyang et al., 2020] and ADNet [Hansen et al., 2022].

EvaluationMetrics. Following the common practice [Roy et al., 2020] in FSMIS, we employ the mean Sorensen-Dice coefficient (DSC) to compare the prediction to the ground truth. It measures the effectiveness of the same area between the predicted mask and the ground-truth.

4.2 Implementation Details

Following the practice in SSL-ALPNet [Ouyang et al., 2020], the experiments are conducted based on a 1-way 1-shot setting with 5-fold cross-validation. Initially, pseudo labels are generated for training by leveraging 3D superpixel clustering. After that, we reshape the size of the 2D slice from the 3D scanned image to 256 × 256 for Abd-MRI and CMR and 257 × 257 for Abd-CT. The total number of generated hard and normal prototypes for foreground and background are set to 100 and 600 respectively. We set the initial learning rate to 0.001, the decay rate to 0.9, and the total number of required iterations to 50K, containing 5000 iterations each epoch. The balancing coefficients λ0 and λ1 are set to 0.5 and 0.3 respectively in our experiment.

4.3 Comparison with SOTAMethods

For the purpose of convincingly demonstrating the superiority and effectiveness of our proposed CoW, we compare our method with several existing state-of-the-art methods, including PA-Net [Wang et al., 2019], SSL-ALPNet [Ouyang et al., 2020], SR&CL [Wang et al., 2022], CAT-Net [Lin et al., 2023], Q-Net [Shen et al., 2023], RPT [Zhu et al., 2023], GMRD [Cheng et al., 2024], DSPNet [Tang et al., 2024] and MSFSeg [Zheng et al., 2024], from the perspective of segmentation performance under two experimental settings. In this paper, considering that most researchers commonly adopt the PA-Net [Wang et al., 2019] as the baseline and the settings in SSL-ALPNet [Ouyang et al., 2020] to evaluate their performance, we also follow them to ensure a fair comparison.

As shown in Table 1 and Table 2, our CoW significantly outperforms all listed methods in terms of the mean dice score on three different datasets.

4.4 Ablation Studies

Effect of each component. To evaluate separate contributions of each component within our model, we present their positive contributions in Table 3. Our SSP module and HPG module significantly improve its performance on top of the baseline model [Wang et al., 2019], achieving a mean dice score of 81.44%, which is attributed to the combination of normal and hard prototypes carrying rich subject and boundary information. In addition, the MSMF module with a dualpath architecture that aims at fusing multiple similarity maps further improves the performance by 2.66%.

Effect of each loss. As can be seen from the last column in Table 4, each loss indeed makes a separate contribution to the final performance. Lintra and Linter achieve a total gain of 1.27% by improving discriminability of generated multiple prototypes. Particularly, Lbound gives a 1.83% increase due to its constraint on the segmentation boundary. Moreover, Lssp brings an additional gain of 0.71% for the reason that we can acquire more robust normal and hard prototypes making support self-prediction more precise.

Number of normal and hard foreground prototypes. We test the effect of the number for Phf and Pnf on the final performance of CoW. As shown in Table 5, the optimal number setting is half for each. By contrast, employing Phf or Pnf alone leads to a performance degradation of 2.01% and 1.04% respectively because they solely leverage singular subject or boundary information, thus causing a less comprehensive representation for each category.

Generatingmultiple prototypes. We compare our method with several previous works [Ouyang et al., 2020; Yang et al., 2020; Hartigan and Wong, 1979; Cheng et al., 2024] in Table 6. It is apparent that our CoWachieves the best performance, surpassing the suboptimal GMRD by 3.97%, for the reason of focusing more attention to weak boundary details.

4.5 Visualization

Qualitative Results. To better demonstrate the superiority of our method, qualitative results of our model and other models, including Q-Net [Shen et al., 2023] and GMRD [Cheng et al., 2024], on Abd-MRI are shown in Fig. 3 (More qualitative results on the CMR dataset and the AbdCT dataset can be found in Appendix).

Distribution of Hard and Normal Prototypes. As depicted in Fig. 4, we visualize the generated hard and normal prototypes with t-SNE [Van der Maaten and Hinton, 2008]. It can be seen that normal foreground or background prototypes alone fail to well represent the entire distribution of their corresponding class. In contrast, the introduction of hard prototypes can compensate for such deficiency.

5 Conclusion

In this paper, we have proposed to mine hard prototypes by concentrating on weak features for few-shot medical image segmentation. It alleviates the challenge of current FSMIS methods that generate multiple prototypes solely through random sampling and local averaging, often failing to give a complete representation of corresponding category distribution. Additionally, our model includes a powerful edge constraint, which can further ensure a clearer segmentation boundary. Numerous experimental results also demonstrate that our method can capture more comprehensive class distribution and precisely segment the intricate organs in medical images. We believe that our model has the potential to assist in the diagnosis of rare disease and improve the efficiency of medical treatment.

Acknowledgments

This work was partly supported by the National Natural Science Foundation of China (NSFC) under Grant No. 62371235, and partly by the Key Research and Development Plan of Jiangsu Province (Industry Foresight and Key Core Technology Project) under Grant BE2023008-2.