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

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.

There are two types of approaches for anomaly segmentation, including uncertainty-based methods and generative model-based 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.

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.

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), few-shot 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.

Suppose $\mathcal{C}_{in}=\{C_{in,1},C_{in,2},...C_{in,N}\}$ are $N$ in-distribution classes, which are all annotated in training datasets, and $\mathcal{C}_{out}=\{C_{out,1},C_{out,2},...C_{out,M}\}$ are $M$ novel classes not involved in the training datasets. Here, the semantic segmentation network $\mathcal{S}$ is divided into a feature extractor $\mathcal{F}$ and a label predictor $\mathcal{G}$, where $\mathcal{S}=\mathcal{G}\circ\mathcal{F}$.

For the close-set segmentation, we minimize the following loss $\mathcal{L}_{seg}(\mathcal{F},\mathcal{G})$ which guides $\mathcal{S}$ to produce a pixel-level segmentation for in-distribution classes.

where $\ell_{ce}(\cdot,\cdot)$ indicates the multi-class cross entropy loss, $\mathbf{X}\in\mathbb{R}^{3\times H\times W}$ is an input image, $\mathbf{Y}$ is the corresponding label.

After training this module, we obtain a trained feature extractor $\mathcal{F}$ and a trained label predictor $\mathcal{G}$. The feature map $\mathbf{F}=\mathcal{F}(\mathbf{X})\in\mathbb{R}^{N_{1}\times H\times W}$ and the non-normalized logit $\mathbf{U}=\bar{\mathcal{G}}(\mathbf{F})\in\mathbb{R}^{N\times H\times W}$ can then be generated for in-distribution classes, where $\bar{\mathcal{G}}$ is obtained by removing the softmax layer of $\mathcal{G}$. The feature map $\mathbf{F}$ and the non-normalized logit $\mathbf{U}$ will be used in later modules.

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 $\mathbf{E}$ can be generated as follow,

where $\mathbf{X}$ is the input image, $\mathbf{U}$ is the non-normalized logit, and $\mathbb{I}(\cdot)$ is an indicator function, $\alpha$ and $\beta$ are hyper-parameters to control the edge prediction. According to $\mathbf{E}$, we use a post-processing sub-module, including the hole filling and connected component algorithms, to generate the candidate regions $\mathcal{R}=\{\mathbf{R}_{1},\mathbf{R}_{2},\cdots,\mathbf{R}_{T}\}$, where $\mathbf{R}_{i}\in\{0,1\}^{H\times W}$ represents the $i$-th region.

We then propose a RAML module for anomaly segmentation to classify the candidate regions $\mathcal{R}$. For each region $\mathbf{R}_{i}\in\{0,1\}^{H\times W}$, the region-aware feature embedding is obtained as below:

where $\mathbf{F}^{j,k}\in\mathbb{R}^{N_{1}}$ is the feature vector of pixel $(j,k)$, $\mathcal{D}(\cdot)$ consists of two fully-connected layers to control the embedding dimension. $f_{object}$ 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 $f_{l}$ can be obtained using the semantic segmentation label. Then, the region-aware anomaly probability of $\mathbf{R}_{i}$ 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 $\mathbf{Q}^{j,k}$ is computed as,

where the pixel $(j,k)$ belongs to region $\mathbf{R}_{i}$, $\mathbf{F}$ is the feature map, $\mathcal{P}(\cdot,\cdot)$ is the region-aware anomaly probabilities. $\mathbf{U}^{j,k}_{(l)}$ is the $l$-th output of pixel $(j,k)$ in the non-normalized logit $\mathbf{U}$. We then normalize the uncertainty intensity $\mathbf{Q}^{j,k}$ for each pixel to obtain the anomalous probability map, which is used to identify the unknown regions in the image.

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 close-set segmentation model under the loss $\mathcal{L}_{seg}$. 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 $\hat{\mathbf{X}}$. MCA takes the prediction of the label predictor $\mathcal{G}$ in the close-set model as its input to output $(N+K)$ channels with softmax activation $\mathbf{C}\in[0,1]^{(N+K)\times H\times 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 $\mathcal{L}_{seg}$ during the fine-tuning, and the overall loss function is,

The first term $\mathcal{L}_{seg}$ is the segmentation loss for all in-distribution 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 $\ell_{dice}(\cdot,\cdot)$ indicates the dice loss and $\mathbf{C}_{i}$, $\mathbf{C}_{j}$ 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 $\mathbf{C}^{j,k}_{i}$ represent $(j,k)$ pixel output of $i$-th channel and $\eta$ is a hyper-parameter to control the separation. $\mathcal{L}_{split}$ 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 $\odot$ is the element-wise multiplication operator and $\mathbbm{1}_{H\times 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 $\mathcal{R}=\{\mathbf{R}_{1},\mathbf{R}_{2},\cdots,\mathbf{R}_{T}\}$ for the final RAML module of incremental few-shot learning.

Similar to Equation 3, the region-aware feature embedding $f_{object}$ for each region $\mathbf{R}_{i}$ could be computed. The prototype of $i$-th unknown class ($1\leq i\leq M$) from $L$-shot newly labeled images is defined as:

where $f_{i}^{(j)}$ represents the feature embedding of $i$-th unknown class in $j$-th annotated image. For each region-aware feature embedding $f_{object}$, 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 $C_{out,i}$ only if the cosine similarities satisfy the following two criteria:

where $\theta_{novel}$ is a hyper-parameter to control classification.