【论文笔记】MGU-Net

【论文笔记】Multi-Scale GCN-Assisted Two-Stage Network for Joint Segmentation of Retinal Layers and Disc in Peripapillary OCT Images

论文下载（途径多种多样）：

[2102.04799] Multi-scale GCN-assisted two-stage network for joint segmentation of retinal layers and disc in peripapillary OCT images (arxiv.org)

Dataset（数据集）

1. Collected dataset: Download our collected dataset.
2. Public dataset: Download Duke SD-OCT dataset

研究背景

青光眼作为世界第一大不可逆致盲性眼病，若未能及时发现和接受，可导致青光眼患者失明。青光眼前期大多没有症状，只有当疾病进展到晚期才可能会被察觉到，但此时往往已“为时已晚”错过了最佳治疗时间。青光眼也因此被称为“沉默的致盲性眼病”。目前，中国约有2180万名青光眼患者，且发病呈现年轻化趋势。当前，预防青光眼最有效的方法是早期筛查、早期诊断和早期干预。

视网膜层的微小形态变化，如视网膜神经纤维层（retinal nerve fiber layer，简写为RNFL）和神经节细胞层（ganglion cell layer, 简写为GCL）变薄，对青光眼的诊断起着至关重要的作用。光学相干断层扫描（optical coherence tomography, 简写为OCT）作为一种非侵入性的三维成像方法可用于视网膜检查。由于OCT的微米级轴向分辨率，它可以直接可视化视网膜的分层结构。如下图诊断报告[2]所示，从OCT图像得到的视神经乳头周围RNFL厚度是早期青光眼诊断的常见指标。因此，一种可以快速、便捷地分割视神经乳头周围视网膜OCT图像的自动化方法对实现青光眼“早发现”尤为重要。

论文摘要

An accurate and automated tissue segmentation algorithm for retinal optical coherence tomography (OCT) images is crucial for the diagnosis of glaucoma. However, due to the presence of the optic disc, the anatomical structure of the peripapillary region of the retina is complicated and is challenging for segmentation. To address this issue, we develop a novel graph convolutional network (GCN)-assisted two-stage framework to simultaneously label the nine retinal layers and the optic disc. Specifically, a multi-scale global reasoning module is inserted between the encoder and decoder of a U-shape neural network to exploit anatomical prior knowledge and perform spatial reasoning. We conduct experiments on human peripapillary retinal OCT images. We also provide public access to the collected dataset, which might contribute to the research in the field of biomedical image processing. The Dice score of the proposed segmentation network is 0.820 ± 0.001 and the pixel accuracy is 0.830 ± 0.002, both of which outperform those from other state-of-the-art techniques.

视网膜光学相干断层扫描（OCT）图像的精确自动组织分割算法对于青光眼的诊断至关重要。然而，由于视盘的存在，视网膜的毛细血管周围区域的解剖结构是复杂的，并且对分割具有挑战性。为了解决这个问题，我们开发了一种新的图形卷积网络（GCN）辅助的两阶段框架，以同时标记九个视网膜层和视盘。具体而言，在U形神经网络的编码器和解码器之间插入多尺度全局推理模块，以利用解剖先验知识并执行空间推理。我们在人的视神经乳头周围视网膜OCT图像上进行了实验。我们还为公众提供对所收集数据集的访问，这可能有助于生物医学图像处理领域的研究。所提出的分割网络的骰子分数为0.820±0.001，像素精度为0.830±0.002，两者都优于其他先进技术。

论文主要内容

首先，该研究将OCT图像输入到第一级的视盘检测网络，通过该网络可以得到指示视盘位置的蒙版和相应的特征图，在原输入图像上应用蒙版以生成无视盘的图像。之后，将其提供给第二级视网膜层分割网络，获取视网膜层的特征图。最后，将第一级网络的视盘特征图和第二级的视网膜层特征图融合得到一个完整的视网膜OCT图像的分割结果。

在这两级网络中，该研究都使用了一种多尺度的图网络辅助的全局推理模块。在第一级网络中获取全局的水平空间特征，而在第二个网络中捕获全局的垂直空间特征。

下面结合论文中给出的图，

图1：（a）黄斑OCT图像和（b）视神经乳头周围OCT图像之间的比较

人工分割视神经乳头周围图像。使用十种标签，包括RNFL、GCL、IPL、INL、OPL、ONL、IS/OS、RPE、脉络膜和视盘，并用不同的颜色标注。层结构遵循视神经头位于图像中心的规律，而更薄的视网膜层在两侧分层。

图2：分割框架的示意图

（a） 分割整个过程包括三个主要步骤：

①第一级为初始视盘检测

② 第二阶段是视网膜层分割

③ 前两级的输出随后进行融合

（b） 多尺度GCN辅助U形网络（MGU网络）的简化说明。

MGU网络由一对编码器和解码器组成，其中插入了多尺度全局推理模块（MGRM）。MGU网络和MGRM的详细结构分别见图6和图5。

图3：视神经乳头周围视网膜OCT图像的基于图形的表示

图像在全局范围内具有“non-disc”-“disc”-“non-disc”的水平布局。“non-disc”区域的放大视图呈现分层结构。我们的分割框架旨在利用视网膜视神经乳头周围区域的解剖先验知识，解决视网膜各层厚度变化引起的分割挑战。

图4：图形推理块的示意图

图推理块由四个操作组成，分别为①投影，②图卷积，③逆投影，④融合。

首先，将原始特征投影到节点空间。然后对节点空间特征进行图卷积，提取全局节点特征。为了将全局节点特征与原始特征融合，在与原始特征融合之前，将全局节点特征反向投影回原始特征空间。

图5：多尺度全局推理模块的示意图

为了解决由于不同层之间视网膜厚度的巨大变化而导致的分割挑战，我们提出了一个多尺度全局推理模块（MGRM）来对所有九层视网膜的高级语义特征进行全局推理。MGRM由多尺度池算子和图推理块组成。它使用多个有效的感受野来捕捉和学习不同厚度视网膜层的特征。

MGRM将输入分为四条不同的路径，其中三条路径配备了不同内核大小的池层，其次是图形推理块：内核大小分别设置为2×2、3×3和5×5，将不同厚度的视网膜层信息编码到全局特征图中。然后，通过双线性插值对全局特征进行上采样，以匹配原始输入特征映射的大小。最后，将所有四条路径重新组合，并连接特征。整个过程如图5所示。

图6：MGU网络的结构示意图

MGU网络的结构包括编码器、MGR和解码器。跳跃连接（the skip-connections）将编码路径中的低级特征连接到解码路径中相应的高级特征。

跳跃连接的作用：在比较深的网络中，解决在训练过程中梯度爆炸和梯度消失问题。

Loss函数：

在两级分割框架中，提出了两个损失函数L_seg1和L_seg2来监督这个两级MGU-Nets，并让它们以端到端的方式更精确地分割视盘和九个视网膜层。总损失函数如下：

其中𝜆 权衡两个损失，Lseg定义为骰子损失和交叉熵损失之和。

𝑔𝑖 及𝑝𝑖 指示像素预测的基本事实和概率𝑥 属于阶级的𝑖，𝑀 是分割网络中的类数。

Experimental details（实验细节）

①初始学习速率为0.001；

②学习速率每20次迭代后降低一个数量级；

③迭代次数为50次；

④动量和权重衰减系数分别设置为0.9和0.0001；

⑤优化算法采用Adam optimizer；

⑥mini batches of size设为1；

⑦公式（1）中的 𝜆 设为2；

⑧训练集：验证集：测试集=6：2：2。

Ablation study（消融实验）

①one-stage baseline；

②two-stage baseline；

③在多尺度全局推理模块中不加图形推理块；

④在两级分割框架中使用单尺度全局推理模块。

Evaluation Metrics（评估指标）

DSC和PA用于定量评估分割性能。数学表达式如下：（X是预测的区域，Y是参考标准的区域）

DSC用来衡量预测和参考标准之间的重叠。

PA用于计算预测结果的真阳性率，并与它们的基本真值进行比较。

图7：正常视神经乳头周围OCT图像的分割对比图

（a） 原始图像（b） 参考标准（c） U-Net的预测（d） ReLayNet的预测（e） DRUNET的预测（f） 我们提出的方法的预测。黄色箭头指出了典型的散射标签，而层间不连续性用白色星星标记。提供了放大视图以实现更好的可视化。

图8：视网膜病变的视神经乳头周围OCT图像分割对比图

（a） 原始图像（b） 参考标准（c） U-Net的预测（d） ReLayNet的预测（e） DRUNET的预测（f） 我们提出的方法的预测。散乱的标签用黄色箭头表示，层间不连续用白色星形标记，视网膜病变用白色虚线圈出。同样提供了放大视图以实现更好的可视化。

图9：Duke SD-OCT数据集上分割结果的可视化对比图

（a） 原始图像（b） 地面真相（c） U-Net的预测（d） ReLayNet的预测（e） DRUNET的预测（f） 我们提出的方法的预测。分散的标签用黄色箭头表示，不连续处用白色星形标记。

局限性

1、收集的数据集的大小有限

合格的人体数据很难获取，并且在一张OCT图像中手动注释10层非常昂贵且耗时。

措施：对训练集进行数据增强，包括水平翻转、加性高斯白噪声和对比度调整。

2、图像伪影和标签噪声对分割结果的影响

众所周知，常见的伪影（包括血管阴影和视网膜病变）可能会影响自动分割算法。这些伪影通常会导致层边界模糊、组织纹理减弱和图像对比度改变，这可能导致分割精度降低。

措施：对伪影区域进行检测、移除和修复，并在未来处理标签噪声。

3、提出的框架依赖于标准化图像

它要求输入图像得到很好的标准化，以便我们所做的所有解剖假设或空间约束都是有效的。以第一级视盘分割网络为例，它依赖于假定disc区域和non-disc区域按照图3所述的适当水平顺序排列。

措施：在分割之前进行一个标记的过程，目的是在视网膜层上标记视盘区域。

实验结果

研究团队在收集的数据集上验证了所提出的多尺度图网络辅助的两级分割网络框架。在这个数据集上的实验结果表明，与其他最新技术相比，该方法可以显著提高分割视网膜组织的性能。除视盘以外，在大多数视神经乳头组织中，所提出的方法均优于最新方法，所得结果的平均Dice比ReLayNet，U-Net和DRUNET分别提高了1.6％，1.5％和1.4％。

研究结论

医学图像中的先验知识对于医学图像的处理非常重要。本研究通过多尺度图网络辅助的两级分割网络学习到的解剖学先验知识，提升了视神经乳头周围的视网膜OCT图像的分割效果。

代码结构

data包：数据集

figs包：有三张论文中的图

model包：神经网络模型

result包：训练及测试结果，还有训练日志

utils包：工具包。logger-运行日志，loss-损失函数，utils-工具，vis-可视化。

main_os.py

main_ts.py

README.md：介绍文档

requirements.txt：环境配置要求

将数据集放进data包里，然后用命令行执行以下指令即可。（相关参数按需调整）

python main_ts.py --name tsmgunet -d ./data/dataset --batch-size 1 --epoch 50 --lr 0.001

python main_os.py --name osmgunet -d ./data/dataset --batch-size 1 --epoch 50 --lr 0.001

结果可以在result包里面看到，其中包括train和test两个文件夹。

1、train里面有model文件夹和dice_epoch.txt。

①model文件夹里面放着最佳权重和最后一次迭代的权重。

②dice_epoch.txt则是记录每次迭代的dice score。

2、test里面有vis文件夹和dice_mpa_epoch.txt。

①vis文件夹里面放着预测结果的图片。

②dice_mpa_epoch.txt则是记录每次迭代的dice和pa score。

代码运行可能遇到的错误

1、RuntimeError: cuda runtime error (38) : no CUDA-capable device is detected at ……

os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # 设置当前使用的GPU设备仅为0号设备  设备名称为'/gpu:0'

2、RuntimeError: DataLoader worker (pid(s) 13628, 7140) exited unexpectedly

错误原因：torch.utils.data.DataLoader中的num_workers错误

解决方法：将num_workers改为0即可，0是默认值。num_workers是用来指定开多进程的数量，默认值为0，表示不启用多进程。

在parse_args()方法中：

parser.add_argument('-j', '--workers', type=int, default=0)

DataLoader的函数定义如下：

DataLoader(dataset, batch_size=1, shuffle=False, sampler=None,num_workers=0, collate_fn=default_collate, pin_memory=False,
drop_last=False)

dataset：加载的数据集(Dataset对象)

batch_size：批量处理大小

shuffle:：是否将数据打乱

sampler： 样本抽样

num_workers：使用多进程加载的进程数，0代表不使用多进程

collate_fn： 如何将多个样本数据拼接成一个batch，一般使用默认的拼接方式即可

pin_memory：是否将数据保存在pin memory区，pin memory中的数据转到GPU会快一些

drop_last：dataset中的数据个数可能不是batch_size的整数倍，drop_last为True会将多出来不足一个batch的数据丢弃
于是将num_workers参数值改成了默认值 0，程序可以运行了

注：

将num_workers设置为0，程序报错，并提示设置环境变量 KMP_DUPLICATE_LIB_OK=TRUE 那你可以在设置环境变量 KMP_DUPLICATE_LIB_OK=TRUE 或者使用临时环境变量：（在代码开始处添加这行代码）即：

os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'

3、RuntimeError: cublas runtime error : the GPU program failed to execute at C:/w/1/s/tmp_conda_3.6_035809/conda/conda-bld/pytorch_1556683229598/work/aten/src/THC/THCBlas.cu:450

这个问题尚未解决，网上有多种解释，其中：①pytorch不正宗，要官网的pytorch；②gpu_ids要设置为0；③电脑显存不足。

可能比较异想天开地拿笔记本电脑跑，所以我认为我应该是③显存不足。

参考资料

最新研究 | 这个“先验知识”，可提升视神经乳头周围视网膜OCT图像的分割效果 (thepaper.cn)