BriNet: Towards Bridging the Intra-class and Inter-class Gaps in One-Shot Segmentation

Information Exchange Module. Given a pair of support and query images, their features are initially extracted by a common CNN model, such as ResNet. Following that, the information exchange modules are introduced to refine the features based on our belief that the common information contained in the support and the query images should be shared and co-weighted during feature extracting. The details of the IEM are given in Fig. 2 (b).

Specifically, before the initial feature maps $F_{s},~{}F_{q}\in\mathbb{R}^{c\times h\times w}$ extracted from ${\mathbf{x}}_{s},~{}{\mathbf{x}}_{q}\in\mathbb{R}^{3\times H\times W}$ by CNN could be embedded by an IEM, we apply the support mask $\mathbf{m}_{s}$ on the support feature map $F_{s}$ to filter out unrelated background information. The resulting pure target object feature map can be written as $F_{s}^{m}=F_{s}\odot\mathbf{m_{s}}$, where $\odot$ denotes element-wise multiplication and $\mathbf{m}_{s}$ is down-sampled from ${W\times H}$ into the identical size ${w\times h}$ with $F_{s}$.

Our IEM takes $F_{s}^{m}$ and $F_{q}$ as inputs. It consists of Non-local block $\mathcal{T}$ [Wang et al.(2018a)Wang, Girshick, Gupta, and He] and Squeeze-and-Excitation (SE) block $\mathcal{C}$ [Hu et al.(2017)Hu, Shen, and Sun], where Non-local block is used for information exchange and SE block aims to boosting channel information of target object. The Non-local block was also adopted by [Hsieh et al.(2019)Hsieh, Lo, Chen, and Liu] for few-shot object detection but without applying foreground mask. IEM outputs the refined feature maps $F_{q}^{{}^{\prime}}$ and $F_{s}^{{}^{\prime}}$, formulated as Eq. 1,

$\displaystyle F_{q}^{{}^{\prime}}=\mathcal{T}(F_{q},F_{s}^{m})\odot\mathcal{C}(F_{s}^{m}),\hskip 11.38109ptF_{s}^{{}^{\prime}}=\mathcal{T}(F_{s}^{m},F_{q})\odot\mathcal{C}(F_{s}^{m}).$

(1)

Non-local block is originally proposed to capture long-range dependencies focusing on the regions of a single image input, where the non-local transformation is modeled by relations among local features. In contrast, in our approach, the Non-local block under support-query input setting is for inter-input transformation, denoted as Transformation Matrix in Fig. 2 (b), where all local features of one input (say, the support) are transformed into one local feature of the other input (say, the query). The support-query transformation is applied as Eq. 2 and Eq. 3,

$$\mathcal{T}(F_{q},F_{s}^{m})_{i}=\sum_{j}^{wh}{\phi(F_{q})_{i}}^{T}\psi(F_{s}^{m})_{j}\cdot g(F_{s}^{m})_{j}+F_{q,i}$$

(2)

$$\mathcal{T}(F_{s}^{m},F_{q})_{i}=\sum_{j}^{wh}{\phi(F_{s}^{m})_{i}}^{T}\psi(F_{q})_{j}\cdot g(F_{q})_{j}+F_{s,i}^{m}$$

(3)

where $\phi,\psi$ and $g$ are $1\times 1$ convolution kernels and $i,j$ are the index of pixel. The SE block generates channel attention by Global Average/Max Pooling, followed by two sequential MLP layers.

Specifically, given IEM ouputs $F_{s}^{{}^{\prime}},F_{q}^{{}^{\prime}}\in\mathbb{R}^{c\times h\times w}$, apart from global average pooling, we apply 2 slide average pooling windows on feature map $F_{s}^{{}^{\prime}}\in\mathbb{R}^{c\times w\times h}$ with size $s\times h,w\times s$ and strides $s,s$ respectively, where $s$ is the stride size. As a result, more fine-grained feature representations $F_{s,c}^{{}^{\prime}}\in\mathbb{R}^{c\times\frac{w}{s}\times 1},F_{s,r}^{{}^{\prime}}\in\mathbb{R}^{c\times 1\times\frac{h}{s}},F_{s,g}^{{}^{\prime}}\in\mathbb{R}^{c\times 1\times 1}$ are obtained. After convoluting the query feature map $F_{q}^{{}^{\prime}}\in\mathbb{R}^{c\times w\times h}$ with these 3 kernels respectively, the three activation maps are summed as Eq. 4,

$$F_{s-q}=F_{s,c}^{{}^{\prime}}*F_{q}^{{}^{\prime}}+F_{s,r}^{{}^{\prime}}*F_{q}^{{}^{\prime}}+F_{s,g}^{{}^{\prime}}*F_{q}^{{}^{\prime}}$$

(4)

where $*$ denotes convolution operation and $F_{s-q}$ is the output of MCM.

Loss function. To boost performance, in addition to the final cross-entropy segmentation loss $\mathcal{L}_{seg}$, we introduce another auxiliary segmentation branch into our proposed architecture before the Decoder to shorten gradient propagation, as shown in Fig. 2. The auxiliary cross-entropy segmentation loss $\mathcal{L}_{aux-seg}$ is minimized together with the final segmentation loss $\mathcal{L}_{seg}$. Thus the overall loss function is

$$\mathcal{L}=\mathcal{L}_{seg}+\mathcal{L}_{aux-seg}$$

In order to make our model adaptive to the agnostic objects in the test stage, we propose to conduct online refinement. Our proposed online refinement algorithm takes advantage of the support-query pair available at the test stage and switches their roles to extract complementary information for the refinement iteratively.

According to the definition of few-shot segmentation, at test stage, the information about the agnostic categories in the support image $x_{s}^{test}$, as well as its corresponding mask $m_{s}^{test}$, and query image $x_{q}^{test}$ are available, but the query mask $m_{q}^{test}$ and other support-query pairs from $\mathcal{D}_{test}$ are unknown. Inspired by self-supervision, the core of our online refinement is to regard the query offline prediction $\hat{m_{q}^{test}}$ as the pseudo (ground-truth) mask to assist the support image segmentation in return. This refinement could be conducted for a few rounds by switching the roles of the query and the support images iteratively.

Specifically, given a query image $x_{q}^{test}$, a support image $x_{s}^{test}$ and a support ground-truth mask $m_{s}^{test}$, we do the online refinement in three steps. First, we feed $x_{q}^{test}$, $x_{s}^{test}$ and $m_{s}^{test}$ into the model obtained at the offline training to estimate the query mask $\hat{m_{q}^{test}}$. Second, $\hat{m_{q}^{test}}$ is treated as the pseudo ground-truth label for $x_{q}^{test}$, which constitutes the next input with $x_{q}^{test}$, $x_{s}^{test}$ to predict the support mask $\hat{m_{s}^{test}}$. Third, the segmentation model is constantly refined by minimizing the cross-entropy loss between the predicted support mask $\hat{m_{s}^{test}}$ and the ground-truth support mask $m_{s}^{test}$. These 3 steps are repeated until the mean IoU between $m_{s}^{test}$ and $\hat{m_{s}^{test}}$ is higher than a threshold $t_{i}=t_{0}*\frac{N-1}{N+i}$ in the $i$-th step or the maximum iteration time $N$ has been reached. This algorithm is formulated as alternatively updating $\hat{m_{s}}$ and $\hat{m_{q}}$ according to Eq. 5,

$\displaystyle\hat{m_{s}}=\mathcal{F}(x_{s},x_{q},\hat{m_{q}}),\hskip 11.38109pt\hat{m_{q}}=\mathcal{F}(x_{q},x_{s},m_{s}).$

(5)

Here $\mathcal{F}$ indicates the embedding function of the model either trained offline or refined online in the last iteration.

We evaluate the performance of our model on two benchmark datasets commonly used for few-shot segmentation.

We adopt ResNet50 [He et al.(2016)He, Zhang, Ren, and Sun] modified from Deeplab V3 [Chen et al.(2017)Chen, Papandreou, Schroff, and Adam] as our model backbone. In view of the task characteristic, we abandon ResBlock-4 and the later layers of ResNet50, which is consistent with the existing works in [Zhang et al.(2019b)Zhang, Lin, Liu, Yao, and Shen, Zhang et al.(2019a)Zhang, Lin, Liu, Guo, Wu, and Yao]. The parameters of ResNet are initialized from the model pre-trained by ImageNet [Russakovsky et al.(2015)Russakovsky, Deng, Su, Krause, Satheesh, Ma, Huang, Karpathy, Khosla, Bernstein, and et al.] and fixed during training.

Our model is trained using SGD for 400 epochs on Nvidia Titan V GPUs. We set the base learning rate to 2e-2 and reduce it to 2e-3 after 150 epochs. Momentum and weight decay of SGD are set to 0.9 and 5e-4, respectively. The input images have the size of $353\times 353$. For data augmentation, we follow CaNet [Zhang et al.(2019b)Zhang, Lin, Liu, Yao, and Shen] to adopt random mirror, random rotation, random resize and random crop for both datasets. For online fine-tune, the iteration number is $20$.

We evaluate the performance of our proposed model with/without online refinement, and compare it with multiple state-of-the-art methods for few-shot segmentation. The results are reported in Tab. 2 and Tab. 2. It is worth mentioning that, for the two closely related methods CaNet [Zhang et al.(2019b)Zhang, Lin, Liu, Yao, and Shen] and PGNet [Zhang et al.(2019a)Zhang, Lin, Liu, Guo, Wu, and Yao], in addition to directly quoting the results from the original papers (for PASCAL-$5^{i}$), we also re-run the models and report the results, indicated as CaNet* and PGNet* in the tables. For CaNet* and PGNet*, the same support-query pairs are used for test as in our model. In this way, a strict comparison that further removes the difference in randomly sampling test pairs is conducted.

The results on COCO-$20^{i}$ are given in Tab. 2. COCO-$20^{i}$ is a very challenging dataset. As can be seen, on this task, even our offline model only has outperformed the four competitors in terms of the mean IoU. With the proposed online refinement, the performance of our model could be further boosted, making its advantage over the other methods in comparison more salient. The results on PASCAL-$5^{i}$ are given in Tab. 2. This is a relatively easy task and all methods in comparison have better performance than what they do on COCO-$20^{i}$. The performance of our offline model is comparable to that of the second best method FW&B which builds on a more powerful backbone network ResNet-101. Compared with CaNet* and PGNet* that use the same test pairs as ours, our offline model wins both of them with a large margin. Again, our online refinement could further improve our performance on this dataset consistently. Moreover, cross-referencing the results in Tab. 2 and Tab. 2, it seems that our online refinement contributes more to the performance improvement when the segmentation task is hard.

Fig. 3 shows six visual examples of segmentation results from our proposed BriNet and previous best models, CaNet and PGNet. Given the same query image, all of CaNet, PGNet and our BriNet are able to segment different classes with different support examples as guidance (the two rightmost columns in Fig. 3). However, our BriNet can generate more accurate and complete segmentation results compared with CaNet and PGNet, even when both of them totally fail (3rd column and 4th column). Our online refinement improves the model adaptation to agnostic object segmentation significantly (the last two rows in Fig. 3).

To single out the contribution of each component proposed in our model, we conduct an ablation study in order to answer two questions: (i) How do IEM and MCM contribute to the performance of the offline model? (ii) Could our online refinement, as a general method, help other few-shot segmentation models improve the performance? 4-fold validation is used and the mean IoU values are reported.

IEM and MCM. To answer the first question, we compare our offline model with its variants that remove IEM and MCM, respectively. The results are given in Tab. 4. As seen, without either IEM or MCM, the performance of the offline model will significantly decrease on both PASCAL-$5^{i}$ and COCO-$20^{i}$, showing the necessity of employing these two modules, as we argued before.

Online refinement. To answer the second question, we apply our online refinement to CaNet* and PGNet*, and the results are in Tab. 4. Significant improvement could be observed on the hard classes in COCO-${20^{i}}$ for both models. On PASCAL-${5^{i}}$, although little effect is observed on CaNet*, our online refinement could help PGNet* to improve further. This experiment demonstrates the value of our online refinement method as a general strategy to improve few-shot segmentation.