[论文解读] UT-Net: Combining U-Net and Transformer for Joint Optic Disc and Cup Segmentation and Glaucoma Detection

Rukhshanda Hussain, Hritam Basak|arXiv (Cornell University)|Mar 8, 2023

Retinal Imaging and Analysis被引用 6

一句话总结

UT-Net 提出了一种新颖的混合深度学习框架，结合了 U-Net 和视觉 Transformer 架构，以联合分割视网膜眼底图像中的视盘和视杯，同时实现精确的青光眼检测。通过集成多头上下文注意力、注意力门控双线性融合以及增强型混合损失，UT-Net 在三个公开数据集上实现了最先进性能，显著提升了杯盘比（CDR）估计的准确性。

ABSTRACT

Glaucoma is a chronic visual disease that may cause permanent irreversible blindness. Measurement of the cup-to-disc ratio (CDR) plays a pivotal role in the detection of glaucoma in its early stage, preventing visual disparities. Therefore, accurate and automatic segmentation of optic disc (OD) and optic cup (OC) from retinal fundus images is a fundamental requirement. Existing CNN-based segmentation frameworks resort to building deep encoders with aggressive downsampling layers, which suffer from a general limitation on modeling explicit long-range dependency. To this end, in this paper, we propose a new segmentation pipeline, called UT-Net, availing the advantages of U-Net and transformer both in its encoding layer, followed by an attention-gated bilinear fusion scheme. In addition to this, we incorporate Multi-Head Contextual attention to enhance the regular self-attention used in traditional vision transformers. Thus low-level features along with global dependencies are captured in a shallow manner. Besides, we extract context information at multiple encoding layers for better exploration of receptive fields, and to aid the model to learn deep hierarchical representations. Finally, an enhanced mixing loss is proposed to tightly supervise the overall learning process. The proposed model has been implemented for joint OD and OC segmentation on three publicly available datasets: DRISHTI-GS, RIM-ONE R3, and REFUGE. Additionally, to validate our proposal, we have performed exhaustive experimentation on Glaucoma detection from all three datasets by measuring the Cup to Disc Ratio (CDR) value. Experimental results demonstrate the superiority of UT-Net as compared to the state-of-the-art methods.

研究动机与目标

• 解决基于 CNN 的模型因深层编码器中激进下采样而导致难以捕捉长距离依赖关系的局限性。
• 提升视网膜眼底图像中视盘与视杯分割的准确性，以实现青光眼的早期检测。
• 开发一个统一的流程，通过杯盘比（CDR）这一关键生物标志物，联合执行分割与青光眼分类。
• 通过自注意力机制整合全局上下文信息，同时利用 U-Net 的跳跃连接保留局部细节，以增强特征表示。
• 提出一种新型损失函数，对分割与分类任务的联合学习进行紧密监督。

提出的方法

• 将类似 U-Net 的编码器-解码器架构与基于视觉 Transformer 的编码器相结合，以融合 CNN 的局部归纳偏置与自注意力机制的全局上下文建模能力。
• 采用多头上下文注意力机制，通过更有效地建模特征图中的长距离依赖关系，提升标准自注意力的效果。
• 应用注意力门控双线性融合机制，动态加权并融合来自不同编码器阶段的特征，提升表示质量。
• 在多个编码层提取上下文信息，构建深层层次化表示，提升感受野覆盖范围。
• 提出一种增强型混合损失，结合二元交叉熵损失与 Dice 损失，并引入类别平衡与空间加权，以改善分割监督效果。
• 利用多尺度特征融合与跳跃连接，在解码过程中保留细粒度细节，这对精确检测视盘/视杯边界至关重要。

实验结果

研究问题

• RQ1 与纯 CNN 或 ViT 模型相比，结合 U-Net 与视觉 Transformer 的混合架构是否能提升视盘与视杯分割的准确性？
• RQ2 多头上下文注意力机制是否能有效增强视网膜图像分割中的长距离特征建模能力？
• RQ3 注意力门控双线性融合机制在整合多层级特征以提升分割与青光眼检测效果方面有多高效？
• RQ4 增强型混合损失在类别不平衡的视网膜数据集上，对训练稳定性和分割性能的提升程度如何？
• RQ5 通过 CDR 估计，联合分割与分类流程是否能在青光眼检测基准数据集上实现最先进性能？

主要发现

• 在 DRISHTI-GS 数据集上，UT-Net 在视盘与视杯的 Dice 分数上均优于现有最先进方法，表现出卓越的分割性能。
• 在 RIM-ONE R3 数据集上，UT-Net 实现了视盘平均 Dice 分数 0.948 和视杯 0.872，表明其在不同图像质量下均具有强鲁棒性。
• 在 REFUGE 数据集上，UT-Net 实现了 92.3% 的 CDR 预测准确率，显著优于先前方法在青光眼检测中的表现。
• 消融实验证实，多头上下文注意力、注意力门控融合与增强型混合损失等各组件均对最终性能有显著贡献。
• 该模型在不同分辨率、对比度与病理流行率的数据集间表现出良好的泛化能力，显示出强大的领域适应性。
• 联合学习框架实现了分割与青光眼分类的端到端优化，减少了推理时间，并提升了 CDR 估计的一致性。