Few-Shot Medical Image Segmentation with High-Fidelity Prototypes

Currently, the deep neural network approaches dominate the medical image segmentation field. The early phase shares models with the natural image semantic segmentation. Fully Convolutional Networks (FCNs) [24] first equipped vanilla Convolutional Networks (CNN) with a segmentation head by introducing Up-sampling and Skip layer. Following that, the encoder-decoder-based methods [3, 28] are developed. Unlike the coarse reconstruction in FCNs, the symmetrical reconstruction of the decoder can capture much richer detailed semantics. With the application of deep learning in the medical field, the medical image-specific models merge correspondingly, among which U-Net [28] is extensively recognized for its superior performance. Besides symmetrical encoder-decoder architecture, U-Net infuses the skipped connections to facilitate the propagation of contextual information to higher resolution hierarchies. Inspired by it, several variants of U-Net are designed, including U-Net 3D [5], Atten-U-Net [29], Edge-U-Net [2], V-Net [27] and Y-Net [26].

These segmentation models mentioned above only work in a supervised fashion, relying on abundant expert-annotated data. Thus, they cannot apply to the few-shot setting where we need to segment an object of an ”unseen” class as only a few labeled images of this class are given.

The key to solving FSS is building a class-wise similarity between the query and support images. Following this view, the existing methods can be divided into three categories. The first category constructs support images-based guidance [35, 34, 44, 39]. For instance, [35] developed a two-branch approach where a conditioning branch imposes controlling on the logistic regression layer of the segmentation branch. [34] introduced squeeze and excitation blocks into the conditioning branch to encourage dense information interaction between the two branches. The second category designed novel network modules, e.g., attention modules [14, 41] and graph networks [11, 45], for discriminative representations, by which the features shared by query and support images were identified. The third mainstream is prototypic network [42, 30, 7, 9] that prototypes bridge the similarity computation in a meta-learning fashion. Here, the prototypes are specific features with semantics extracted from the support images. Recently, PANet [42] achieved impressive performance on the natural image segmentation task, performing dual-directive alignment between the query and support images. SSL-PANet [30] transferred the PANet architecture into the medical image segmentation where self-supervision with superpixels and local representation ensure the unsupervised segmentation. Following that, anomaly detection-inspired methods enhanced the performance by introducing self-supervision with supervoxels[12] or learning mechanism of an expert clinician [36].

Our DSPNet belongs to the third category above, having two significant discrepancies. Compared with methods working with classes-abundant natural images, e.g., PANet, DSPNet is a medical image-specific model with limited labelled support data. On the other hand, DSPNet considers the limitation of local information loss from a new view of detail self-refining, which is totally different from the existing prototypical methods.

For few-shot semantic segmentation tasks, attention mechanism [40] is a popular technique to build the relationship between the support and query images. The existing approaches can be divided into two categories: (i) graph-based [47, 41, 11] and (ii) non-graph-based [14, 48, 25]. The methods belonging to the first category employing a graph model to activate more pixels, such that the correspondence between support and query images is enhanced. For example, [11] fused the graph attention and the last layer feature map to generate an enhanced feature map to solve the problem of foreground pixel loss in the attention map. The core idea of the second category methods is to build the correspondence based on feature interaction between the support and query data, e.g., multi-scale contextual features [14], affinity constraint [48], mix attention map [25].

Unlike the spatial attention-based methods above, both FSPA and BCMA in DSPNet are channel attention-like methods. For FSPA, the channel-wise fusion ensures the deeper semantic fusion from local to global, whilst in BCMA, the detail self-refining relies on the channel structural information.

In the FSS field, Resemblance Attention Network (RAN) [43] is a classic module to integrate the support and query features [10, 15, 6]. In the proposed DSPNet, RAN engages in filtering irrelevant texture and objects between $F_{s}$ and $F_{q}$. Fig. 3 presents its network architecture. When support and query feature maps $F_{s}$, $F_{q}$ are input, they are first reshaped to feature vector $A_{s}$ and $A_{q}$, respectively. After that, in a Query-Key-Value attention manner with residual connection, the $A_{s}$, $A_{q}$ are fused to $\hat{F}_{s}$ where ${\rm{\bf Q}}={\rm{\bf V}}=A_{s}$, ${\rm{\bf K}}=A_{q}$. The process can be formulated by Eq. (1).

where $\phi(\cdot)$ stands for softmax operation, $\times$ means matrix multiplication, ${\phi}\left({A}_{s}^{T}\times{A}_{q}\right)$ means the similarity-based probability matrix weighting $A_{s}$.

To obtain high-fidelity class prototype for the semantic foreground, we explore the local semantics in the foregroud and fuse them to form global semantics without local semantics loss. We accomplish this idea using the cluster-based detail prototypes and channel-wise attention with local semantics guidance.

Overview. As shown in Fig. 4(a), to get more local semantics, we first employ the superpixel-guided clustering method [18] to mine $N_{s}$ cluster prototypes, denoted by ${P}_{c}=\{{P}_{c}^{i}\}_{i=1}^{N_{s}}$ where ${P}_{c}^{i}\in{\mathbb{R}}^{1\times D}$, $D$ is the dimension of prototype. The intuitive fusion, e.g., vanilla weighting without prior knowledge [18], can obtain the global semantics but suffers from confusing detail semantics. Therefore, we propose an attention-like cluster prototype fusion to address this issue, implementing detail self-refining and foreground tailoring sequentially.

Attention-like cluster prototype fusion. As shown in the middle of Fig. 4 (marked by grey box), this attention can be implemented in the fashion of Query-Key-Value. Taking ${\rm{\bf Q}}={\hat{F}_{s}}$, ${\rm{\bf K}}={\rm{\bf V}}={{P}_{c}}$, we can summarize this module to the following equation.

where $\phi(\cdot)$ is softmax computation; operator and respectively means the computation for cosine similarity measurement and channel-wise prototype fusion, whose details are presented in the following.

Since the size of $\hat{F}_{s}$ and $P_{c}$ are different, computation of does not follow cosine similarity’s definition, but performing in prototype-wise. Specifically, each prototype in ${P}_{c}$, i.e., ${P}_{c}^{i}$, is used to compute similarity with the supporting feature maps $\hat{F}_{s}$ in a one-dimensional convolution manner, in which the convolution calculation is replaced by cosine similarity computation. Thus, the $N_{s}$ prototypes lead to $N_{s}$ similarity maps, which can be collectively written as $S_{s}=\hat{\rm{F}}_{s}~{}{\scriptsize\leavevmode\hbox to10.16pt{\vbox to10.16pt{\pgfpicture\makeatletter\hbox{\hskip 5.08159pt\lower-5.08159pt\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ }\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{ }\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\pgfsys@setlinewidth{0.4pt}\pgfsys@invoke{ }\nullfont\hbox to0.0pt{\pgfsys@beginscope\pgfsys@invoke{ }{ {{}}\hbox{\hbox{{\pgfsys@beginscope\pgfsys@invoke{ }{{}{{{}}}{{}}{}{}{}{}{}{}{}{}{}{{}\pgfsys@moveto{4.88159pt}{0.0pt}\pgfsys@curveto{4.88159pt}{2.69606pt}{2.69606pt}{4.88159pt}{0.0pt}{4.88159pt}\pgfsys@curveto{-2.69606pt}{4.88159pt}{-4.88159pt}{2.69606pt}{-4.88159pt}{0.0pt}\pgfsys@curveto{-4.88159pt}{-2.69606pt}{-2.69606pt}{-4.88159pt}{0.0pt}{-4.88159pt}\pgfsys@curveto{2.69606pt}{-4.88159pt}{4.88159pt}{-2.69606pt}{4.88159pt}{0.0pt}\pgfsys@closepath\pgfsys@moveto{0.0pt}{0.0pt}\pgfsys@stroke\pgfsys@invoke{ } }{{{{}}\pgfsys@beginscope\pgfsys@invoke{ }\pgfsys@transformcm{1.0}{0.0}{0.0}{1.0}{-2.52777pt}{-2.39166pt}\pgfsys@invoke{ }\hbox{{\definecolor{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@rgb@stroke{0}{0}{0}\pgfsys@invoke{ }\pgfsys@color@rgb@fill{0}{0}{0}\pgfsys@invoke{ }\hbox{{C}} }}\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope}}} } \pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope{{{}}}{}{}\hss}\pgfsys@discardpath\pgfsys@invoke{\lxSVG@closescope }\pgfsys@endscope\hss}}\lxSVG@closescope\endpgfpicture}}}~{}{P}_{c}\in{\mathbb{R}}^{(H_{s}\times W_{s})\times N_{s}}$ where $(H_{s},W_{s})$ is map size. For any map in $S_{s}$, denoted by $S_{s}^{i}$, its computation can be expressed as

where function ${\rm{sim1D}}~{}(\cdot,\cdot)$ stands for the similarity computation working in the one-dimensional convolution fashion. The value of ${{S}_{s}^{i}}$ at position $(h,w)$ is the cosine similarity between ${P}_{c}^{i}$ and ${\hat{F}}_{s}$ at position $(h,w)$. Namely,

where ${\hat{F}}_{s}{(h,w)}\in{\mathbb{R}}^{1\times D}$ represents the feature vector of the feature maps ${\hat{F}}_{s}$ at position $(h,w)$ along channel dimension.

To incorporate the knowledge represented by the similarity maps $S_{s}$ into the cluster prototypes $P_{c}$, we also adopt a one-dimensional convolution to implement the computation of , as illustrated in Fig. 4(b) and (c). Specifically, the computation begins with the channel-wise generation of convolution filters. Given that the cluster prototypes $P_{c}$ are arranged as shown in Fig. 4(b). We slice $P_{c}$ along the channel-dimension and obtain $D$ convolution vectors $\{K_{i}\}_{i=1}^{D}$ where $K_{i}\in{\mathbb{R}}^{1\times N_{s}}$ contains cluster prototypes’ semantic component on the i-th channel. After that, as done in Fig. 4(c), we conduct one-dimensional convolution to obtain fused maps ${\bar{F}_{s}}\in{\mathbb{R}}^{D\times H\times W}$. The computation for the $i$-th map can be expressed as:

where $\phi(\cdot)$ is softmax operation, $\phi({S}_{s})$ stands for probability map, $K_{i}$ work as the convolution filter. In this end, to suppress introduced noise in the fusion step, we tailor the fused map ${\bar{F}_{s}}$ to high-fidelity foreground prototype $P_{f}$ by Mask Average Pooling.

where $m_{s}$ is the given mask of the support image and resized to the same as ${\bar{F}}_{s}$.

Remark: In Eq. (5), $S_{s}$ is essentially semantic response maps concerning the cluster prototypes, such that probability map $\phi({S}_{s})$ is noticeably relational to the detail semantics. That is, this fusion ensured by computation is guided by the detail semantics represented in $S_{s}$. Equivalently, the fusion process preserves the detail semantics, as our expectation.

Besides, two designs differs FSPA from the previous work. First, FSPA reduces the mined cluster prototypes to a fused one instead of using these prototypes separately [8, 21]. Second, the proposed channel-wise attention leads to global semantics preserving the local semantics unlike spatially weighting [18].

Unlike the foreground taking the cluster prototypes as the local details, the background in medical images is usually semantic-less in a large scope. Therefore, in this paper, we do not mine from the spatial dimension but deem the structural information in the channel dimension as the local details. Within this context, we design a controllable channel attention mechanism to jointly model the channel-specific structural information and incorporate them into the raw background prototypes.

Overview. As illustrated in Fig. 5(a), BCMA begins with generating raw detail prototypes. By Average Pooling and reshaping, ${\hat{F}}_{s}$ is converted to $P_{n}\in{\mathbb{R}}^{(H\times W)\times D}$. Following that, the controllable multi-head channel attention module refreshes $P_{n}$ to high-fidelity background prototypes $P_{a}$. Finally, $P_{a}$ is reshaped to feature maps ${\widetilde{F}}_{s}$ and further tailored to the high-fidelity background prototypes $P_{b}$ by the background zone in pooled support mask ${M}_{r}$.

Controllable multi-head channel attention. The proposed channel attention mechanism encodes the channel-structural information into raw background prototypes in an element manner. For a raw prototype, their elements are independently refined by the structural information of different channels. We achieve this by the D-way architecture illustrated in Fig. 5(b). Suppose that for any raw prototype in $P_{n}$, denoted by $P_{n}^{k}$, the converted high-fidelity prototype is $P_{a}^{k}$. In the Q-K-V fashion, we set ${\rm{\bf Q}}\!=\!Q_{n}$, ${\rm{\bf K}}\!=\!P_{n}$, ${\rm{\bf V}}\!=\!P_{n}^{k}$ where $Q_{n}$ is transformed by channel-wise slicing $P_{n}$. In the proposed module, $P_{n}$ and $P_{n}^{k}$ are copied D times and inputted into the D heads, respectively. At the same time, $Q_{n}$ is inputted channel-wise. That is, the $j$-th head takes the $j$-th component $Q_{n}^{j}$ as the input. The multi-head module refining $P_{n}^{k}$ can be formulated as

where ${\rm{cat}}(\cdot)$ concatenates the input set to a vector according to their indices; ${\rm{h}}_{j}(\cdot,\cdot,\cdot)$ is the $j$-th channel attention head generating $P_{a}^{k,j}$ (the $j$-th element in $P_{a}^{k}$) that we elaborate as follows.

In Eq. (7), the objective of the attention head $h_{j}$ is encoding the $j$-th channel-specific structural information, denoted by $\boldsymbol{a}^{\prime}_{j}$, into the raw $P_{n}^{k}$. Under the attention framework, the encoding can be implemented by a weighting operation $P_{n}^{k}\times\boldsymbol{a}^{\prime}_{j}$, whilst the $\boldsymbol{a}^{\prime}_{j}$ generation is the core problem we need to address. For this issue, as depicted in Fig. 5(c), we provide a controllable design consisting of (i) global exploration module and (ii) sparse channel-aware regulating module. Among them, the former predicts the global channel structural information of the $j$-th channel $\boldsymbol{a}_{j}$, whilst the latter serves as a controller by injecting the $j$-th channel-specific adjustment $\boldsymbol{r}$. Thus, the working mechanism ${\rm{h}}_{j}$ can be formulated as

where parameter ${\boldsymbol{a}_{j}}$ is learnable; operator $\odot$ means element-wise multiplying.

In Eq. (8), the generation of adjustment $\boldsymbol{r}$ involves two blocks in the sparse channel-aware regulating module (see the two dark grey box in Fig. 5(c)). First, the channel similarity computation, formulated by Eq. (9), captures the dynamics of the relationship between $j$-th channel and other channels.

where $\boldsymbol{w}_{c}\in{\mathbb{R}}^{D}$ is the channel similarity whose $i$-th element is $\boldsymbol{w}_{c,i}$, function ${\rm{cossim}}(\cdot,\cdot)$ measures the cosine similarity of vector ${Q_{n}^{j}}$ over set ${Q_{n}}$, $\phi$ is softmax operation. Subsequently, the incorporation unit generates adjustment coefficients by high-lighting the sparse-relative channels indexed by masked frozen vector ${\boldsymbol{m}_{w}}$. This process can be formulated as

where rade-off parameter $\beta$ stands for the control strength.

As mentioned above, the proposed $h_{j}$ involves two important parameters, i.e., $\boldsymbol{a}_{j}$ (Eq. (8)) and $\boldsymbol{m}_{w}$ (Eq. (10)). In our design, both of them are initiated by a pre-set sparse vector $\boldsymbol{w}_{i}$ that represents a prior knowledge about the channel structural pattern. Specifically, at the beginning of model training, we set $\boldsymbol{m}_{w}={\rm{mask}}(\boldsymbol{w}_{i})$ and $\boldsymbol{a}_{j}=\boldsymbol{w}_{i}$ where function ${\rm{mask}}(\cdot)$ outputs Boolean vector whose locations of 1 corresponds to the non-zero places in input vector.

Remark: In our controllable attention mechanism, the core idea is imposing sparse channel-aware regulating to adjust the learnt global channel relation, leading to channel-specific structural information. Here, the sparse constraint is motivated by the ubiquitous sparse nature of neural connections, whose rationality is verified by much work [20, 4].

Also, from a methodological point of view, our structure can be understood as a piece of work of structural learning-based attention [33, 23, 13], but in the channel dimension. For instance, shifted window partitioning in Swin Transformer [23] introduces spatial relation constraint to self-attention. Similarly, our sparse channel-aware regulating introduces a channel structural constraint, i.e., sparse relation ($\boldsymbol{r}$), to the predicted global channel structural information ($\boldsymbol{a}_{j}$).

We regulate cross-entropy regularization to supervise this model training process:

where $\hat{m}_{q}^{j}(h,w)$ is the predicted results of the query mask label ${m}_{q}^{j}(h,w)$; in $\{f,b\}$, $f$ and $b$ means foreground and background, respectively. Also, following [42, 30, 36], we perform another inverse learning where the query images serve as the support set to predict labels of the support images. Thus, we encourage a prototypical alignment formulated by

Overall, for each training episode, the final objective of DSPNet is defined as follows:

To demonstrate the effectiveness of DSPNet, we conduct evaluation on three challenging datasets with different segmentation scenarios. Their details are presented as follows.

Abdominal CT dataset [17], termed ABD-CT, was acquired from the Multi-Atlas Abdomen Labeling challenge at the Medical Image Computing and Computer Assisted Intervention Society (MICCAI) in 2015. This dataset contains 30 3D abdominal CT scans. Of note, this is a clinical dataset containing patients with various pathologies and variations in intensity distributions between scans.

Abdominal MRI dataset [16], termed ABD-MRI, was obtained from the Combined Healthy Abdominal Organ Segmentation (CHAOS) challenge held at the IEEE International Symposium on Biomedical Imaging (ISBI) in 2019. This dataset consists of 20 3D MRI scans with a total of four different labels representing different abdominal organs.

Cardiac MRI dataset [50], termed CMR, was obtained from the Automatic Cardiac Chamber and Myocardium Segmentation Challenge held at the Conference on Medical Image Computing and Computer Assisted Intervention (MICCAI) in 2019. It contains 35 clinical 3D cardiac MRI scans.

In our experiment settings, to ensure fair comparison, we adopted the same image preprocessing solution as SSL-ALPNet [30]. Specifically, we sampled the images into slices along the channel dimension, and resized each slice to 256$\times$256 pixels. Moreover, we repeated each slice three times along the channel dimension to fit into the network. We employ 5-fold cross-validation as our evaluation method, where each dataset is evenly divided into 5 parts.

To evaluate the performance of the segmentation model, we utilized the conventional Dice score scheme. The Dice score has a range from 0 to 100, where 0 represents a complete mismatch between the prediction and ground truth, while 100 signifies a perfect match. The Dice calculation formula is

where $A$ represents the predicted mask and $B$ represents the ground truth.

To evaluate the model’s performance, we follow the experimental settings in [30, 12], considering two cases. Setting-1 is the initial setting proposed in [34], where test classes may appear in the background of training images. We train and test on all classes in the dataset without any partitioning. Setting-2 is a strict version of Setting-1, proposed in [30], where we adopted a stricter approach. In this setting, test classes do not appear in any training images. For instance, when segmenting Liver during training, the support and query images do not contain the Spleen, which is the segmenting target for testing. We directly removed the images containing test classes during the training phase to ensure that the test classes are truly ”unseen” for the model.

We implemented our model using the Pytorch framework with a pre-trained fully convolutional Resnet101 model as the feature extractor. The Resnet-101 model was pre-trained on the MS-COCO dataset. Given that the superpixel pseudo-labels contain rich clustering information, which are helpful to alleviate the annotation absence. We generate the superpixel pseudo-label in an offline manner as the support image mask before starting the model training, following [30, 12, 36].

In DSPNet, there is one hyper-parameter: Local adjustment intensity $\alpha$ in Eq. (8). As another important factor, the sparse pattern of $\boldsymbol{w}_{i}$ follows the neighbour channel constraint, namely, $\boldsymbol{w}_{i}=\left[\boldsymbol{0},w_{1},w_{2},w_{3},\boldsymbol{0}\right]$ where $w_{2}$ is the $j$-th element of $\boldsymbol{w}_{i}$, $({w}_{1},{w}_{2},{w}_{3})\!\in\!\left[0,1.0\right]$, ${w}_{2}\!>\!{w}_{1}={w}_{3}$. Specifically, in the ABD-MRI dataset, Setting-1 adopts $\beta\!=\!0.3,(w_{1},w_{2},w_{3})\!=\!(0.2,0.8,0.2)$, whilst $\beta\!=\!0.2,(w_{1},w_{2},w_{3})\!=\!(0.3,0.6,0.3)$ are used in Setting-2. In the ABD-CT dataset, Setting-1 adopts $\beta\!=\!0.2,(w_{1},w_{2},w_{3})\!=\!(0.3,0.7,0.3)$, and $\beta\!=\!0.4,(w_{1},w_{2},w_{3})\!=\!(0.1,0.7,0.1)$ are used in Setting-2. For the CMR dataset, Setting-1 selects $\beta\!=\!0.3,(w_{1},w_{2},w_{3})\!=\!(0.1,0.9,0.1)$.

For our experimental results, we used stochastic gradient descent algorithm with a batch size of 1 for 100k iterations to minimize the objective in Eq. (13). The self-supervised training took around 4.5 hours on a single Nvidia TITAN V GPU, and the memory consumption was about 8.1GB.

To evaluate our approach, we compared it with six state-of-the-art medical image semantic segmentation methods, including SE-Net [34], PANet [42], SSL-ALPNet [30], Q-Net [36], and CAT-Net [19]. Among them, SE-Net belongs to the category of constructing support images-based guidance, whilst the rest comparisons all follow the clue of prototypic network. For a fair comparison, we obtain the results of all prototype-based methods, i.e., PANet, SSL-ALPNet, Q-Net and CAT-Net, by re-running their official codes on the same evaluating bed with DSPNet. The results of SE-Net are cited from the publication.

The same as the previous methods, we perform the evaluation on ABD-MRI and ABD-CT under both Setting-1 and Setting-2 whilst the CMR is based on Setting-1. Tab. 1 reports the results in Dice score on ABD-MRI and ABD-CT. The results showed that DSPNet outperforms the previous methods in the two settings. On ABD-MRI dataset, compared with the second-best method Q-Net in mean score, DSPNet achieves an improvement of 2.6 under Setting-1. Meanwhile, as for the strict Setting-2 testing model for ”unknown” classes, DSPNet demonstrates impressive performance with 7.8 increase, especially with a dice score of approximately 82 for Right kidney. The reason is discussed in Section 5.7: Ablation study. On the ABD-CT dataset, in average score, DSPNet also surpasses the second-best method SSL-ALPNet by 1.2 in Setting-1 and 2.1 in Setting-2, respectively. For an intuitive observation, we present the visual segmentation results in Fig. 6. As shown in this figure, DSPNet has much better segmentation for large objects (see Liver), while predicting the finer boundary for small objects (see Spleen).

Tab. 2 shows the comparison results on CMR with adjacent organs. In this scenario, DSPNet exhibited better segmenting performance in all three classes, obtaining 2.0 improvement in mean score compared with the previous best method SSL-ALPNet. The right side of Fig. 6 depicts three toy experimental results. It is seen that the DSPNet can generate complicated boundaries (see LV-MYO, RV), implying more details are captured by DSPNet compared with the previous methods. For the objects with relatively regular shapes, e.g., LV-BP, DSPNet achieves fuller segmentation near the boundary.

As illustrated in the middle of Fig. 2, DSPNet involves three components, i.e., RAN, FSPA and BCMA. In this part, we carry out an ablation study to isolate their effect as follows. All experimental results are obtained based on the ABD-MRI dataset under strict Setting-2.

This work is partly funded by the German Research Foundation (DFG) and National Natural Science Foundation of China (NSFC) in project Crossmodal Learning under contract Sonderforschungsbereich Transregio 169, the Hamburg Landesforschungsförderungsprojekt Cross, NSFC (61773083); NSFC (62206168, 62276048, 52375035).