Affinity Attention Graph Neural Network for Weakly Supervised Semantic Segmentation

According to the definition of supervision signals, weakly supervised semantic segmentation can be generally divided into the following categories: based on scribble label [1, 2, 3], bounding box label [4, 5, 6], point label [8] and image-level class label [9, 10, 11]. Scribble, bounding box and point labels are stronger supervision signals compared to the image-level class label since both class and localization information are provided. Whereas image-level labels only provide image class tags with the lowest annotation cost.

Different supervisions are processed with different methods to generate pseudo labels. For the scribble supervision, Lin et al. [2] used superpixel based approach (e.g., SLIC [23]) to expand the initial scribbles, and then used an FCN model [13] to get the final predictions. Tang et al. proposed two regularized losses [1, 3] using the constraint energy loss function to expand the scribble information. For the point supervision, Bearman et al. [8] directly incorporated a generic object prior in the loss function. For image-level supervision, class activation map (CAM) [24] is usually used as seeds to get pseudo labels. For example, Ahn and Kwak designed an affinity net [9] to obtain the transition probability matrix, and used random walk [25] to get pseudo labels. Huang et al. [26] proposed to use the seed region growing algorithm [27] to get pseudo labels from the initial confident class activation map. For the bounding box supervision, SDI [5] used the segmentation proposal by combining MCG [28] with GrabCut [29] to generate the pseudo labels. Song et al. proposed a box-driven method [6], using box-driven class-wise masking and filling rate guided adaptive loss to generate pseudo labels. Box2Seg [30] attempt to design a segmentation network which is suitable to utilize the noisy labels as supervision.

Following the same definition, there are different level sub-tasks for weakly supervised instance segmentation: image-level [12] and bounding box level [5, 7]. For the image-level task, Ahn et al. tried to generate the instance pseudo label using an affinity network  [12]. For the bounding box task, SDI [5] produced pseudo labels using the segmentation proposal by combining MCG [28] with GrabCut [29] while Hsu et al. [7] tried to design a new loss function which samples positive and negative pixels relying on bounding box supervision.

Generally, there are many different GNN [31, 32, 33, 22] methods designed for semi-supervised task and they have achieved satisfying performances. Recently, GNN has been successfully used in various computer vision tasks such as person search [34], image recognition [35], 3D pose estimation [36] and video object segmentation [37], etc.

For weakly supervised semantic segmentation, GraphNet [21] was proposed for bounding box task using GCN [22]. Specifically, based on the CAM technique [24], bounding box supervision was converted to initial pixel-level supervision firstly. And then the superpixel method [23] was applied to produce the graph nodes. They used a pre-trained CNN model to obtain the node feature, which was computed using average pooling for all its pixels. After that an adjacency matrix was computed based on the L1 distance between a node and its 8-neighbor nodes. Later, GCN was used to propagate the initial pixel-level labels to integral pseudo labels. Finally, all the pseudo labels were input into a semantic segmentation model for training. GraphNet proved that GNN was one possible solution for weakly supervised semantic segmentation. However, this method has some limitations. First of all, using superpixels as nodes introduce incorrect node labels, and building an adjacency matrix by a threshold loses some important detailed information. Secondly, the performance of this method is limited by the usage of GNN [22], which only considers the non-weight adjacency matrix. Finally, only cross entropy loss is used for GraphNet, which cannot mitigate the influence of incorrect nodes, edges and labels.

In order to overcome the limitations of previous graph-based learning approach, we design a new approach, $A^{2}$GNN, which takes a more accurate weighted graph as input and aggregates feature by considering both attention mechanism and edge weights. Meanwhile, we propose a new loss function to provide extra supervision and impose restrictions on the feature aggregation, thus our $A^{2}$GNN can generate high quality pseudo labels.

In order to utilize GNN to generate the accurate pixel-level pseudo labels, there are three main problems: (1) How to provide useful supervision information and reduce the label noise as much as possible. (2) How to convert the image data to accurate graph data. (3) How to generate accurate pseudo labels based on the built graph and the supervision.

In this section, we will elaborate on the proposed $A^{2}$GNN to address the above mentioned three main problems. To generate an accurate graph, we propose a new affinity CNN to convert an image to a graph. To provide accurately labeled nodes for the graph, we select highly confident pixel-level seed labels as node labels, and at the same time, we introduce extra online updated labels based on the bounding box supervision, meanwhile, the pixel-level seed labels are further refined by consistency-checking. To generate accurate pseudo labels, we design a new GNN layer since the previous GNN, such as GCN [22] or AGNN [32] is designed based on the assumption that labels are 100% accurate, while in this case, there is no foreground pixel label being 100% reliable.

In Fig. 3, we show the main process of our approach, which can be divided into three steps:

1. (1)

   Generating confident seed labels. In this step, both image-level labels and bounding box-level labels are converted to initial pixel-level seed labels, as explained in section 3. Then the pixel-level seed labels with high confidence will be selected as confident seed labels (section 4.2).

2. (2)

   Converting images to graphs. In this step, we propose a new affinity CNN to generate the graph. Meanwhile, the selected confident seed labels will be converted to corresponding node labels.

3. (3)

   Generating final pixel-level pseudo labels. $A^{2}$GNN is trained using the converted graph as input, and it makes the prediction for all nodes in the graph. After converting node pseudo labels to pixel labels, we generate the final pixel-level pseudo labels.

After that, a FCN model such as Deeplab [14, 15] for BSSS or MaskR-CNN [16] for BSIS is trained using above pixel-level pseudo labels as supervision.

In the following section, we will first introduce how to provide useful supervision, and then we will give an explanation about how to build a graph from the image (section 4.3). Finally, we will introduce $A^{2}$GNN, including its affinity attention layer (section 4.4) and its loss function (section 4.5).

An intuitive solution is to use the pixel-level seed label $M_{F}$ obtained from Eq. (4) as the seed labels. However, $M_{F}$ is noisy and directly using it will be harmful to train a CNN/GNN. As a result, in this paper, we only select those highly confident pixel-level seed labels in $M_{F}$ as the final seed labels. Specifically, we use a dynamic threshold to select top 40% confident pixel labels $M^{\prime}_{I}$ following [40] from the pixel label $M_{I}$ in Eq. (1). Then the selected seed labels are defined as:

where 255 means that the label is unknown. $M_{B}$ and $M_{F}$ are obtained from Eq. (3) and Eq. (4), respectively. Fig. 3 (top-right) illustrates the confident label selection.

Although noisy labels can be removed considerably, the confident label selection has two main limitations: 1) it also removes some correct labels, making the rest labels scarce and mainly focus on discriminative object parts (e.g., human head) rather than uniformly distributed in the object; 2) there still exist non-accurate labels.

To tackle the label scarcity problem in the BSSS task, we propose to mine extra supervision information from the available bounding box. Assuming all bounding boxes are tight, for a random row or column pixels inside a bounding box, there is at least one pixel belonging to the object. Identifying these nodes can provide extra foreground labels. And using the online updated labels, we introduce a new consistency-checking mechanism to further remove some noisy labels from $M_{g}$. We will describe the detailed process in section 4.5 since they rely on the output of our $A^{2}$GNN.

Effective GNN architectures have been studied in existing works  [22, 32], where most of them are designed based on the assumption that the graph node and edge information is 100% accurate. However, in the BSSS task, it is not the case. We propose a new GNN layer with attention mechanism to mitigate this issue. As shown in Fig. 3, in the proposed $A^{2}$GNN, an affinity attention module is applied after the embedding layer. The affinity attention module includes three new GNN layers named affinity attention layers. Finally, an output layer is followed to predict class labels for all nodes.

Specifically, we use a feature embedding layer followed by a ReLU activation function in the first layer to map the initial node features to the same dimension of the assigned feature:

where $X$ is the feature matrix defined in section 4.3.2 and $W^{0}$ is the parameter set of the embedding layer. Then we design several affinity attention layers to leverage the edge weights:

where $P^{l}\in\mathbb{R}^{N_{G}\times N_{G}}$, $N_{G}$ is the number of nodes. For node $v_{i}$, the affinity attention $P^{l}(i,j)$ from node $v_{j}$ is defined as:

where $l\in\left\{1,2,...,L\right\}$ is the layer index ($L$ is set as $3$ in our model) of $A^{2}$GNN and $w^{l}$ is the learning parameter. $S(i)$ is the set of all the nodes connected with $v_{i}$ (including itself). $\left[\cdot\right]$ equals 1 when $\cos(\cdot,\cdot)>0$ and otherwise equals 0. $H^{l}(i)$ and $H^{l}(j)$ correspond to the features of $v_{i}$ and $v_{j}$ at layer $l$, respectively. $\cos(\cdot,\cdot)$ is used to compute the cosine similarity, which is a self-attention module. $E_{ij}$ is the predicted edge in Eq. (14). $\beta$ is a weighting factor. The final output is:

where $W^{L+1}$ is the parameter set of the output layer.

Fig. 5 shows the flowchart of our affinity attention layer. Compared to GCN layer [22] and AGNN layer [43], our affinity attention layer makes full advantage of node similarity and edge weighting information.

As described in section 4.2, we only select confident labels as supervision, which is insufficient for the network optimization. In order to address this problem, we impose multiple supervision on our $A^{2}$GNN. Specifically, we design a new joint loss function, including a cross-entropy loss, a regularized shallow loss [3] and a multi-point (MP) loss:

where $\mathcal{L}_{ce}$ is the cross-entropy loss to use the labeled nodes $M_{g}$ generated in section 4.2. $\mathcal{L}_{reg}$ is the regularized loss using the shallow feature, i.e., color and spatial position. $\mathcal{L}_{mp}$ is the newly proposed MP loss to use the bounding box supervision.

We evaluate our method on PASCAL VOC 2012 [52] and COCO [53] dataset. For PASCAL VOC 2012, the augmented data SBD [54] is also used, and the whole dataset includes 10,582 images for training and 1,449 images for validating and 1,456 images for testing. For COCO dataset, we train on the default train split (80K images) and then test on the test-dev set.

For Pascal VOC 2012 dataset, mean intersection over union (mIoU) is applied as the evaluation criterion for weakly supervised semantic segmentation, and the mean average precision (mAP) [55] is adopted for weakly supervised instance segmentation. Following the same evaluation protocol as prior works, we reported mAP with three thresholds (0.5, 0.7, 0.75), denoting as $\text{mAP}_{0.5}^{r}$, $\text{mAP}_{0.7}^{r}$ and $\text{mAP}_{0.75}^{r}$, respectively. For COCO dataset, following [56], mAP, $\text{mAP}_{0.5}^{r}$, $\text{mAP}_{0.75}^{r}$, $\text{mAP}_{s}$, $\text{mAP}_{m}$ and $\text{mAP}_{l}$ are reported.

Weakly supervised semantic segmentation: In Table I, we compare the performance between our method and other state-of-the-art approaches for BSSS. For using deeplab as the segmentation model, it can be seen that our approach obtains 96.1% of our upper-bound with pixel-level supervision (Deeplab-Resnet101 [14] with CRF). Compared to the other approaches, our approach gives a new state-of-the-art performance. Specifically, our approach with deeplab-resnet101 [49] outperforms Box-Sup [4], WSSL [17] by big margins, approximately 11.8% and 13.2%, respectively. Besides, compared to GraphNet [21], the only graph learning solution, our method with Deeplab-Resnet101 performs much better than it, with an improvement of 10.9% for mIoU (without CRF). We can also observe that our performance is even better than SDI [5], which uses MCG [18] and BSDS [19] as extra pixel-level supervision. When using PSPNet [45] as the segmentation model, our approach obtains 74.4% mIoU without CRF as post-processing, which is even higher than the results in [47] with CRF. Finally, our method with Tree-FCN [49] outperforms the state-of-the-art Box2Seg [30] in this task. Note that Box2Seg focused on designing a segmentation network using noisy label from bounding box, thus our performance could be further improved using their network as the final segmentation network.

Weakly supervised instance segmentation: In Table II, we compare our approach to other state-of-the-art approaches on BSIS. It can be seen that our approach achieves a new state-of-the-art performance among all evaluation criteria. Specifically, our approach performs much better than SDI [5], increasing 14.3% and 11.1% on $\text{mAP}_{0.5}^{r}$ and $\text{mAP}_{0.75}^{r}$, respectively. It can also be found that compared to BBTP [7], which is the state-of-the-art approach on this task, our approach significantly outperforms it by large margins, around 5.1% on $\text{mAP}_{0.7}^{r}$ and 5.8% on $\text{mAP}_{0.75}^{r}$. The performance is increased more on $\text{mAP}_{0.75}^{r}$ than $\text{mAP}_{0.7}^{r}$ and $\text{mAP}_{0.5}^{r}$, which also indicates that our approach can produce masks that preserve the object structure details. One interesting observation is that our approach even achieves better performance than the fully supervised method SDS [55].

In Fig. 6, we compare some qualitative results between our approach and other state-of-the-art approaches for which the source code is publicly available. Specifically, we compare our results with SDI [5]¹¹1we use a re-implement code from: github.com/johnnylu305 for the BSSS task and BBTP [7]²²2github.com/chengchunhsu/WSIS_BBTP for the BSIS task. It can be seen that compared to other approaches, our approach produces better segmentation masks covering object details.

In Table III, we make a comparison between our approach and others on COCO test-dev dataset. It can be seen that our approach performs much better than LIID [61], with an increase of 16.8% on mAP${}^{r}_{\text{50}}$. Furthermore, our approach even performs competitive with fully-supervised approach MNC [60], which also indicates the effectiveness of our approach.

Since the pseudo labels for BSIS are generated from the BSSS task, in this section, we will conduct ablation studies only on the BSSS task. We simply evaluate the pseudo label mIoU on the training set, without touching the val and test set.

In Fig 7, we make a comparison between our $A^{2}$GNN and others for BSSS. It can be seen that our $A^{2}$GNN performs much better than other GNNs, with an improvement of 1.9% mIoU over AGNN [32] when only using the cross-entropy loss, and the full $A^{2}$GNN outperforms AGNN [32] by a large margin (6.9%).

In Table V(a), we explore the influence of different modules in our approaches to generate pseudo labels. Baseline means that we use SEAM [38] to generate the foreground seed labels and then use bounding box supervision to generate the background. RW means that we follow SEAM [38] to use random walk for pseudo label generation. It can be seen that the proposed approach outperforms the baseline by a large margin. And each module significantly improves the performance.

In Table V(b), we study the effectiveness of our joint loss function. It can be seen that compared to $A^{2}$GNN which only adopts cross-entropy loss, our MP loss can improve its performance by 2.1%, validating the effectiveness of our MP loss. With consistency-checking, the performance is improved to 77.2%, indicating the effectiveness of our proposed consistency-checking mechanism. When jointly optimized by these three losses with our consistency-checking mechanism, the performance is further improved to 78.8%.

In Table V(c), we study different ways to build our graph. Superpixel (S.P.) means that we adopt [21] to produce graph nodes and their features. Distance (Dis.) means that we build the graph edge using $L_{1}$ distance of feature map [21]. It can be seen that the performance is improved when directly using pixel in the feature map as the node, suggesting that it is more accurate than using superpixel. When we use our affinity CNN to build the graph, the performance is significantly improved by 4.1%, which shows that our approach can build a more accurate graph than other approaches.

In Table V(d), we study the effectiveness of our affinity attention layer. It can be found that if we use either the attention module or the affinity module separately, the mIoU score is lower than that of the full $A^{2}$GNN, which indicates the effectiveness of our designed GNN layer.

In Table V(e), we show the joint influence of the loss functions and labels for our proposed affinity CNN. It can be seen that when only using $L_{Ac}$, $M_{F}$ labels perform better than $M_{g}$. This is because that $M_{g}$ only provides limited pixels and these pixels are usually located at the discriminative part of an object (such as the human head). Such limited labels are not sufficient when only using $L_{Ac}$. When we use both $L_{Ac}$ and $L_{Ar}$, $M_{g}$ performs much better than $M_{F}$, indicating that $L_{Ar}$ can accurately propagate the labeled regions to unlabeled regions.

In addition, we also analyze the influence of supervision for our $A^{2}$GNN. Specifically, we make a comparison of the results when using $M_{F}$ (in Eq. (4)) and $M_{g}$ (in Eq. (5)) as supervision for our $A^{2}$GNN, respectively. Compared to $M_{F}$, $M_{g}$ has fewer annotated nodes but each annotation is more reliable. The mIoU score on Pascal VOC 2012 training set is 73.2% and 78.8% for $M_{F}$ and $M_{g}$, respectively. This result validates the effectiveness of the leverage of the high-confident labels.

As mentioned in section 3, the common practice to initialize weakly supervised task is to generate pixel-level seed labels from the given weak supervision. For different weakly supervision, we use different approaches to convert them to pixel-level seed labels.

Image-level supervision: we directly use $M^{\prime}_{I}$ defined in Eq. (5) to train our affinity CNN and use it as $M_{g}$ to train our $A^{2}$GNN. The final pseudo labels are generated using the ratio (1:3) to fuse our results and the results of random walk.

Scribble supervision: For the scribble supervised semantic segmentation task, for each class in an image (including background), it provides one or more scribbles as labels. Superpixel method [23] is used to get the expanded labels $M_{S}$ from the initial scribbles. To get seed label to train our affinity CNN, we merge $M_{S}$ with $M^{\prime}_{I}$ using the following rule: if the pixel label in $M_{S}$ is known (not 255), the corresponding label in $M_{g}$ will be the same label as $M_{S}$. Otherwise, the pixel label will be treated as the same label as $M^{\prime}_{I}$. To generate the node labels for $A^{2}$GNN, we directly use $M_{g}=M_{S}$ since it provides accurate labels for around 10% pixels in an image.

Point supervision: For point supervised semantic segmentation, for each object in an image, it provides one point as supervision and there is no annotation for background. To train our affinity CNN, we used $M^{\prime}_{I}$ directly. To generate node supervision for $A^{2}$GNN, we use a superpixel method [23] to get the expanded label $M_{P}$ from initial point labels. Then $M_{g}$ is generated using the same setting with the scribble task.

For our affinity CNN and our $A^{2}$GNN, we use the same setting with our bounding box task.

In Table V, we compare the performance between our method and other state-of-the-art weakly supervised semantic segmentation approaches.

For point supervision, our method achieves state-of-the-art performance with 66.8% and 67.7% mIoU on the val and test set of PASCAL VOC, respectively. Compared to other two approaches [8] and [3], our method increases 23.4% and 9.8% in mIoU on PASCAL VOC 2012 val dataset, respectively.

For the image-level supervision task, our $A^{2}$GNN achieves mIoU of 66.8% and 67.4% on val and test set, respectively. It should be noticed that PSA [9], SEAM [38] and CONTA [70] apply Wider ResNet-38 [50] as segmentation model, which has a higher upper-bound than Deeplab-Resnet101 [14]. Using Deeplab-Resnet101 [14] as the segmentation, Subcat [69] is the state-of-the-art approach on this task, but it require multi-round training processes. Moreover, our method achieves 66.8% mIoU using Deeplab-Resnet101 [14], being 87.0% of our upper-bound (76.8% mIoU score with Deeplab-Resnet [14]) on val set.

Besides, for the image-level supervision, some approaches [26, 64, 39, 61] used salient model with extra pixel-level salient dataset [72] or instance pixel-level salient dataset [58] to generate more accurate pseudo labels. Follow these approaches, we also use saliency models. Specifically, we use the saliency approach [73] following ICD [39] to produce the initial seed labels, and then use our approach to produce the final pseudo labels. It can be seen from Table V that our approach outperforms other approaches (using ResNet101 as the backbone).Following ILLD [61], we also evaluate our approach using Res2Net [71] as the segmentation backbone, and our performance is further improved to 69.0% and 69.6%. For this setting, we have not designed any specific denoising scheme for the seed labels. Nevertheless, our performance is comparable with other state-of-the-art methods, e.g., [61], which also proves that our method can be well generalized to all weakly supervised tasks.

For the scribble supervision task, our method also achieves a new state-of-the-art performance.

In Table VI, we present a comparison to evaluate the pseudo labels on the PASCAL VOC training set. It can be seen that our approach outperforms other approaches. Compared to the state-of-the-art approach SEAM [38], our approach obtains 1.7% mIoU improvement. We also compare the quality of the pseudo labels between our approach and Box2Seg [30]. It can be seen that our method outperforms Box2Seg [30] by a large margin, with 5.2% mIoU improvement.

In Fig. 8, we also present more qualitative results for the above three tasks. It can be seen that stronger supervision leads to better performance and preserves more segmentation details.