Region-Aware Metric Learning for Open World Semantic Segmentation via Meta-Channel Aggregation

1 Introduction

The breakthrough of deep learning in many fields of computer vision is based on the closed set assumption, which means that all classes in the test should be covered in the training set. However, this assumption rarely holds in the open world. Since most computer vision applications have to deal with unknown classes, models, especially the deep models, must

∗ Equal contribution.

†Correspondence to Li Zhang: zhangli pku@pku.edu.cn

(b) Proposed region-aware method

Figure 1: Main idea of our proposed method. (a) Existing methods focus on pixel-level which may result in fine-grained segmentation errors. (b) Our proposed Region-Aware Metric Learning (RAML) method maintains the semantic integrity of the OOD objects.

handle the out-of-distribution (OOD) data. Quite a number of work for image recognition and classification in the open world has been proposed since the first introduction of the concept "open world" in [Bendale and Boult, 2015]. However, the work about open world segmentation is scarce. It is not until recently that [Cen et al., 2021] proposes a two-step framework to achieve open world semantic segmentation. The framework consists of (1) an anomaly segmentation module that extends the close-set model of in-distribution objects to delineate the unknown regions of the OOD objects correctly, and (2) an incremental few-shot learning module that separates the unknown regions into OOD objects with novel classes. They also introduce metric learning into both stages of open world segmentation, and the results prove that their proposed criteria of metric learning can improve the model's segmentation of OOD objects.

Although this pilot work provides a good framework for open world segmentation tasks, the model can be improved in two aspects for better performance. First, the metric learning in [Cen et al., 2021] relies on the pixel-wise feature embeddings, which may falsely split the object into pieces and result in numerous fine-grained segmentation errors. For example, as shown in Figure 1, the bus wheels and the car wheels have similar feature embeddings and are highly likely to be classified into one group according to the pixel-wise feature embeddings, but they apparently belong to different classes in semantic segmentation. To solve this kind of problems, we propose region-aware metric learning (RAML) for open world segmentation, which significantly outperforms pixelwise metric learning (PML) in multiple experiments.

Moreover, we improve the model performance, especially for the incremental few-shot learning stage, by introducing a novel region separation module named meta-channel aggregation (MCA). MCA first aims at over-segmenting the unknown regions into several meta channels. Regions belonging to different meta channels are aggregated to form a segmentation of the objects and then evaluated by the Region-aware Metric Learning module.

In addition, [Cen et al., 2021] sets a fixed center embedding for each in-distribution class, i.e., a one-hot vector in the feature space. Although the fixed center embedding can effectively create a distance between the distribution of different classes, it ignores the relative similarity between them. For example, in the Cityscapes dataset, the method fails to reveal that the difference between person and rider is smaller than the difference between either of them and sky. This paper aims to overcome the drawback by exploiting a more natural metric learning to constrain the distance between the interclass region-aware features. Specifically, we replace the onehot setting in [Cen et al., 2021] with Circleloss [Sun et al., 2020] as the objective of the metric learning, which not only maintains a fine inter-class distance but also shapes the intraclass distribution more concentrated. Experiments show that such division of the feature space is more conducive to segmenting the OOD data.

In summary, we propose a region-aware metric learning method for open world semantic segmentation. Our contributions are as follows:

• We propose using the region-aware over pixel-wise features for open world semantic segmentation to ensure better semantic integrity of the segmented OOD objects.
• We introduce the MCA module as a novel region separation method that suits incremental few-shot learning.
• We adopt Circleloss [Sun et al., 2020] to enlarge the inter-class distance and reduce the intra-class distance of the data samples, improving the performance of the RAML module.

2 RelatedWork

2.1 Region-aware Semantic Segmentation

The ideas of how to apply regional information to improve semantic segmentation have been discussed by many research groups recently, including two main threads. First, several works have shown that region-aware information has better contextual representation than pixel-level information to achieve pixel labeling [Yuan et al., 2020]. Secondly, for image segmentation tasks, region-aware information can be better combined with metric or contrastive learning to manipulate the feature space more effectively [Wang et al., 2021; Hu et al., 2021]. These ideas inspire our paper, but the above works require a sufficient number of training samples to obtain the reasonable region-aware feature representation, while our work is in an open world setting that can only access a few images with unseen class labels. Therefore, we have to design novel region-separation modules (such as MCA) that fit the open world segmentation tasks.

2.2 Anomaly Segmentation

There are two types of approaches for anomaly segmentation, including uncertainty-based methods and generative modelbased methods. Uncertainty refers to the level of not belonging to known classes, widely used to determine abnormal states. The baseline of uncertainty-based methods is maximum softmax probability (MSP) reported by [Hendrycks and Gimpel, 2017]. [Hendrycks et al., 2019] then improves MSP using maximum logit (MaxLogit) for better performance on large-scale datasets. Other uncertainty-based methods include using Bayesian neural networks [Gal and Ghahramani, 2016] and maximizing the entropy of OOD objects in the images [Chan et al., 2021]. On the other hand, generative model-based methods also perform well, including autoencoder (AE) [Baur et al., 2018] and GAN-based methods [Xia et al., 2020]. However, generative models suffer from unstable training and usually have complex network backbones.

In this work, we follow the idea of MaxLogit and develop our anomaly segmentation based on non-normalized logit.

2.3 OpenWorld Problem

[Bendale and Boult, 2015] is the first research that gives the formal definition of "open world", i.e., an open world model must incrementally learn and extend its generality, thereby making the objects with novel classes "known" to itself. Since then, the research on open world problems has increased, including classification [Zhong et al., 2021], object detection [Joseph et al., 2021], instance segmentation [Saito et al., 2021], among others. However, it is not until recently that [Cen et al., 2021] proposes the first framework of open world semantic segmentation. Our work follows the settings in [Cen et al., 2021] and divides the problem into anomaly segmentation and incremental few-shot learning. However, to ensure semantic integrity and improve the segmentation performance, we use region-aware feature embedding instead of pixel-wise feature extraction in their original method.

2.4 Metric Learning

Deep metric learning constrains the distance between feature embedding of learning samples to manipulate the feature distribution. Its applications are seen in various computer vision tasks, such as open set recognition [Chen et al., 2020], fewshot learning [Oreshkin et al., 2018] and open world semantic segmentation [Cen et al., 2021]. Classic metric learning includes two paradigms. The first is to learn with pair-wise labels, under the guidance of triplet loss [Schroff et al., 2015] and center loss [Wen et al., 2016]. The second consists of softmax cross-entropy and variants that train the model with class-level labels. A recently proposed method called Circle loss [Sun et al., 2020] unifies the above two paradigms and forms the feature space with large inter-class distances and small intra-class distances. We thus adopt Circle loss as the key objective of our proposed RAML module.

3 Methods

As shown in Figure 2, our proposed method contains: 1) a backbone model for close-set segmentation, 2) an anomaly segmentation process to delineate the unknown regions of OOD data, and 3) an incremental few-shot learning step for splitting the unknown regions into objects with novel classes.

3.1 Close-set SegmentationModule

Suppose Cin = {Cin,1, Cin,2, ...Cin,N} are N in-distribution classes, which are all annotated in training datasets, and Cout = {Cout,1,Cout,2, ...Cout,M} are M novel classes not involved in the training datasets. Here, the semantic segmentation network S is divided into a feature extractor F and a label predictor G, where S = G ◦ F.

For the close-set segmentation, we minimize the following loss Lseg(F, G) which guides S to produce a pixel-level segmentation for in-distribution classes.

where ℓce(·, ·) indicates the multi-class cross entropy loss, X ∈ R3×H×W is an input image, Y is the corresponding label.

After training this module, we obtain a trained feature extractor F and a trained label predictor G. The feature map F = F(X) ∈ RN1×H×W and the non-normalized logit U = ¯ G(F) ∈ RN×H×W can then be generated for in-distribution classes, where ¯ G is obtained by removing the softmax layer of G. The feature map F and the non-normalized logit U will be used in later modules.

3.2 Anomaly Segmentation

To identify the candidate regions of region-aware anomaly segmentation, we adopt an uncertainty-based OOD objects detection method, MSP [Hendrycks and Gimpel, 2017], as our region separation module, named Uncertainty-based Region Separation (URS). Its high uncertainty response around the object edges could be used as a promising initialization of the region separation, as shown in Figure 3.

To further enhance the edges, we introduce Sobel filtering over the original input image. The final edge prediction map E can be generated as follow,

where X is the input image, U is the non-normalized logit, and I(·) is an indicator function, α and β are hyperparameters to control the edge prediction. According to E, we use a post-processing sub-module, including the hole filling and connected component algorithms, to generate the candidate regions R = {R1,R2, · · · ,RT}, where Ri ∈ {0, 1}H×W represents the i-th region.

We then propose a RAML module for anomaly segmentation to classify the candidate regions R. For each region Ri ∈ {0, 1}H×W , the region-aware feature embedding is obtained as below:

where Fj,k ∈ RN1 is the feature vector of pixel (j, k), D(·) consists of two fully-connected layers to control the embedding dimension. fobject is compared to all the prototypes of the known classes by metric learning constrained by circle loss [Sun et al., 2020]. Specifically, the prototype of l-th known class fl can be obtained using the semantic segmentation label. Then, the region-aware anomaly probability of Ri can be expressed as below,

Finally, to generate a pixel-level anomalous probability map, we combine the information from the non-normalized logit and the above region-aware anomaly probabilities. For each pixel (j, k), uncertainty intensity Qj,k is computed as,

where the pixel (j, k) belongs to region Ri, F is the feature map, P(·, ·) is the region-aware anomaly probabilities. Uj,k (l) is the l-th output of pixel (j, k) in the non-normalized logit U. We then normalize the uncertainty intensity Qj,k for each pixel to obtain the anomalous probability map, which is used to identify the unknown regions in the image.

3.3 Incremental Few-shot Learning viaMCA

After the anomaly segmentation, open world semantic segmentation requires the model to identify all objects of M novel classes in the unknown regions. One way to realize the incremental few-shot learning is to use a few labeled images containing objects with novel classes to fine-tune the closeset segmentation model under the loss Lseg. However, experiments show that this improvement is trivial. We thus propose an innovative MCA module for further creating sub-regions in the unknown regions from anomaly images ˆX. MCA takes the prediction of the label predictor G in the close-set model as its input to output (N +K) channels with softmax activation C ∈ [0, 1](N+K)×H×W . The first N channels are the segmentation results for all in-distribution classes, while the last K(K > M) channels are meta channels to overly segment the unknown regions. Several MCA-related losses are integrated into Lseg during the fine-tuning, and the overall loss function is,

The first term Lseg is the segmentation loss for all indistribution classes from Equation 1. The second term utilizes the negative of Dices to minimize the intersection between any pairs of output channels, which is defined as:

where ℓdice(·, ·) indicates the dice loss and Ci, Cj are the i-th and j-th channels of the segmentation output.

The third term aims to avoid the sub-regions (candidates of OOD objects) gathering in a few certain channels:

where Cj,k i represent (j, k) pixel output of i-th channel and η is a hyper-parameter to control the separation. Lsplit reaches the minimum when the sub-regions scatter across the output channels according to Jenson's inequality.

The last term encourages the outputs of all channels to reconstruct the entire image, further avoiding loss of information:

where ⊙ is the element-wise multiplication operator and 1H×W is a matrix with all ones.

As shown in Figure 4, we observe that MCA tends to segment objects based on local semantic information. One unknown object may be segmented into more than one channel and lose completeness. (e.g., The windows and wheels of cars may be divided into different channels.) Therefore, we aggregate the sub-regions from certain meta channels according to few-shot (here L-shot) labeled images, which generates the candidate regions R = {R1,R2, · · · ,RT} for the final RAML module of incremental few-shot learning.

Similar to Equation 3, the region-aware feature embedding fobject for each region Ri could be computed. The prototype of i-th unknown class (1 ≤ i ≤ M) from L-shot newly labeled images is defined as:

where f (j) i represents the feature embedding of i-th unknown class in j-th annotated image. For each region-aware feature embedding fobject, we use cosine similarity to measure the distance between this candidate region and every unknown class:

The candidate region can be classified as the i-th novel class Cout,i only if the cosine similarities satisfy the following two criteria:

where θnovel is a hyper-parameter to control classification.

4 Experiments

Our experiments include three parts: (1) experimental results of anomaly segmentation in subsection 4.1; (2) experimental results of incremental few-shot learning results in subsection 4.2; (3) ablation studies in subsection 4.3 and Appendix.

4.1 Anomaly Segmentation

Datasets. 7000 full-frame annotated driving scenes from BDD100k [Yu et al., 2020] are used to train the close-set segmentation model, containing 19 categories of objects as in-distribution objects. For anomaly segmentation, we use another two road scene datasets, Lost and Found [Pinggera et al., 2016] and Road Anomaly [Lis et al., 2019], with anomalous objects other than ones in BBD100k.

Implementation details. We follow [Hendrycks et al., 2019; Cen et al., 2021] to use PSPNet as the network backbone of our close-set segmentation module and apply two fully connected layers for RAML. We follow [Hendrycks and Gimpel, 2017] to use three metrics to evaluate the performance of anomaly segmentation, including area under ROC curve (AUROC), area under the precision-recall curve (AUPR), and the false-positive rate at 95% recall (FPR95).

Results. As shown in Table 1, our proposed RAML module achieves the SOTA performance on Lost and Found and Road Anomaly for anomaly segmentation. Figure 5 presents some visual examples to compare RAML and the pixel-wise method. The proposed RAML module produces higher response values and better integrity within the anomalous objects, significantly reducing the false-negative cases.

4.2 Incremental Few-shot Learning

Datasets. we use Cityscapes dataset to train and evaluate our RAML module in the incremental few-shot learning step. Cityscapes consists of 2975 real-world images in the training set and 500 in the validation set with a resolution of 2048 × 1024. The division of training set and test set in our experiments is consistent with this division.

Implementation details. We follow [Cen et al., 2021] to train a DeeplabV3+ model as the close-set model, which is followed by two fully connected layers for RAML and use mean Intersection-over-Union (mIoU) to evaluate the performance of segmentation results. Specifically, mIoUold and mIoUnovel are the mIoUs of known and unknown classes, respectively. The metric mIoUharm is a comprehensive index [Xian et al., 2019] that balances mIoUold and mIoUnovel.

Results. We test our method on CityScapes and compare our method to pixel-wise NPM and PLM proposed by [Cen et al., 2021]. In our experiment, car, truck, and bus are 3 OOD classes not involved in the training stage while the other 16 classes are regarded as in-distribution classes. As shown in Table 2, our proposed RAML module outperforms the previous methods with a relatively large margin. According to Figure 8, pixel-wise metric learning shows erroneous broken segmentation results on OOD objects, while the proposed RAML demonstrates a remarkable ability to maintain the integrity of these results. In addition, Figure 6 shows that the feature embeddings produced by the proposed RAML maintain a reasonable inter-class distance and their intra-class distributions are also more concentrated. Such feature distribution could foster the model to obtain a robust decision boundary.

4.3 Ablation Study

Ratio of unknown classes to known classes. The performance of the trained segmentation model has highly correlated with the amount of training information. We compare our proposed RAML method with the current SOTA method, NPM [Cen et al., 2021], under the different ratios of unknown classes to known classes. As shown in Figure 7, although our RAML method has a decline in performance as the ratio increases, it outperforms NPM in all ratio settings.

Losses in MCA. This section evaluates the losses of our MCA module. As shown in Table 3, the reconstruction loss ensures that our model obtains all information for the unknown classes, significantly improving the validity of MCA. The intersection loss and split loss also bring relatively smaller gains by improving the distribution of candidate regions in meta channels.

5 Conclusion

We have proposed RAML to enhance the performance of open world semantic segmentation. The main reason is that the region-aware feature outperforms the pixel-wise feature on maintaining the semantic integrity of the segmented OOD objects. Effective region separation methods are needed to realize RAML on anomaly segmentation and incremental fewshot learning. We, therefore, adopt the classic uncertaintybased methods to extract candidate regions for anomaly segmentation and propose an MCA module to further separate the anomaly regions for incremental few-shot learning. Experimental results show that our proposed method achieves the SOTA performance on the anomaly segmentation and the overall open world semantic segmentation. Our method has the potential to boost the use of open world semantic segmentation in practical applications.

6 Acknowledgments

This work is supported by the Grants under the National Natural Science Foundation of China (NSFC) under Grants 12090022, 11831002, 71704023, and Beijing Natural Science Foundation (Z180001).