[论文阅读] Self-supervised Correction Learning for Semi-supervised Biomedical Image Segmentation

论文地址：https://doi.org/10.1007/978-3-030-87196-3_13
代码：https://github.com/ReaFly/SemiMedSeg
发表于：MICCAI’21

Abstract

生物医学图像分割在计算机辅助诊断中起着重要作用。然而，现有的基于CNN的方法在很大程度上依赖于大量的人工标注，这些标注非常昂贵，需要大量的人力资源。在这项工作中，我们采用了从粗到细的策略，并提出了一种用于半监督生物医学图像分割的自监督校正学习范式。具体来说，我们设计了一个双任务网络，包括一个共享的编码器和两个独立的解码器，分别用于分割以及病变区域的补全。在第一阶段，只使用分割分支，以获得一个相对粗糙的分割结果。在第二步，我们根据初始分割图对原始图像上检测到的病变区域进行屏蔽，并将其与原始图像一起再次送入网络，同时分别进行补全和分割。对于有标签的数据，这个过程由分割标注监督，而对于无标签的数据，则由被屏蔽的病变区域的补全损失指导。由于这两个任务依赖于相似的特征信息，未标注的数据有效地增强了网络对病变区域的表示，进一步提高了分割性能。此外，还设计了一个门控特征融合(GFF)模块，以纳入这两个任务的互补特征。在三个不同任务的医学图像分割数据集上进行的实验，包括息肉、皮肤病变和眼底视盘分割，很好地证明了我们的方法与其他半监督方法相比的出色性能。

I. Motivation

本文属于一种半监督进行图像分割的范式。与无监督，弱监督，主动学习等其他范式类似，本文的思想也为使用尽可能少的标注训练样本来实现更高效的模型训练。

II. Network Architecture

本文网络的结构如上所示。由于引入了一个inpainting任务作为辅助，因此会有两个Encoder-Decoder子网络，其中一个用来进行分割任务，而另一个用来进行inpainting任务。不过，在本质上分割网络和inpaint网络学到的都是与医学图像有关的特征，完成不同的任务只需要设计不同的decoder，因此文中指出这两个encoder是"共享权重的"。从代码实现角度来讲，可能确实只使用了一个共享encoder。

落实到具体的结构上，本文的设计重心并不在设计新颖的网络结构，Encoder采用的是ResNet-34(没用最近特别"火"的Transformer)，而Decoder也非常类似于原始的U-Net，各Decoder Block仅由两个Conv-BN-ReLU块构成，没有加其他复杂的结构设计。

接下来将分析本文的思路亮点。

III. Training Pipleline

记分割网络的编码器，解码器分别为$E_{seg}$，$D_{seg}$，inpainting网络的编码器，解码器分别为$E_{inp}$，$D_{inp}$。给定输入图像$x\in {\mathbb{R}}^{H\times W\times C}$，首先将其送入分割网络，这时候可以得到一个(粗的)分割结果${\hat{y}}_{\text{coarse }}$：$${\hat{y}}_{\text{coarse }}=D_{\text{seg }}\left(E_{\text{seg }}(x)\right)$$ 基于此预测结果，我们可以把输入图像x中的相应部分给mask掉，即：$$x_{\text{mask }}=x\times \left(1-{\bar{y}}_{\text{coarse }}\right)$$ 此时，就可以将$x_{\text{mask}}$视为输入，$x$作为真值来训练inpainting网络。

现在的关键问题就在于，inpainting网络怎么去辅助改善segmentation网络的性能？首先不考虑前面谈到的权重共享问题，两个encoder所学到的特征是有区别的，但也会比较相似。换而言之，两者的特征是"互补"的，因此需要设计相应的模块来进行交互。

具体来说，本文通过提出的GFF模块来解决特征交互的问题，GFF的详细结构将在下一节中阐述，这里首先给出数学定义。对于第$i$层分割编码器特征$e_{seg}^i$与inpainting编码器特征$e_{inp}^i$，将其送入一个GFF模块$G_i$进行"融合"：$$e^i=G_i(e_{seg}^i,e_{inp}^i)$$ 然后在送入相应各自的解码器中。

额外值得一提的是，这是使用了deep supervision策略，防止inpainting网络迅速过拟合(因为$e_{seg}^i$)中本身就包含学到的完整图像信息，如果只对最终结果监督的话容易导致inpainting网络"走捷径"。

IV. Gated Feature Fusion

门控特征融合单元的结构如下所示：

至于这个东西的作用，其实可以简单地将其理解为一种注意力模块。因为$e_{seg}^i$与$e_{inp}^i$中的很多特征都是重复的，那么直接concat融合的话会导致我们真正需要的互补特征被削弱，因此需要使用一些注意力机制对融合后的特征进行一些"后处理"。

形式化地讲，对于融合后的特征，我们分别将其通过两个不同的卷积$W_r$与$W_s$，并进行sigmoid处理，得到(空间)注意力图$r_i$与$s_i$：$$r_i=\sigma \left(W_r\left[e_{seg}^i,e_{inp}^i\right]\right)$$ $$s_i=\sigma \left(W_s\left[e_{seg}^i,e_{inp}^i\right]\right)$$ 对于$r_i$，其作为应用于$e_{inp}^i$上的权重，称为reset gate，用于抑制$e_{inp}^i$中的背景信息；而$s_i$作为应用于$e_{seg}^i$上的权重，称为select gate，用于更好地融合互补信息。形式化地，有：$${\tilde{e}}^i=W\left[r_i\times e_{inp}^i,e_{seg}^i\right])$$ $$e^i=s_i\times {\tilde{e}}^i+\left(1-s_i\right)\times e_{seg}^i$$

V. Loss

对于分割网络，其损失函数使用的是经典的BCE Loss&Dice Loss；对于inpainting网络，由于我们只关注被mask区域的复原效果，因此损失函数为mask区域的L1 Loss。