[论文解读] Semi-supervised Medical Image Segmentation through Dual-task Consistency

Xiangde Luo, Jieneng Chen+4 more|arXiv (Cornell University)|Sep 9, 2020

Advanced Neural Network Applications被引用 29

一句话总结

本文提出了一种基于双任务一致性的新型半监督医学图像分割框架，其中深度网络联合预测像素级分割图与几何感知的水平集函数。通过可微分的任务转换层，将水平集表示转换为分割图，双输出之间的一致性正则化利用未标注数据提升了性能，在两个3D医学影像数据集上实现了最先进结果，且计算成本更低。

ABSTRACT

Deep learning-based semi-supervised learning (SSL) algorithms have led to promising results in medical images segmentation and can alleviate doctors' expensive annotations by leveraging unlabeled data. However, most of the existing SSL algorithms in literature tend to regularize the model training by perturbing networks and/or data. Observing that multi/dual-task learning attends to various levels of information which have inherent prediction perturbation, we ask the question in this work: can we explicitly build task-level regularization rather than implicitly constructing networks- and/or data-level perturbation-and-transformation for SSL? To answer this question, we propose a novel dual-task-consistency semi-supervised framework for the first time. Concretely, we use a dual-task deep network that jointly predicts a pixel-wise segmentation map and a geometry-aware level set representation of the target. The level set representation is converted to an approximated segmentation map through a differentiable task transform layer. Simultaneously, we introduce a dual-task consistency regularization between the level set-derived segmentation maps and directly predicted segmentation maps for both labeled and unlabeled data. Extensive experiments on two public datasets show that our method can largely improve the performance by incorporating the unlabeled data. Meanwhile, our framework outperforms the state-of-the-art semi-supervised medical image segmentation methods. Code is available at: https://github.com/Luoxd1996/DTC

研究动机与目标

• 通过有效利用未标注数据，解决医学图像分割中的高标注成本问题。
• 探究任务级正则化是否能在半监督学习中优于传统的基于数据或网络级扰动的正则化方法。
• 通过水平集表示，将几何与全局形状约束整合到半监督分割中。
• 开发一种简单而高效的框架，避免复杂架构或迭代推理。
• 与现有方法相比，实现最先进性能的同时降低计算成本。

提出的方法

• 设计了一个双任务深度网络，同时预测像素级分割图与表示分割轮廓的全局水平集函数（LSF）。
• 通过平滑Heaviside层将水平集函数可微分地转换为分割概率图，支持端到端训练。
• 引入双任务一致性损失，最小化已标注和未标注数据中直接预测的分割图与由水平集函数推导出的分割图之间的差异。
• 一致性损失作为正则化项，强制两个不同但相关的任务分支之间的预测保持一致。
• 该框架每张图像仅需一次前向传播，避免迭代优化或多次推理阶段。
• 该方法可扩展至多类别分割，并可集成额外的可微任务，如边缘检测或关键点估计。

实验结果

研究问题

• RQ1 在半监督医学图像分割中，任务级一致性正则化是否能优于基于数据级或网络级扰动的正则化方法？
• RQ2 在半监督设置下，将几何感知的水平集表示与像素级分类结合，是否能提升分割性能？
• RQ3 在两个不同任务分支（像素级与全局级）之间强制预测一致性，是否能提升未标注数据上的泛化能力与鲁棒性？
• RQ4 一个简单、单次前向传播的框架是否能实现最先进性能，而无需复杂架构或训练方案？
• RQ5 与现有最先进半监督分割方法相比，该方法在准确性、参数效率和推理时间方面表现如何？

主要发现

• 所提方法在3D左心房MRI数据集和胰腺CT数据集上均实现了最先进性能，所有评估指标均优于现有半监督方法。
• 在左心房数据集中，该方法在Dice、ASD和95HD指标上均取得最佳结果，Dice分数显著高于次优方法。
• 在胰腺数据集中，该方法在所有对比方法中实现了最高的Dice分数和最低的95百分位Hausdorff距离，表明边界精度更优。
• 与SASSNet、UAMT和CCT等方法相比，该框架训练时间更短、参数量更少，同时性能更优。
• 可视化结果表明，与基线方法相比，该方法与真实标签重叠更高，假阳性更少，且能更好地保留解剖细节。
• 消融实验确认，双任务一致性是性能提升的关键因素，且该方法在不同解剖结构和成像模态间具有良好的泛化能力。