Deep Watershed Transform for Instance Segmentation

Min Bai & Raquel Urtasun
UfT

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

在图像处理领域中广泛使用的分水岭算法通常用于实现图像分割。其基本思路在于构建一个能量映射来表征整个图像，在这种映射中存在显著差异的部分通常对应着重要的边界特征。这些显著差异形成了界定物体边界的关键轮廓线——即所谓的分水岭线；而被定义为吸引子盆（通常假设这些盆内的能量相对较低）的核心区域则被认为是目标对象所在的捕获 basin。如图所示：

该算法采用类似于注水过程的方式来识别物体区域；从而阻止相邻区域之间水流的汇合，最终实现对物体的分离。

分水岭算法多采用梯度值构建能量图；这样做的结果是将平缓的区域划分得更为清晰。然而分水岭算法的主要缺点在于其容易导致过分割现象的发生。通过使用标记控制的方法来实现图像分割时选择合适的标记仍然是一个启发式的问题。

2. Deep Watershed Transform

作者的核心策略是通过网络结构来训练分水岭算法中的能量映射。其目标在于所学得的能量映射能够明确标识出实例间的边界。从而使得直接运用该算法即可实现对实例分割的目标。如图所示，则展示了传统方法基于梯度大小生成能量映射与本研究中自监督学习所得能量映射之间的对比分析。

从图中可以看出，在上图中有3个实例存在。如果以梯度大小作为能量图（Energy Map），可能会导致出现较多的分割结果（共7个），而通过学习获得的能量图能够更精确地描述一个实例。

直接学习实例间的边界较为具有挑战性，在具体实施中需要采取特殊策略以提高效率。具体而言, 作者的方法可分为两个阶段: 首先通过Direction Network模型来学习实例内部点到边界的指向; 接着利用生成的方向图对Watershed Transform网络进行训练以获取实例内的能量分布; 最后将获得的能量图谱应用于分水岭算法实现实例分割

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

2.1. Direction Network

该网络接收RGB图像数据以及对应的语义分割图（通过引入分割结果使模型聚焦于目标区域而非背景部分）。同时旨在通过学习能量变化方向来确定优化路径（具体而言，在实例边界处能量较低而在中心区域较高）。因此优化过程将沿着实例边界进行调整以达到最低能量状态。

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

具体而言，在数学表达式中，
D_{gt}(p) 表示像素 p 到其对应 instance 边界的最短距离。
以下是关于能量梯度方向的示意图：

由于选取作者方向时考虑了数值连续性的问题，在此方案中并未选择基于绝对角度的真实值作为基准。这是因为绝对角度在某些情况下（如0°与360°）虽然在几何上是等价的但在计算过程中可能会引入较大的偏差。为了克服这一限制我们采用了基于相对角度的能量梯度方法以确保数值上的连续性和稳定性。通过采用基于能量梯度的方向法可以发现当两个实例之间存在相互遮挡的情况时损失函数会特别关注边界区域的能量梯度变化从而增强了模型对边缘检测的能力。此外当某个像素单元被错误分类到某个实例时此时的能量梯度将与实际真值的方向呈现相反的趋势这样的损失函数能够有效地引导网络在像素级别上实现精确的目标识别。

Direction network负责生成一个双通道的方向性map。其结构示意图如图所示。该网络基于改进版的VGG架构进行设计。训练损失函数的具体计算公式如下所列：

其中符号\langle u, v \rangle用于计算两个向量之间的内积；而\cos^{-1}运算的结果则是这两个向量之间夹角；权重系数w_p与每个instance所占区域面积呈反比例关系；总损失即为角度误差平方之和。

2.2. Watershed Transform Network

基于前面所述的方向网络模型，在获得每个像素点的能量梯度方向后,wtn系统将通过二维映射图的方式学习每个像素对应的能量水平分布。其中(level) 0标记背景区域及距离边界不超过2个像素的像素点,随着(level) 值增加,该区域与物体中心的距离逐渐减小。作者共手动设置了K=16个不同级别的特征层

本节将介绍WTN模块的学习流程，在输入数据方面采用方向图表示（direction map），输出结果则基于能量的变化计算得到（energy map）。该学习框架所采用的损失函数即基于能量的变化计算得到（based on energy variation）交叉熵（cross-entropy）。

其中 仍然是与原来相同的，并与 instance 面积呈负相关性的系数。c_k 表示与 energy level 相关的系数。由于我们最终选择 level 大于 0 或 1 来分割 instance，因此我们希望对于那些 level 较小的实例能够尽可能地准确学习。

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

2.3. Energy Cut and Instance Extraction

在最后阶段进行图像分割时，在处理低级别的目标类别（如person和rider等）时，将相应的区域划分设置为level=0；而对于高级别的目标类别（如car和truck等），同样将其区域划分设置为level=1。

3. Experiments

该研究仅限于CityScapes数据集上的实验验证。在具体实施过程中，在使用RGB图像作为基础输入后，在分割任务中分别提取ground truth（简称gt）以及instance segmentation任务中的gt信息被用来训练direction network；随后，在基于ground truth中的direction map和instance segmentation进行WTN（Wavelet Transform Network）模型的设计与求解；最后将这两个模块串联起来，并结合RGB图像及PSPNet输出进行微调优化。

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

该系统对分割效果的敏感度并不显著；然而当实现优化时，instance segmentation的表现会进一步提升

定性分析结果如下：