点击下方卡片，关注 「AI视界引擎」 公众号
视网膜血管分割是医学领域图像处理中一个众所周知的问题。良好的分割可以帮助医生在诊断眼部疾病时做出更好的决策。
本文描述了在DRIVE挑战中的工作，其中包括对视网膜血管进行分割。作者发明了一种新的方法，结合了StyleGAN2和SA-UNet的使用。本文的创新可以帮助解决任何小数据集分割问题。
1、简介
作者提出了一种改进视网膜图像分割的方法，通过创建图像及其相应的分割地图来实现。作者的解决方案包括使用DRIVE数据集1对StylGAN2进行训练，并使用目前在分割DRIVE图像方面取得最先进结果的SA-UNet模型对新合成的图像进行分割。作者使用新生成的图像及其分割结果来重新训练SA-UNet，以获得更好的结果。
在过去几年中，神经网络（NN）已成为处理许多计算机视觉问题的最佳方法。尽管在那之前，分割问题通常通过边缘检测和其他标准图像处理算法来解决，但神经网络表现出更好的性能。
在过去的几年中，生成对抗网络（GAN）对图像生成器产生了巨大的影响。2019年，NVIDIA团队的Kerras等人展示了一些突破性的进展，他们提出了基于风格的生成对抗网络（StyleGAN）。StyleGAN生成的图像非常逼真，分辨率高，并且可以通过潜空间向量轻松控制。
Changlu Guo等人来自布达佩斯大学，他们提出了一个DRIVE挑战的解决方案。这个解决方案是当前最成功的解决方案之一。在本文中，研究人员提出了对U-Net架构的改变，并将新模型命名为SA-UNet。
U-Net和SA-UNet的主要区别在于卷积层的主干网络结构，以及编码器和解码器之间的空间注意力块，如图1所示。这些改变提高了SA-UNet的性能，特别是在小血管分割方面。SA-UNet在DRIVE数据集上的性能如下（表1）。
2、相关工作
小数据集一直是机器学习算法，尤其是神经网络算法中最棘手的领域之一。增加数据集大小的主要常见方法是使用增强技术，例如裁剪、翻转等。视网膜血管的分割是一项非常具有挑战性的任务，因为血管的大小可能从单个像素到十个像素甚至更大不等。即使对于正常人来说，这个任务也被认为是一项艰巨的任务。
U-Net在数据分割方面表现出色。IterNet网络利用U-Net，通过将多个小U-Net连接起来，取得了优势。IterNet的工作方式是首先生成粗糙的分割地图，然后通过在小区域上应用小型U-Net来改进分割地图。
SA-UNet采用了U-Net架构，但改变了卷积层的主干网络结构，并在编码器和解码器之间添加了空间注意力机制。
Gal Ofir等人在这个领域迈出了一步，他们将StyleGAN与神经网络相结合，旨在增加数据集的大小。他们采用预训练的IterNet，在使用StyleGAN生成合成图像后对图像进行分割，然后对分割数据进行手动修正。通过这种方法，他们解决了小数据集的问题，但代价是图像不够逼真，并且需要手动纠正。
3、数据
DRIVE数据集包含20张训练图像和20张测试图像，图像分辨率为565*584，采用RGB格式。每个图像都有一个覆盖视网膜实际数据的Mask。训练图像包含一次手动分割的血管。所有手动分割血管的人类观察者都接受过经验丰富的眼科医生的指导和培训。他们被要求标记那些他们至少有70％把握是血管的像素。数据还包含许多不同的健康和感染的眼睛图像。在作者的工作中，作者没有利用眼睛的状态信息。
4、本文方法
4.1 SA-UNet
DRIVE数据集是一个众所周知的挑战，已经有许多DRIVE数据集是一个众所周知的挑战，已经有许多方法来解决它。SA-UNet是其中一个最先进的解决方案。SA-UNet的主要优势在于它结合了U-Net、SD-UNet和空间注意力块。SA-UNet的架构与U-Net非常相似，但SA-UNet的参数更少，因此训练过程更短，过拟合问题更不容易发生。
在SA-UNet中，Dropout被DropBlock取代，DropBlock会丢弃感兴趣区域，而不是随机丢弃权重。此外，在任何一对卷积层之间，都有Batch Normalization，如图2所示。此外，SA-UNet还在编码器和解码器之间引入了空间注意力块。
4.2 StyleGAN
数据增强一直是计算机视觉领域的一个重要问题。有许多方法可以进行数据增强，从更简单的方法如图像的传统变换（如翻转、裁剪、旋转等）开始。这些方法在一开始就取得了不错的性能，但仍然需要创建接近真实图像但有着重要区别的图像。神经网络提供了创造"deep fake"图像的机会，生成对抗网络（GAN）提供了创建这种图像的替代方法。目前最先进的GAN是StyleGAN2，相比于SW-GAN和DCGAN，它取得了更好的性能。
4.3 Transfer Learning
训练神经网络模型的最佳方法之一是采用预训练网络，仅训练最后几层。作者对迁移学习和从头开始训练全新网络进行了比较。作者选择使用FID指标来评估网络图像，其在公式1中给出。
FID指标通过将原始图像x和生成图像转换到特征空间来衡量它们之间的距离。在这个过程中，代表平均值，代表协方差，类似于正态分布中的参数。而代表迹（trace）。
根据作者的测试结果，迁移学习和从头开始训练StyleGAN几乎是相同的。StyleGAN生成的图像与真实图像的比较结果可以在图4中展示。
4.4 Suggested Pipeline
在作者的工作中，作者建议使用一个集成方法来改善小数据集问题，该方法基于相同的技术（见图3）。这个集成方法可以帮助获得更好的性能。
- 用原始数据集训练最好的神经网络。
- 使用原始数据集训练StyleGAN2。迁移学习也是一个不错的选择。
- 使用第一部分中的网络对生成的合成图像进行分割。
- 从头开始使用所有数据训练相同的网络。
- （可选）重复步骤2-4，直到获得足够好的结果。
根据作者的推断，作者建议使用生成的数据开始训练最先进的网络，并在最后几个Epoch使用原始数据进行训练。这是有道理的，因为原始数据是最准确的分割数据，这将导致最好的微调效果。
通过在生成数据上进行初始训练，网络可以学习到整体的结构和特征。然后，在最后几个Epoch使用原始数据进行训练，可以进一步优化网络，并使其更好地适应原始数据的细节和特点。这种策略可以帮助提高网络的性能和泛化能力。
4.5 Ensamle
在这篇论文中，作者还测试了另一种方法，即将上述流程与原始的SA-UNet结合使用，以创建网络集成。下一部分将详细描述这种方法。
5、实验
5.1 评价指标
为了评估模型的成功性，作者将分割结果与地面真实数据进行比较，并对每个像素进行检查，确定它是真正阳性（TP）、假阳性（FP）、假阴性（FN）还是真阴性（TN）。然后，作者定义了敏感性（SE），它表示在所有预测为正样本中真正阳性样本的比例（TP/TP+FP），也被称为阳性预测值（PPV）。
特异性（SP）是一种衡量正确识别阳性的比例的指标（TP/TP+FN），也被称为真正阳性率（TPR）。ACC是一种衡量被正确分类的像素与数据集中总像素之间比例的指标（TP+TN / TP+TN+FP+FN）。最后，有一种用于评估分割的独特指标称为F-度量（F1），其定义如下：
该度量标准比较了测试数据集的相似性和多样性。而对于最后一个指标，也是最重要的指标是ROC曲线下的面积（AUC）。
5.2 训练细节
这两个模型（StylGAN2和SA-UNet）的训练是在Google Colab的Pro+版上完成的。预处理训练包括调整图像大小，使其成为2的幂次方的矩形大小。在这种情况下，选择的大小是512×512像素。
StyleGAN2的训练是使用ADATorch的Patch 3进行的，采用旋转和缩放进行数据增强。损失度量选择了"fid50k-full"指标，并使用ADAM优化器进行了4K个迭代周期的训练。训练开始时，使用预训练的StyleGAN模型在"ffhq512"数据集上进行了训练。训练完成后，作者使用截断值为0.7生成了1000张图像。
SA-UNet的训练是从原始论文中的参数开始进行的，共进行了150个Epoch。前100个Epoch使用自适应学习率开始为0.001进行训练，后50个Epoch使用学习率开始为0.0001进行训练。作者选择了ADAM优化器，并选择平均Dice系数和二元交叉熵作为损失函数。DropBlock的丢弃块大小设置为7，与原始论文中相同。
在训练完StyleGAN2和SA-UNet之后，作者使用个合成图像从头开始对SA-UNet进行了100个Epoch的训练，学习率为0.001。由于数据内存限制，作者的批量大小为2。然后，作者使用原始DRIVE数据集图像进行了50个Epoch的微调。仅使用合成数据，作者已经达到了0.98的准确率，而使用原始数据时准确率为0.96，F1得分为0.89。
在微调之后，作者获得了更好的结果。作者将作者的结果与SA-UNet的原始结果进行了比较，并在表2中进行了展示。
接下来，作者尝试了所有可能的训练组合。首先是真实数据，然后是合成数据（或者反之），分别包括使用增强技术和不使用增强技术（分别应用于真实数据、合成数据以及两者同时应用），以及不同数量的合成图像（50、200、500或1000张）。
5.3 结果
作者取得了一些有趣的工作结果，每个结果在某些指标上都创造了新的SOTA（State-of-the-Art），但在其他指标上也出现了退化。
SA-UNet训练的最佳配置是首先使用真实的DRIVE图像进行带有增强（颜色和仿射变换）的100个Epoch训练，然后再使用1000个合成图像进行50个Epoch的训练，但不进行增强。作者得到的AUC指标为0.977，而原始论文中为0.987；ACC指标得到了相同的结果0.968；而F1得分为0.9363，而在SA-UNet原始论文中为0.8820。训练过程如图5所示。
在精确度方面，作者获得了0.887的高分，远高于SA-UNet原始论文中的0.802。但在敏感度方面，作者的模型性能从0.853下降到了0.7295。当作者寻找最佳F1得分时，作者获得了0.983的F1得分，以及0.94的AUC、0.955的准确率和0.969的精确度。另一个有趣的工作结果是0.982的AUC，0.92的F1得分，0.9674的准确率和0.8726的精确度。
在尝试了上述所有配置之后，作者将所有网络再次与原始的SA-UNet模型一起进行了测试，这次是以上面描述的集成模型的形式。由于时间限制，作者只尝试了非常简单的集成技术，将两个模型的预测组合成一个分割图。这种集成对结果的改进很小。所有的结果都在表2中。
6、结论
添加通过GAN生成的合成图像可以在缺乏真实标注数据时有助于获得更好的结果。令人惊奇的是，仅使用合成数据对SA-UNet进行训练可以达到与使用真实数据训练的SOTA模型相媲美的结果。
借助StyleGAN的帮助，作者无需对每张图像进行手动处理，成功在某些指标上达到了SOTA水平，同时在其他指标上几乎保持了相同的结果。与此同时，作者不能忽视的是，由StyleGAN生成的图像在细节上没有达到足够好的分辨率，这导致了潜在更好结果的丧失。这是因为原始的SA-UNet在非常小的血管方面存在一些弱点。
7、参考
[1].SA-UNet for Retinal Vessel improvment using StyleGAN2.
8、推荐阅读
[
即插即用 | Uni-Head这个学霸让ATSS/TOOD/RetinaNet/FreeAnchor通通涨点！！！](https://mp.weixin.qq.com/s?__biz=MzkyNjU0NjA4OA==&mid=2247484813&idx=1&sn=f0fa56686d07cd8dcc03d43ed548dc1d&chksm=c234e44af5436d5cc5080c1788391efb539230ab624999f1f3a319d97f7f883fd440a4f5f4df&scene=21#wechat_redirect)
[
港科大重组YOLOv5模块 | HIC-YOLOv5优化小目标问题，虽不新颖但很有用!](https://mp.weixin.qq.com/s?__biz=MzkyNjU0NjA4OA==&mid=2247484777&idx=1&sn=4fb70b7fa5eccac6fd0300baae45227f&chksm=c234e4aef5436db87bcd8f90a0203a9f154fa9b266ab037365abf5fab09d1fe90f2aaf2b1a58&scene=21#wechat_redirect)
[
可以取代NMS？IOU感知校准，取代NMS，成就YOLOX等目标检测更快的后处理方法](https://mp.weixin.qq.com/s?__biz=MzkyNjU0NjA4OA==&mid=2247484712&idx=1&sn=60c1498bef4d69806001458c142d64c4&chksm=c234e4eff5436df9ef20b0bba1ce96a529f8c2f7ecaef355273090b7825ffc571da1d0e60f1f&scene=21#wechat_redirect)
DMKD蒸馏 | 模型学习，空间/通道信息我都要！超越FKD、FGD、MGD以及AMD方法！
点击上方卡片，关注 「AI视界引擎」 公众号