Retinal Vessel Segmentation under Extreme Low Annotation: A Generative Adversarial Network Approach

Retinal Vessel Segmentation under Extreme Low Annotation: A Generative Adversarial Network Approach

严格意义来说不算StyleTransfer，但还没开新的分类，暂且放到风格迁移类吧。
基于生成对抗网络完成超低标注的视网膜血管分割

Abstract

借助深度学习进行医学图像分割之前往往需要该领域的专家提供详尽的标注信息，而一些医学图像分割任务中受限于有限的训练数据使得很多模型性能不佳。本文基于对抗生成网络搭建了一个数据高效的网络，即使在极少标注信息的参与下也可以取得比监督学习更好的分割结果。

本文的前期工作的一个延续，主要创新点在于提出了一种新的对抗损失函数和无监督网络框架，也是首次搭建了了基于结构预测的对抗网络用于医学图像分割。
实验结果表明在DRIVE和STARE两种眼底血管图像数据集中，本文基于极少标注信息的分割效果均超过了前面基于监督学习的方法；此外消融实验的一系列结果也揭露了成功训练 GAN的一些关键recipe。

Section I Introduction

卷积神经网络和深度学习在诸多计算机视觉任务中取得了巨大的成功，如物体检测、语义分割、视频理解、VQA等。近期人们逐渐把研究兴趣放到了医学图像分析上，如理解、分析医学图像序列、图像分割、病灶检测、图像修复等。
但是深度学习取得如此骄人的成绩是以动辄上百万训练参数组成的复杂网络为代价的，比如（当时）较为流行的AlecNet和VGG系列参量均上百万，训练过程中很容易过拟合；而获取带有详细标注信息的数据集也是一件费时费力且成本高昂的事情，医学图像更是需要专家进行标注。因此需求逐渐转变成在有限的训练数据训练网络。


所以本文聚焦于解决这样一个问题：是否可以从标注数据和未标注数据中均能学习到有用的信息。以视网膜血管分割为例，是否可以在标注样本仅占0.8%-1.6%的比重下有效学习图像特征从而完成血管分割。比如一共60000个patch用于训练，其中只有500-1000张带有标注信息。


本文的主要工作可以概括为：
（1）基于前期半监督的GAN框架增加了新的无监督损失函数项，这样对于标注图像和无标注图像均可以进行学习；

（2）将分割任务当做“多标签推断任务”来解决，因此将中心像素预测框架（CP,CenterPixel)改进编程结构预测框架（SP,Structure Prediction)

（3）探究了应如何权衡标注图像比例以及标注信息多样性两者之间的关系

（4）在DRIVE和STARE两个视网膜血管图像分割数据集上进行了测试，主要评价指标为ROC，超过了前任使用监督学习进行分割的结果。

Section II Related Work

传统的血管分割框架可分为无监督学习和监督学习两类，无监督学习主要通过不同算法检测出特定结构的局部性质，如line-matrix，co-occurrence matrix,co-linearly aligned difference-of-Gaussian filters；监督学习则会输入图像patch以及对应的标注信息patch，从而学习图像到类别标签的映射，这方面的综述可以参考[9]。

Liskowski等人首次用CNN实现视网膜血管分割任务，将血管的分割转化为一个二分类问题，其中血管部分为正类，背景非血管部分为负类。随后又有许多学者进行了改进，有的使用一系列堆叠的autoencoder；还有的需要借助一部分带有标注信息的patch进行学习，带标注的patch数量高达3.810^6 或1.210^5.而本文与此相反，致力于使用尽量少的标注图像（低于500张）进行学习。

Section III Methods

Part A GAN



首先回顾一下GAN的结构，包含两个网络：生成网络和判别网络。


生成网络通过输入服从某一分布的随机噪声输出生成图像；判别网络负责区分输入的样本是真实的还是合成的，近期GAN已经被用于写实性图像的生成、图像去噪等，本文将探索半监督学习的GAN进行医学图像分割的效果。





Part B Semi-Supervised Learning



监督学习训练的分类器一般输出具有Nc个结点对应Nc个类别，经过softmax转化为概率输出，使用交叉熵损失函数进行优化。
如果要使用GAN进行半监督学习，判别网络分类器的输出就要增加一个node，也就是一共Nc+1个，额外多出的一个用来判别输入样本是real还是generated。因此判别网络的损失函数包含supervised、adversarial、unsupervised三部分。










其中Lsup就是交叉熵损失函数，因为监督学习部分有对应的标注信息可以参考；



Ladv部分当输入generated样本时会输出更高的fake类别的概率；




Lunsup则会对输入real样本但判别为fake的情况施加惩罚。
这样就可以从大量无标注的样本中进行学习。






Fig1中展示了生成网络和判别网络的详细结构，可以看到生成网络还是常规的结构，但是判别网络分成了监督训练和无监督训练两部分，监督网络通过输入带有标注信息的patch产生分割图谱；无监督网络部分判断样本是real或fake。






Part C Center Pixel vs Structured Prediction







医学图像分割主要由两类策略：中心像素预测（ center-pixel prediction，CP）和结构化预测（structure prediction,SP）。CP的目标是寻找一个拟合函数完成数据到类别的映射，比如返回one-hot类别向量或者一个概率标量；而SP结构化预测并不是输出一个标量或类别而是产生有结构的输出，比如一个序列，一个句子、一棵树、一幅图等，而不再仅局限于中心像素点的类别。







CP








SP

Part D Generator Training








原始的GAN生成网络如果应用于半监督学习，需要进行一定的调整，[40]的研究表明GAN生成器采用特征匹配（feature matching)是一种较好的方法，特征匹配的关键思想在于，一个成功的生成器对于real类和fake类，鉴别器中间层的激活应该是相同的。因此生成器的损失函数为：









Part E Joint Training of G and D

两个网络的联合训练参考paper[11]中的设定，比如在更新判别器参数时会固定生成器的参数，只根据Ld更新判别器的参数，反之亦然。

Section IV Implimentation Details

Part A Generator


采用UNet的Encoder-Decoder结构，因此UNet已经成功地用于多种医学图像分析任务中。输入图像经过Encoder编码得到特征表达，再经过Decoder恢复到原分辨率，同层的Decoder和Encoder通过skip connection在通道上进行级联，这样更便于梯度的传递。TableI具体展示了Unet的结构，可以看到经过Encoder和Decoder均使用了两层，中间C5是bottleneck，以及中间层随机失活概率为0.8.


Part B Discriminator

D的结构输入一个100D的向量，经过线性层整形，随后通过转置卷积逐步恢复至原始图像分辨率。
Part C Optimization
采用mini-batch SGD，Adam Optimizer
Lr=10^-4
batch_size = 64

Section V Experiments

Part A Datasets



DRIVE:包含40张视网膜血管图像，test:train=20:20



STARE:train :test=19:1



首先所有的视网膜血管图像被转换为灰度图，因为前人研究表明绿色通道辨别能力最佳，因此再对RGB图进行灰度转化时绿色通道赋予了更高的权重；随后还对灰度图像通过CLAHE进一步增强了对比度；patch大小：48*48.



Part B 消融实验



Table III展示了使用feature matching和原始GAN在DRIVE数据集上的对比，以及GAN使用不同设定（均值池化、最大池化、归一化）等消融实验的结果。




Max Pool vs Avg Pool




Unet中使用的最大池化可能不利于GAN的训练，因为取的是感受野中的最大值，而均值池化可以在压缩参数的同时取得更好的分割效果，因此在本文中用均值池化替换了原始UNet中的最大池化。





Normalization:
对中间激活层或权重进行归一化与否对GAN的训练至关重要。DCGAN提出后权重归一化就已经成为GAN训练的标配了，因此本文也测试了不同归一化策略对分割效果的提升（Instance Norm vs Weight Norm)

Selecting layers for Feature Matching

在[40]中选择的是倒数第二层用于特种匹配，但本文认为对于一些低层次的视觉任务使用网络中如此深层次的特征不是一个明智之举，就像本文只是对每一个像素点实现一个二分类，因此低层信息对于准确分类是必不可少的。

通过分析TableIII的实验结果我们发现仅选择较低层次或者较深层次效果都不是最佳的，Conv1和Conv2级联后进行特征匹配得到的分割效果是最佳的。


More Patches or more Images?


标注数目的重要性
标注采集也分好坏，最坏的情况下可能500个标注patch均来自于同一张图像，最好的情况是从所有训练图像中均匀采集patch。



Table II 展示了不同patch数量在DRIVE数据集上分割的AUC指标对比，一个有趣的发现就是，即使采用较少的采样预算，但是来源于更大的采样范围，分割效果也趋于更佳。比如在STARE数据集中，在15张图像中采样500个patch的AUC达到了0.85，而从10张图像采样1000patch的AUC只有0.81，在DRIVE数据集上也是如此。
Table IV展示了与其他SOTA框架的对比，可以看到有全监督和半监督两类，以及不同的标注预算（0.5K-10K），本文以半监督的方式取得了最佳的分割效果。

Section VI Conclusion and Discussion

本文以半监督学习的方式完成视网膜血管图像的分割，通过无监督对抗损失函数的约束可以有效学习无标注样本的信息用于提升分割效果。主要创新之处在于无监督对抗损失函数以及编解码结构的GAN网络中采用不同的特征匹配方式，最终在很少标注样本的参与下取得了SOTA的分割结果。通过实验还发现标注样本的多样性对于精确分割是十分重要的，未来还可以进一步探索本文框架在其他医学图像上的应用。