[论文解读] Localise to segment: crop to improve organ at risk segmentation accuracy

Abraham George Smith, Denis Kutnár+3 more|arXiv (Cornell University)|Apr 10, 2023

Advanced Neural Network Applications被引用 4

一句话总结

本文提出一种两阶段深度学习方法——先定位后分割——通过聚焦于相关解剖结构的图像裁剪，提升了放射治疗计划中危及器官分割的准确性。该方法减少了类别不平衡问题，加快了训练速度，并显著提升了Dice分数，尤其在脾脏、胰腺和心脏等较小器官上效果更明显。

ABSTRACT

Increased organ at risk segmentation accuracy is required to reduce cost and complications for patients receiving radiotherapy treatment. Some deep learning methods for the segmentation of organs at risk use a two stage process where a localisation network first crops an image to the relevant region and then a locally specialised network segments the cropped organ of interest. We investigate the accuracy improvements brought about by such a localisation stage by comparing to a single-stage baseline network trained on full resolution images. We find that localisation approaches can improve both training time and stability and a two stage process involving both a localisation and organ segmentation network provides a significant increase in segmentation accuracy for the spleen, pancreas and heart from the Medical Segmentation Decathlon dataset. We also observe increased benefits of localisation for smaller organs. Source code that recreates the main results is available at \href{https://github.com/Abe404/localise_to_segment}{this https URL}.

研究动机与目标

• 通过减轻医学图像数据集中类别不平衡的影响，提升放射治疗中危及器官分割的准确性。
• 探究两阶段定位-分割流程相较于单阶段全图分割是否能提升性能。
• 评估定位对不同器官大小下训练效率、收敛速度和模型稳定性的影。
• 评估定位是否对较小器官更具优势，因其面临更高的类别不平衡问题。
• 确定基于真实标注的定位是否在性能上持续优于在全分辨率图像上进行端到端训练。

提出的方法

• 定位和分割网络均采用3D U-Net架构，共享模型结构以简化实现。
• 在卷积层中应用零填充，以在训练过程中保持输入和输出的空间维度一致。
• 使用组归一化（group normalisation）替代批量归一化，以在小批量大小（每张GPU 2个）下稳定训练。
• 结合Dice损失与交叉熵损失，以应对类别不平衡问题，尤其对小器官至关重要。
• 在64×256×256大小的裁剪图像块上进行训练，以在64GB GPU内存限制下最大化上下文信息。
• 实施早停和模型检查点机制，以防止过拟合并提升训练效率。

Figure 1: Illustration showing the three different methods compared, including the baseline, two stage involving both a localisation network and organ segmentation network and the organ segmentation network that uses the ground truth to localise.

实验结果

研究问题

• RQ1 与单阶段全图训练相比，两阶段定位-分割流程是否能提升危及器官分割的准确性？
• RQ2 定位如何影响不同器官大小下的训练时间、收敛速度和模型稳定性？
• RQ3 定位的优势是否在类别不平衡更高的小器官中更为显著？
• RQ4 基于真实标注边界框的定位是否能持续优于在全图上进行端到端训练？
• RQ5 定位在多大程度上缓解了医学图像分割中类别不平衡带来的负面影响？

主要发现

• 两阶段定位方法显著提升了脾脏、胰腺和心脏的Dice分数，尤其在较小器官上增益最为明显。
• 使用定位裁剪图像进行训练，尽管训练两个网络，但相比全图块训练，显著缩短了训练时间并加快了收敛速度。
• 基线单阶段模型表现出不稳定的训练动态，包括波动的Dice分数，可能源于严重的类别不平衡以及图像块中前景区域的不一致性。
• 心脏（左心房）收敛速度最慢，归因于数据集较小和类别不平衡程度高，尽管定位提升了其性能。
• 胰腺尽管拥有最大的训练数据集，收敛率仍仅为85%，表明仅靠数据集大小无法克服类别不平衡的挑战。
• 基于真实标注的定位提供了持续的性能提升，表明在临床工作流中采用人工或交互式定位可能对高精度分割具有重要价值。