W-Net:一种无监督的图像分割方法.

本文提出了一种无监督的图像分割方法,堆叠两个UNet分别作为编码器和解码器。编码器输入原始图像,输出对应的隐空间表示;解码器将隐层表示重构成原始图像。经过训练,对隐层表示使用一些后处理方法(条件随机场平滑和层次分割)后可以作为图像分割的结果。网络由UNet堆积成,因此得名WNet。整个网络的流程图如下:

1. W-Net结构

WNet网络的结构如图所示,使用一个UNet作为编码器将输入图像映射到紧凑的特征表示,然后使用一个UNet作为解码器从其低维表示中重构输入图像。

UNet编码器由捕捉上下文的收缩路径和相应的扩展路径组成,这种结构能够实现精确的定位。收缩路径首先使用卷积层处理图像,然后在每个下采样层(最大池化层)将特征的空间尺寸减半,通道的数量加倍。在扩展路径中,模型通过上采样层(转置卷积层)将特征的通道数量减半、空间尺寸加倍。收缩路径中每个模块的输出也通过残差连接补充到到扩展路径中相应模块的输入中,以恢复由于下采样而丢失的空间信息。编码器最后使用一个具有softmax激活的$1\times 1$卷积层将特征向量映射为类别$k$的通道数。UNet解码器的结构与编码器一致,它读取编码器的输出作为输入,最后一个卷积层的输出是原始输入的重建结果。

⚪ 重构损失

重构损失强制使得编码表示包含尽可能多的原始输入信息,同时可以使分割预测结果更好地与输入图像对齐。

对于输入尺寸为$224 \times 224 \times 3$的输入图像$X$,通过一个UNet编码器$U_{Enc}$变换到尺寸为$224 \times 224 \times K$的隐藏层特征表示,其中$K$表示分割的类别。将该特征通过另一个UNet解码器$U_{Dec}$重构成图像;构造两者的重构损失:

\[J_{reconstr} = || X- U_{Dec}(U_{Enc}(X;W_{Enc});W_{Dec})||^2_2\]

⚪ 归一化割损失

重构损失会使隐藏层特征随机地落入特征空间中,为使得该特征包含尽可能多的信息,引入归一化割(normalized cut,Ncut)作为衡量特征质量的全局分割标准:

\[Ncut_{K}(V) = \sum_{k=1}^{K} \frac{cut(A_k,V-A_k)}{assoc(A_k,V)} = \sum_{k=1}^{K} \frac{\sum_{u \in A_k,v \in V-A_k}^{} w(u,v)}{\sum_{u \in A_k,t \in V}^{} w(u,t)}\]

其中$A_k$是第$k$个分割结果中的像素集合,$V$是所有像素的集合,$w$衡量两个像素之间的权重。归一化割损失可以最小化不同类别之间的总归一化分离,并最大化同一类别内的总归一化关联。

在确定像素的分割类别$k$时会使用到argmax函数,由于该函数是不可导的,因此无法参与反向传播过程。作者设计了一种可微的soft-Ncut损失:

\[J_{soft-Ncut}(V,K) = \sum_{k=1}^{K} \frac{cut(A_k,V-A_k)}{assoc(A_k,V)} = K- \sum_{k=1}^{K} \frac{assoc(A_k,A_k)}{assoc(A_k,V)} \\ = K- \sum_{k=1}^{K} \frac{\sum_{u \in V,v \in V}^{} w(u,v) p(u=A_k)p(v=A_k)}{\sum_{u \in A_k,t \in V}^{} w(u,t)p(u=A_k)} \\ = K- \sum_{k=1}^{K} \frac{\sum_{u \in V,v \in V}^{} w(u,v) p(u=A_k)p(v=A_k)}{\sum_{u\in V}^{} p(u=A_k) \sum_{t \in V}^{} w(u,t)}\]

归一化割的权重矩阵构造如下:

\[w_{i,j} = e^{-\frac{-||F(i)-F(j)||_2^2}{\sigma^2_I}} * \begin{cases} e^{-\frac{-||X(i)-X(j)||_2^2}{\sigma^2_X}} & \text{if }||X(i)-X(j)||_2<r \\ 0 & \text{otherwise} \end{cases}\]

其中$X(i)$和$F(i)$分别是结点$i$的空间位置和像素值。实现时选择$\sigma_I = 10, \sigma_X=4, r=5$。

⚪ 训练过程

通过交替优化重构损失和soft-Ncut损失,模型能够获得准确的重构结果,并学习到适用于分割的特征表示。其中归一化割损失只作用于编码器,而重构损失同时作用于编码器和解码器。

训练过程中两个损失的变化如图所示。当引入soft-Ncut损失后,重构损失下降缓慢且不太稳定。这是因为归一化割损失强迫隐藏表示空间与图像的较好分割结果更加一致,因此其重建原始图像的能力减弱。最后损失函数都收敛,这意味着模型在最小化重构损失和最大化隐藏特征组内的总关联之间进行权衡。

2. 后处理

对从编码器获得的隐藏层特征进行一系列后处理即可得到图像分割的结果。后处理流程如下所示,首先应用条件随机场CRF进行图像平滑,再使用层次分割生成最终的分割结果。

⚪ 条件随机场平滑

在像素级的标记任务中,卷积神经网络的空间不变性和较大的感受野可能会导致定位精度降低,并且缺乏平滑性约束可能会导致对象描绘不佳。为了解决这个问题,作者使用完全连接的条件随机场(CRF)模型处理编码器的输出,能够改进具有细粒度边界的分割。

全连接CRF模型采用能量函数:

\[E(X) = \sum_{u} \Phi(u) +\sum_{u,v} \Psi(u,v)\]

其中$u,v$是图像$X$上的像素。定义一元势函数$\Phi(u) = -\log p(u)$,二元势函数$\Psi(u,v)$通过在不同的特征空间中使用两个高斯核测量具有不同标签的两个像素$u,v$的加权惩罚。

下图展示了原始图像、编码器的输出结果和通过CRF的平滑结果。编码器的输出使用粗略的边界描述了目标的大致位置,CRF使得边界更加尖锐,且较小的虚假区域被平滑或去除。

⚪ 层次分割

全连接CRF的输出通常是输入图像的过度分割结果,最后一步需要通过适当地合并片段以形成最终的图像分割。作者采用了一种高效的层次分割方法,首先将过度分割的分区转换为加权边界映射,然后迭代合并最相似的区域。

通过计算基于谱聚类的多尺度局部线索和全局边界测量的加权组合来测量初始过度分割边界上每个像素的“重要性”:

\[gPb(x,y,\theta) = \sum_{s} \sum_{i} \beta_{i,s} G_{i,\sigma(s)}(x,y,\theta)+\gamma sPb(x,y,\theta)\]

其中$s$是像素索引的比例,$i$是索引特征通道(包括亮度、颜色和纹理),$G_{i,\sigma(s)}(x,y,\theta)$测量以半径$\sigma(s)$为中心、位于通道$i$中角度$θ$处$(x,y)$的圆盘两半之间的相异性。$mPb$信号测量图像中的所有边缘,$sPb$信号捕捉图像中最显著的曲线。

下图b给出了原始输入图像a由CRF生成的边界为红线的初始区域示例,图c给出了初始边界的相应加权边界图。

之后从该加权边界图根据轮廓检测构建层次分割。该算法有两个步骤:一个建立初始过度分割区域的定向分水岭变换(Oriented Watershed Transform, OWT)和一个超度量等高线图(Ultrametric Contour Map,UCM)。这是一种基于贪婪图的区域合并算法。

3. 实验分析

作者在PASCAL VOC2012数据集(包含$11530$幅图像和$6929$个分割)上训练网络,然后使用Berkeley分割数据库(BSDS300BSDS500,分别包含$300,500$个图像)评估网络。

结果表明,当不使用归一化割损失时,模型能够进行高质量的重构,然而编码器学习的隐特征具有噪声且更加离散。而引入归一化割损失后能够获得更一致的特征表示,尽管重构结果变差。通过隐藏表示的一致性和重建质量之间的权衡证明了在训练期间使用归一化割损失的合理性。

为评估无监督图像分割方法的性能,作者选用信息变量(Variation of Information, VI)、概率随机指数(Probabilistic Rand Index,PRI)和分段覆盖(Segmentation Covering,SC)作为评估指标,分别在在最优数据集尺度(Optimal Dataset Scale,ODS)和最优图像尺度(Optimal Image Scale,OIS)下报告了结果。

作者也给出了一些在BSDS500图像上的分割结果。第一行显示原始输入;第二行显示了编码器产生的初始边界经过CRF平滑的结果;第三行和第四行分别显示了在ODSOIS上生成的超等高线图。