Joint Learning of Saliency Detection and Weakly Supervised Semantic Segmentation

论文地址

摘要

本文通过多任务联合框架用一个网络（SSNet）联合解决WSSS和SD任务。SSNet由segmentation network (SN)和saliency aggregation module (SAM)组成。对于一张输入图像，SN生成分割结果；SAM预测每个类别的saliency然后把所有类别的分割掩码合并到一个saliency map中去。SSNet用image-level标注和saliency标注监督。

1. Introduction

介绍了WSSS和SD任务。这两个任务都要求生成pixel-level的掩码，因此两者有紧密联系。一方面有了saliency map，分割网络的的计算量会大大减少；另一方面有了分割结果，saliency map可以直接通过选择显著性类别得到。但是现有的方法大多通过SD方法生成伪掩码训练分割模型，而忽略了WSSS模型充分利用显著性标注的分割信息。以往的方法利用预训练的显著性模型，而本文直接充分使用显著性标注。

如Fig. 1(b)所示，单独使用image-level标注训练的分割结果并不好，本文用分割结果导出的saliency map，和真实显著性标注求loss来解决这个问题。SSNet的SAM预测每个类别的saliency score，激活显著性类别（如上图的人），抑制非显著性的类别（如上图的瓶子）。
本文主要的contributions如下：

• 提出可以联合解决SD和WSSS任务的端对端框架，其分割任务受image-level标注和显著性标注的监督。
• 设计了一个saliency aggregation module模块来桥接两个任务，两个任务互相提供优化。
• 实验。

3. The proposed approach

3.1. Network overview

我们设计了两个SSNet，分别处于训练的两个阶段，分别叫做SSNet-1和SSNet-2。1的监督信息是image-level标注和saliency标注，2的监督信息是image-level标注和saliency标注和1预测的分割结果。

3.2. Segmentation networks

分割网络由两部分组成，一个是特征提取器，另一个是多层卷积层，根据提取的特征得到分割结果。特征提取器基于VGG和DenseNet设计，有5个卷积blocks，但去掉了全连接分类器和后两个blocks的下采样。

在第一个训练阶段，只有image-level labels作为监督信息。这样训练出来的模型还不够，因此我们在SSNet-1中加入了1×1的卷积层，预测得到C-channel segmentation map和one-channel background map，map上每个点的值在[0,1]之间，每个点每个类别的value加在一起为1。在第二个训练阶段，SSNet-1的分割结果被用作监督信息，我们在SSNet-2中用4个3×3的卷积层， dilation rate分别为6,12,18,24来得到分割结果。

3.3. Saliency aggregation

我们设计SAM模块来利用saliency标注。这个模块利用特征提取器生成的32×32的特征$F$，通过32×32的卷积层和sigmoid函数生成C维向量$v$，$v_i$表示第i个类别的saliency score，这样一来saliency map $S$如下表示：
$$S=\sum _{i=1}^C=v_i\cdot H_i$$
$H_i$表示第i个channel的分割结果。

3.4. Jointly learning of saliency and segmentation

我们用两个训练集（image-level和saliency）来训练SSNet。saliency数据集记作 D s = { ( X n , Y n ) } n = 1 N s \mathcal D_s=\{(X^n,Y^n)\}_{n=1}^{N_s} Ds​={(Xn,Yn)}n=1Ns​​，X为输入图像，Y为真值。classification数据集记作 D c = { ( X n , t n ) } n = 1 N c \mathcal D_c=\{(X^n,t^n)\}_{n=1}^{N_c} Dc​={(Xn,tn)}n=1Nc​​，X为输入图像，t为真值。
分类损失定义如下：
$${\mathcal{L}}_c=-\frac{1}{N_c}\sum _{n=1}^{N_c}[\sum _{i=1}^C{t}_i^{n}log{\hat{t}}_i^n+(1-t_i^n)log(1-{\hat{t}}_i^n)]$$
$t_i^n=1$说明图像中包含第i类目标，否则值为0。${\hat{t}}_i^n$值在[0,1]之间。
image-level category labels可以让分割网络获得语义信息，但是无法学到空间信息。我们用pixel-level
saliency labels解决这个问题。如3.3节所讲，SAM合成saliency map，然后和真值求得loss：
$${\mathcal{L}}_{s1}=-\frac{1}{N_s}\sum _{n=1}^{N_s}[\sum _m{y}_m^{n}log{s}_m^n+(1-y_m^n)log(1-s_m^n)]$$
y是真值，s是预测值。
在训练的第一阶段，我们用以上两个loss训练SSNet-1。训练之后我们在classification dataset ${\mathcal{D}}_c$上得到C+1-channel分割结果。其中C个channel分割结果通过$t^n$去抑制错误类别的预测，最后我们得到了伪掩码。接下来是训练的第二阶段，第二阶段的loss如下：
$${\mathcal{L}}_{s2}=-\frac{1}{N_c}\sum _{n=1}^{N_c}[\sum _{i=1}^{C+1}\sum _{m\in {\mathcal{C}}_i}log{h}_{im}^n]$$
这里$h_{im}^n\in [0,1],i=1,\cdots ,C$表示在分割图像$H^n$第m个像素i通道处的值，即m点属于i类的概率。