[论文解读] A two-stage 3D Unet framework for multi-class segmentation on full resolution image

Chengjia Wang, Tom MacGillivray+3 more|arXiv (Cornell University)|Apr 12, 2018

Advanced Image Processing Techniques参考文献 16被引用 60

一句话总结

本文提出一个两阶段拼接的 U-Net 框架，通过先定位 ROI 再在不重采样的情况下细化分割，对全分辨率的三维医学图像进行多类分割，在 MM-WHS 数据上优于最先进的 U-Net。

ABSTRACT

Deep convolutional neural networks (CNNs) have been intensively used for multi-class segmentation of data from different modalities and achieved state-of-the-art performances. However, a common problem when dealing with large, high resolution 3D data is that the volumes input into the deep CNNs has to be either cropped or downsampled due to limited memory capacity of computing devices. These operations lead to loss of resolution and increment of class imbalance in the input data batches, which can downgrade the performances of segmentation algorithms. Inspired by the architecture of image super-resolution CNN (SRCNN) and self-normalization network (SNN), we developed a two-stage modified Unet framework that simultaneously learns to detect a ROI within the full volume and to classify voxels without losing the original resolution. Experiments on a variety of multi-modal volumes demonstrated that, when trained with a simply weighted dice coefficients and our customized learning procedure, this framework shows better segmentation performances than state-of-the-art Deep CNNs with advanced similarity metrics.

研究动机与目标

• 推动在不通过下采样或后处理丢失细节的前提下实现高分辨率的 3D 全心分割。
• 开发一个两阶段的深度卷积神经网络，首先定位一个 ROI，然后在全分辨率下细化体素级标签。
• 利用受 SRCNN 启发的近端细化，结合跳连，提升分割精度。
• 在 MM-WHS 挑战中的多模态心脏 CT/MR 数据上证明有效性。
• 与最先进的 3D U-Nets 进行比较，并评估在数据受限情况下的训练策略以实现良好性能。

提出的方法

• 一个两级拼接的 U-Net 架构，其中 Net1 在下采样体积上提供粗略、以 ROI 为焦点的预测，Net2 在全分辨率下细化标签。
• Net1 在收缩路径中使用扩张的 5x5x5 卷积，扩张率逐步增大，以增大有效感受野。
• Net2 采用受 SRCNN 启发的设计，带跳连和类似递归的结构，输入为 Net1 输出和原始数据。
• 一个四步训练过程包括： (1) 使用全体积进行 ROI 定位，(2) 结合部分体积进行粗略多类分割， (3) 粗+细分割的端到端联合训练，(4) 在轴向切片的子体积上对 Net2 进行微调（K 个切片）。
• 使用加权多类 Dice 基损失（Eq. 2）以及用于 ROI 定位的前景导向 Dice 损失（Eq. 3–4）引导训练；使用 Adam 优化器，学习率为 0.0001。
• 数据增强和分层/子体积采样解决数据不平衡与内存限制；批次设计聚焦于中间体素类别，以平衡低频类别。

实验结果

研究问题

• RQ1 两阶段、全分辨率的 3D 分割框架是否能够在不进行后处理或重采样的情况下超越单次前向的 3D U-Net？
• RQ2 基于 ROI 的定位结合细化阶段是否能提升高分辨率心脏 CT/MR 数据的多类分割？
• RQ3 分步训练过程与专门的损失函数如何影响 MM-WHS 数据上的收敛性和准确性？
• RQ4 所提出的框架在 MM-WHS 基准测试中对于 CT 和 MR 两种模态是否与先进的 U-Net 相竞争？

主要发现

• Net2 在 CT 结构上的 Dice 和 Jaccard 分数均高于基线 3D U-Net。
• Net1 相对于下采样 U-Net 在若干结构上带来提升，尤其是在与 Net2 联用时。
• 在 MM-WHS 测试数据上，框架在 MR 结果非常具有竞争力（MR 的平均 Dice 约为 0.83），CT 结果相比基线显示显著提升；结果在较少的训练迭代和无额外后处理的情况下接近比赛冠军。
• 该方法在全分辨率下展现出有效的分割，验证了以 ROI 为驱动的两阶段方法用于高分辨率 3D 医学图像分割。
• 端到端训练结合粗损失和细损失比单独训练各自网络更快收敛、性能更好。