Iterative Few-shot Semantic Segmentation from Image Label Text

1 Introduction

In recent years, semantic segmentation has reached remarkable progress. However, a large amount of data is required for training such a task to segment unseen class objects. It differs from the human perception that a human is able to handle novel objects after learning from only a few samples.

Few-shot semantic segmentation is proposed to mimic the property of humans. The segmentation model is expected to segment objects with an unseen class in the query image with a few support images and pixel-wise annotations, as shown in Figure 1 (a). Although it is more like humans and helps reduce the work of data annotation, pixel-wise mask annotation is still unavoidable. In general, compared with image-level or box-level annotation, generating the precise pixel-level label is more time-consuming. Therefore, once faced with unseen categories, it takes much time to prepare the dense annotations, which restricts the practical application in the real-time open world, such as robotic engineering or consumer electronics products.

∗Equal contribution.

†Corresponding authors.

In this work, we concentrate on the task of few-shot segmentation with image-level label support, namely weakly supervised few-shot semantic segmentation, as depicted in Figure 1 (b). We argue that pixel-wise examples are not necessary for pixel-wise recognition, as humans are able to imagine the shape and structure of objects with the help of prior knowledge such as attribute and low-level feature. Therefore, it is meaningful to study the weakly supervised setting.

However, due to the severer data scarcity, i.e., the lack of precise support mask, weakly supervised few-shot semantic segmentation is rarely explored. [Raza et al., 2019] simplify the problem into the condition where only one novel class exists in the support image, but the model would collapse in the more complex scene. Recent works employ the Word2Vec embedding to associate the image feature with the class label explicitly [Lee et al., 2022] or implicitly [Siam et al., 2020]. Nevertheless, the performance still falls far behind that in the standard few-shot setting.

We summarize that the performance gap is mainly because a model designed for weakly supervised few-shot segmentation is faced with two daunting challenges. The first challenge is to locate specific objects in the support image. Given a novel class label, the model has to filter out the corresponding objects of unseen categories from the background and other foreground instances in the support images. It always results in an imprecise or even wrong support mask, rather than the ground-truth mask. Therefore, another challenge is to learn from imprecise labels. In other words, the model has to adapt to the imprecise input support masks.

In this work, we propose a two-stage framework to address the challenges. In the first stage, we tackle the first challenge, i.e., locating specific objects guided by the image label text. We fall back on the recent progress of vision-language model CLIP [Radford et al., 2021], which is trained with massive image-text pairs from the web in a self-supervised way. CLIP encodes images and texts into the joint embedding space, and thus it builds a strong correlation between them. To leverage its potential, we inject the class name into a simple prompt template, as shown in Figure 1 (c), and compute the cosine similarity of the output from vision and text encoder in CLIP. This process could be regarded as the classification task, and hence we can generate the support mask according to the high-response regions of its Grad CAM [Selvaraju et al., 2017].

Then in the second stage, we aim at handling the second challenge, i.e., the imprecise input support mask. We experimentally find that due to the robustness of the well-trained few-shot segmentation model, when segmenting the query mask with an imprecise support mask, the quality of the predicted query mask is much higher than that of the input support mask. Based on this observation, we propose to refine the support and query mask iteratively and mutually. Specifically, we continue to refine the support mask with the help of a higher-quality predicted query mask, and then update the query mask with the better support mask in turn. However, due to the varying quality of the input object masks, the regular few-shot segmentation model would struggle in judging the confidence of the input mask, which leads to a suboptimal result. Hence we propose a progressive mask refinement module Iterative Mutual Refinement (IMR), to capture and adapt to the varying input distribution, i.e., howwell the input mask has been estimated. During the refinement, IMR provides an effective strategy to leverage the information of all past steps with memory, so that the support and query mask could be optimized adaptively. We incorporate the IMR into the pipeline of few-shot segmentation, so as to better explore the potential of the iterative refinement. Our contributions are summarized as follows:

• We propose a two-stage framework for weakly supervised few-shot segmentation, consisting of initial mask generation from the image label text and iterative support and query mask refinement.
• We present an Iterative Mutual Refinement module to progressively refine the coarse masks, which is better for the case where the input mask quality varies. It is designed to be model-agnostic, and could serve to extend common few-shot segmentation models to the weakly supervised setting with minor modifications.
• Extensive experiments show that our method could not

only outperform the state-of-the-art weakly supervised counterparts by a significant margin, but also close to or beyond recent regular few-shot segmentation methods.

2 RelatedWork

Few-shot Semantic Segmentation. The dominant approaches of few-shot segmentation are devised based on meta-learning. The information of support images is aggregated into the query feature for further prediction in different ways. Most methods achieve this by first extracting single [Tian et al., 2020; Xie et al., 2021] or multiple [Zhang et al., 2021; Li et al., 2021a] class prototypes from support feature by masked average pooling, then concatenating the prototypes and query feature as the input of a welldesigned decoder. However, this simple strategy is faced with severe support information loss unavoidably. Some recent works alleviate this by exploring the potential of the pair-wise feature correlations between support and query feature, and the more powerful transformer architecture [Sun et al., 2021] or 4D convolutions [Min et al., 2021] is thus utilized to boost the performance further. Despite these signs of progress, few-shot segmentation requires the precise support mask during testing, which is superfluous and hardly available in customer-oriented real-time applications. In this work, we tackle this problem by replacing the ground-truth mask with the image label text.

Weakly Supervised Few-shot Segmentation. Weakly supervised few-shot segmentation is rarely explored due to the severer challenge of data scarcity. [Wang et al., 2019; Zhang et al., 2019] utilize bounding box as the annotation of support images to segment the query images during testing. However, their performance is suboptimal as these models are not specifically designed for the weakly supervised condition. The first few-shot segmentation method with imagelevel support annotation is proposed by [Raza et al., 2019], where they generate the pseudo masks of support images by retaining the pixels not classified as background. Nevertheless, it cannot deal with the support images containing more than one novel class. Recently, [Siam et al., 2020] incorporate the semantic word embedding of the novel category into the visual feature, and exploited a multi-stage attention mechanism to predict the query mask. [Lee et al., 2022] regard the pseudo Classification Activation Map (CAM) of the unseen classes as the support mask and ground truth query mask during training, which is estimated by the CAMs of semantically similar categories. However, the performance of these approaches is far away from that using a precise support mask, which restricts its practical application. On the contrary, our method could perform close to or even beyond the state-of-the-art few-shot segmentation models, and show strong generalization ability for images in the wild.

CLIP in Classification and Segmentation. CLIP [Radford et al., 2021] is a milestone of vision-language pretraining model, which learns a multi-modal embedding space by jointly encoding images and texts from the web in contrastive learning. CLIP shows strong capacity in zero-shot or fewshot recognition [Gao et al., 2021; Radford et al., 2021], and recent works have focused on the extension into dense prediction tasks like semantic segmentation. [Rao et al., 2021] first employ CLIP in fully supervised semantic segmentation in a model-agnostic way, and [Zhou et al., 2021] extend CLIP into zero-shot segmentation. [Li et al., 2022] combine CLIP with a well-designed transformer to further improve the precision of segmentation. However, these methods rely heavily on complicated prompt engineering, and the upper bound of their segmentation performance is limited due to the lack of support images. Therefore, in this work, we explore the potential of CLIP in few-shot segmentation for the first time with a simple text prompting only, along with an Iterative Mutual Refinement module to obtain the precise prediction.

3 ProblemDefinition

Few-shot semantic segmentation aims at training a model which can segment unseen class objects given extremely limited annotated samples with the same class. Formally, the dominant episode based meta-learning paradigm employs a base dataset Dbase with category set Cbase for training and a novel dataset Dnovel with unseen category set Cnovel for testing. During training, a large number of episodes are sampled from Dbase. Each episode consists of a support set S and a query set Q for a specific class c ∈ Cbase as follows:

where I∗ and M∗ denote image and mask respectively. After training, the model f(., θ) learned to predict a binary mask ˆMq for query image Iq with the guidance of K support images and masks for a specific novel class c ∈ Cnovel,

In the standard few-shot segmentation setting [Min et al., 2021; Tian et al., 2020], both support masks Ms and query masks Mq for class c are available for training, while query masks are inaccessible for testing. On the contrary, we focus on a more challenging setting, named weakly supervised few-shot segmentation, in which even support masks are inaccessible in testing. Thus, the prediction becomes:

Although a weaker annotation replaces the pixel-wise annotated mask, we show that it can still reach a comparable or even better performance.

4 Method

We propose a framework for weakly supervised few-shot segmentation. The pipeline is shown in Figure 2, which contains a CLIP-CAM module to generate coarse masks for support and query images (Section 4.1), an Iterative Mutual Refinement module to refine the coarse masks progressively (Section 4.2), and some basic blocks from previous well-designed models. Notably, our framework is plug-and-play and modelagnostic, and we show how to extend regular few-shot segmentation models to the weakly supervised setting in Section 4.3.

4.1 Coarse Mask fromLabel Text

In the weakly supervised setting, dense annotations of supported images are not available, so we propose to generate segmentation maps as an alternative from images and class labels only. To some extent, this subtask degrades into zeroshot segmentation, for which the goal comes to segment novel objects in an image with the class label only. Nevertheless, most existing zero-shot approaches require complex architectures and training strategies, which might be inefficient to serve as only one stage of our whole pipeline.

Recent vision-language models, especially CLIP[Radford et al., 2021], have shown surprisingly powerful performance in associating visual and text concepts. CLIP consists of a vision encoder Ev(·) and a text encoder Et(·), where the image and text are encoded into fv, ft ∈ R1×1×C . Since it learns a joint embedding space for images and texts, one can simply compute the cosine similarity of fv and ft for image-level zero-shot classification without extra training and modules.

In this work, we extend it to pixel-level zero-shot segmentation in a training-free way.

The original architecture of CLIP is not suitable for dense prediction tasks, so a naive idea is to adjust the architecture of the vision encoder Ev(·). Specifically, similar to [Zhou et al., 2021], we drop the query and key embedding layers, and reformulate the value embedding layer and the last linear layer into two convolutional layers, so as to project the image feature of each position to the embedding space of the text feature. After that, we can conduct pixel-level prediction by computing the cosine similarity between text feature and dense visual feature.

An alternative is the heatmap of Grad CAM [Selvaraju et al., 2017], which aims at looking for the region contributing most to the classification. Since CLIP computes the cosine distance between fv and ft, it could be regarded as the classification output of a certain category, where ft works as the weight of the classifier. Thus we utilize the heatmap output by Grad CAM as the coarse object mask, as the region with a high response makes fv resemble ft.

With the guidance of the class label text, both strategies above could generate the coarse object masks from the images. We select the latter one due to its better performance. Please refer to the supplementary materials for more details of the implementation and comparison. It is worth mentioning that it is unnecessary to pursue a precise initial mask, since we adopt an iterative strategy to refine the estimation of the mask. Therefore, we do not depend on the complicated prompt engineering and other various zero-shot methods, and still achieve a satisfying performance.

4.2 Iterative Mutual Mask Refinement

With a coarse support mask ˆM 0 s , an intuitive solution is to feed it into the few-shot segmentation model to obtain a more delicate query mask. Similarly, the finer query mask could then result in the finer support mask, which is more suitable as the indeed support input. Therefore, an iterative strategy could improve the final prediction gradually by support-query mutual refinement, as illustrated in Figure 3.

Nevertheless, the distribution of input masks from different steps varies, which degrades the performance of the model. Therefore, we propose an Iterative Mutual Refinement (IMR) module to capture and adapt to various input distributions.

Specifically, we first introduce a hidden state h to reflect how well the input mask is estimated. Thus in the t step, we already have the image feature Fs, Fq, the prior estimation of support and query masks ˆM t-1 s , ˆM t-1 q , as well as the stored previous state ht-1 s , ht-1 q . In order to predict ˆM t q , we first encode the information of support images and masks into the query feature, and obtain the correlation volume C t q , which is consistent with the regular few-shot segmentation. Then Fq , C t q , ˆM t-1 q are concatenated as xtq , along with ht-1 q as the input to update the htq . The whole process of state updating derived from the gated recurrent unit is formulated as:

where Wconv1,1,Wconv1,2,Wconv2 are implemented by 1 × 1 convolutions followed by 3 × 3 group convolutions with 16 groups, and [·],⊗, σ denote concatenation, Hadamard product and sigmoid function, respectively. For briefness, we only list how htq is updated, and hts is updated in the symmetrical way. After that, the updated htq , as well as the prior estimation M t-1 q , are utilized for further prediction. Detailed analysis is provided in the supplementary materials.

4.3 Model Architecture and Objective

Our main pipeline consists of the Iterative Mutual Refinement module, backbone, neck, and head. The last three modules are abstract from the regular few-shot segmentation methods according to their respective roles, and could be replaced with corresponding modules in different basic methods.

In this work we reuse the architecture of HSNet [Min et al., 2021] for its excellent performance. Concretely, we adopt ResNet50 [He et al., 2016] or VGG16 [Simonyan and Zisserman, 2015] as backbone to extract feature Fs, Fq. Then we adopt the well-designed 4D convolutions followed by multiscale feature fusion as the neck to model the correlation of Fs and Fq, and obtain the intermediate output Cq. Please refer to [Min et al., 2021] for more details.

For the mask prediction, we make a minor modification based on the original HSNet head, in order to leverage the prior mask estimation ˆM t-1 , which is formulated as follows:

where Head0 is the original head in HSNet, and Wh1,Wh2 is implemented by 1 × 1 convolutional layers.

In order to accelerate convergence, we add the intermediate supervision to the output mask of each step. As a result, the overall loss function is defined as follows:

where N denotes the number of the total steps in the refinement, and Ms,Mq are the ground-truth support and query mask. We adopt cross entropy loss Lseg during training and λt and ws are the hyperparameters.

5 Experimental Results

5.1 Datasets and EvaluationMetrics

We follow almost the same settings with regular few-shot segmentation, and the only difference is that the ground-truth masks of support images are replaced by the label text.we evaluate our framework in two widely-used datasets, i.e., Pascal-5i [Shaban et al., 2017], COCO-20i [Lin et al., 2014], and conduct the cross-validation over all the folds in each dataset. For each fold i, samples in the remaining folds serve as the training data, while the 1000 episodes (support-query pairs) are sampled randomly as the testing data. We adopt the mean intersection over union (mIoU) to measure the performance, which averages the IoU of all classes in the testing data to avoid the bias to the specific categories.

5.2 Implementation Details

We iteratively refine the estimation of support and query masks twice in Pascal-5i , and thrice in COCO-20i due to the more complex scene. It is actually a trade-off between effectiveness and efficiency, as we empirically find that the gain of performance gradually converges with the step number increasing. For a fair comparison, we utilize the ImageNet pretrained VGG16 and ResNet50 as our backbone to extract feature, and the pre-trained CLIP model for initial mask generation. For the CLIP model, the vision and text encoder are a modified ResNet50 and transformer, respectively. The parameters of the backbone and CLIP are fixed during training to avoid the undesirable representation bias towards the base class.

Owing to our refinement in the second stage, we adopt the most superficial prompt text "a photo of a [class]" for the input to the text encoder, where [class] is replaced by the label of specific categories.

Consistent with HSNet, we resize the input into 400 × 400 in both the training and testing stage, and we do not adopt any data augmentation strategies. The network is trained with Adam optimizer with learning rate 2e-4 , and all hyperparameters λt and ws are set to one by default. However, our method is robust to the choice of these hyperparameters, as shown in the supplementary materials.

5.3 Comparison with State-of-the-art

In this subsection, we compare our method with the state-ofthe-art weakly supervised and regular few-shot segmentation models in the 1-shot setting. Table 1 presents the results in the Pascal-5i . It could be seen that our method sets new state of the arts with both VGG16 and ResNet50 backbone in the weakly supervised setting, and outperforms its counterparts with either box-level or image-level annotation, by a large margin (11% with VGG16 and 7.6% with ResNet50). Furthermore, benefiting from our iterative strategy, our weakly supervised method surpasses most regular few-shot models.

Table 2 summarizes the results in the more challenge benchmark COCO-20i . When more objects with various scales appear in the query image, the performance of other weakly supervised models falls much more behind that of fully supervised models. However, our method turns out to be robust to this complex scene, and thus accordingly, the superiority of our method is more significant. Specifically, we achieve more than larger performance gains (20% and 27% for VGG16 and ResNet) compared to the state-of-the-art ones using weak annotation. Furthermore, adopting ResNet50 as the backbone, our method in the weakly supervised setting achieves 42.4 mIoU, which even outperforms the state-of-theart model in the standard setting.

In order to show the comparison more intuitively, we visualize the results of some of the challenging query images, which are segmented by state-of-the-art HSNet (standard few-shot setting) and our method (weakly supervised setting). As shown in Figure 4, our method could perform better in segmenting tiny objects (1st and 2nd columns), semantically confusing objects (3rd column), multiple objects (4th column), and objects in the wild (last two columns). This is probably because our method is able to associate support and query objects with a large difference in appearance with the help of label text, and more likely to predict the complete masks with complex shapes thanks to the prior mask estimation in each step of the refinement.

5.4 Comparison with Zero-shot Segmentation

Weakly supervised few-shot segmentation requires support label text rather than the support mask to segment the query image. From this perspective, it has a lot in common with zero-shot semantic segmentation (ZSS), where only the label text without the support image is available. Therefore, it is natural to compare our method with zero-shot counterparts. However, the comparison is non-trivial due to the inconsistent testing protocol between zero-shot and few-shot segmentation as well as the lack of zero-shot datasets and baselines.

Thus, we still adopt the few-shot benchmark, i.e., Pascal-5i and COCO-20i for evaluation.

In order to apply our framework to the zero-shot segmentation, we need to make some modifications. Therefore, we take the query image as the input of where the support image is required. Concretely, we regard {Iq, ˆM t-1 q } as the {Is, ˆM t-1 s } to predict ˆM t q , and ˆM t s is equal to ˆM t q , then repeat the refinement to obtain final ˆM N q . In this way, the query image is segmented with the guidance by itself. We compare our method with competitive zero-shot methods, e.g. ZS3Net [Bucher et al., 2019], and L-Seg [Li et al., 2022].

As reported in Table 3, our method achieves superior performance in the zero-shot setting with a mean mIoU score of 41.1. Without being designed and trained specifically for zero-shot segmentation and strong backbone ResNet101, our method could outperform the recent zero-shot approaches by a significant margin. However, the zero-shot performance still falls behind that of the few-shot version, which proves the importance of the support image, even without the precise object mask.

5.5 Ablation Study

We conduct the ablation study to investigate the effects of components in our framework. We adopt the ResNet50 as our backbone, and evaluate in the Pascal-5i dataset.

Initial MaskGeneration. In Section 4.1, we introduce two kinds of ways to generate the initial support and query mask ˆM 0 s , ˆM 0 q , i.e., the pixel-level cosine similarity between text and dense visual feature, and the Grad CAM of similarity between text and global visual feature. We further compare the two strategies in the zero-shot and few-shot settings. For zero-shot segmentation, we adopt the generated query mask ˆM 0 q as the predicted query mask. And for few-shot segmentation, we feed the generated support mask ˆM 0 s into the pretrained HSNet as the ground-truth support mask to segment the query images.

Table 4 reports the results of the comparison. From the first two rows (zero-shot setting), it can be seen that the quality of ˆM 0 is similar. Nevertheless, when using the generated mask as the support mask, the results of Grad CAM show clear superiority with nearly 10%. The reason could be that the Grad CAM better leverages the architecture and weights of the pre-trained vision encoder, and thus reflects the strong generality. As a result, we utilize this strategy to generate the initial support and query mask estimation.

Iterative Mutual Refinement. We further compare different iterative strategies in Table 5 to evaluate the effectiveness of IMR. We set the method without Iterative Mutual Refinement as the baseline, which is implemented by feeding ˆM 0 s to the pretrained HSNet, denoted as +HSNet (pretrained). We start with evaluating a simple iterative strategy, where we first predict ˆM 1 q with ˆM 0 s , then update ˆM 2 s with ˆM 1 q , and get the final prediction ˆM 3 q guided by ˆM 2 s , denoted as +HSNet (mutual). Simultaneously, in order to adapt to the condition where the clean ground-truth support mask is unavailable, we trained HSNet with coarse support masks generated by CLIP-CAM. Similarly, we feed ˆM 0 s to the retrained HSNet (+HSNet (retrained)), and refine the estimation of support and query mask iteratively (+HSNet (retrained + mutual)).

As shown in Table 5, both iterative strategies outperform the simple baseline.Therefore, iterative refining mask estimation may be the key to weakly supervised few-shot segmentation. Moreover, our method with IMR still shows great superiority due to the architecture designed with hidden states to fit the varying input distribution during the refinement.

Generalization of Our Method. As mentioned above, our framework is plug-and-play and model-agnostic. In order to verify this, we evaluate our framework with another few-shot segmentation model PFENet [Tian et al., 2020]. Concretely, we adopt the last convolutional layers of PFENet as the head in our method, and the other modules, including the prior generation and feature enrichment module (FEM), as the neck. We keep all the hyperparameters the same with the origin implementation of PFENet. We compare our method built on PFENet with corresponding baselines, where we feed the ˆM 0 s as the support mask to the pre-trained PFENet.

As shown in Table 6, with the well-designed iterative strategy, our method outperforms the baseline by around 6%. Hence, it could be concluded that our method is robust to the selection of the few-shot segmentation model.

6 Conclusion

In this work, we present a general framework to tackle the weakly supervised few-shot segmentation, which contains coarse masks generation from label text and iteratively and mutually refine support and query mask. Extensive experiments show that our method outperforms the state-of-the-art counterparts by a large margin. One limitation is that although our method is able to refine segmentation from coarse masks, it is unclear whether finer initial masks from prompt engineering or post-process such as dense CRF could help performance. We leave it for future work.

Acknowledgements

This work is partially supported by the NSFC fund (61831014), in part by the Shenzhen Science and Technology Project under Grant (ZDYBH201900000002, JCYJ20180508152042002, CJGJZD20200617102601004)