Deep Watershed Transform for Instance Segmentation

Min Bai & Raquel Urtasun
UfT

1. 传统的分水岭算法简介

图像处理中的分水岭算法常用来做图像区域分割(segmentation)，基本的思路是计算一张energy map来表示图像，其中物体区域的energy和其他非物体区域的energy差别较大，从而形成包络物体的分水岭，称包络线为watershed line，称物体区域为catchment basins(一般假设物体区域的energy低)。如下图所示：

分水岭算法以一种类似注水的方法找到物体的区域，并通过抑制相邻catchment basins的水交汇，达到分割物体的目的。

分水岭算法通常采用梯度作为energy map，这样的话basin就对应比较平滑的区域。但是分水岭算法的一大局限在于其很容易导致过分割。采用marker-controlled的方法，选择marker又是一个很heuristic的问题。

2. Deep Watershed Transform

作者的基本思路是利用网络来学习分水岭算法里的energy map，目的是希望学习得到的watershed line就表示instance的边界。这样直接应用分水岭算法就可以实现instance segmentation了。下图表示传统分水岭算法使用梯度大小得到的energy map和作者学习得到的energy map的对比：

可以看到，上图中一共有3个instance，如果使用梯度大小作为energy map，容易受到噪声的干扰，从而分割出7个instance，采用学习得到的energy map，能够更好的表示一个instance。

直接学习instance之间的边界是比较困难的，作者的做法分为两步，首先训练一个Direction Network学习instance内部点到边界的方向，然后根据这个direction map再训练一个Watershed Transform Network学习instance内部点的energy level，最后拿energy map应用分水岭算法。

整体的网络框架如下所示：

2.1. Direction Network

该网络的输入为RGB图像以及语义分割的map（加入分割的结果作为输入是为了让网络只关注目标区域，忽略背景区域）。
目标是学习energy下降的方向（instance边缘的能量低，instance中心的能量高，所以下降的方向指向instance的边界）

对于每一个pixel，作者用二维单位向量表示该下降方向的真值，如下：

$$u_{p, gt} = \frac{\nabla D_{gt}(p)}{|\nabla D_{gt}(p)|}$$
其中 Dgt(p) D g t ( p ) $D_{gt}(p)$表示像素p到其对应instance边界的最短距离。
能量梯度方向的示意图如下所示：

作者方向的真值没有选取角度的原因是因为角度的数值和方向不具有连续等价的性质，比如0度和360度方向一致但是数值差太大。采用能量梯度方向还有一个好处是：如果两个instance之间存在相互遮挡，那么loss就会对边界处的能量梯度方向比较敏感，即如果一个pixel分错了instance，那么能量梯度方向肯定是和真值相反的，这样loss会变得很大。这种loss能强迫网络学习pixel级别的精确定位。

Direction network输出的是一个2-channel的direction map，网络部分如上图所示，采用的是修改的vgg模型。这部分网络训练的loss如下所示：

$$l_{direction} = \sum_{p \in P_{obj}} w_p ||cos^{-1} <u_{p, gt}, u_{p, pred}>||^2$$
其中 <> <> <script type="math/tex" id="MathJax-Element-29"><></script>表示求两个向量内积， cos−1 c o s − 1 $cos^{-1}$得到的是两个向量的夹角， wp w p $w_p$是和instance面积成反比的系数， Pobj P o b j $P_{obj}$表示每一个instance。整体的loss是角度误差的平方和。

2.2. Watershed Transform Network

通过前面的direction network得到每个像素点的能量梯度方向之后，wtn再根据这张二维的map学习每一个pixel对应的energy level。其中level 0表示背景以及和距离边界2个pixel以内的像素点。level越高，像素点越靠近物体的中心。作者一共手动选择了K = 16个level。

WTN模块的训练过程如上所示，输入为direction map，输出为energy map。其训练的loss为改动的交叉熵：

$$l_{watershed} = \sum_{p \in P_{obj}} \sum_{k=1}^K w_p c_k (\overline{t}_{p,k} log\overline{y}_{p,k} + t_{p,k}logy_{p,k})$$
其中 wp w p $w_p$还是和原来一样的和instance面积成负相关的系数， ck c k $c_k$表示和energy level相关的系数，因为我们最后是选择level大于0或1来分割instance，因此我们希望level小的部分学习尽量准确，所以 ck c k $c_k$对于小的k，值比较大。

整体训练过程中各个步骤的输入输出如下所示：

2.3. Energy Cut and Instance Extraction

最后分割的时候，对于小目标类别（person, rider, etc），选择level=0进行分割，对于大目标类别（car, truck, etc），采用level=1进行分割。

3. Experiments

作者只在CityScapes上做了实验。
训练的时候，首先先拿RGB图像，分割的gt以及instance segmentation的gt训练direction network；
然后拿gt的direction map和gt的instance segmentation训练WTN。
最后把两个模块级联在一起，用RGB图像和PSPNet的输出做fine-tune。

在CityScapes上比之前的instance segmentation效果好很多

网络对于分割的效果也不是特别敏感，但是分割好，instance segmentation也会更好

定性分析结果如下：