[论文解读] CT organ segmentation using GPU data augmentation, unsupervised labels and IOU loss

Blaine Rister, Darvin Yi+3 more|arXiv (Cornell University)|Nov 27, 2018

Radiomics and Machine Learning in Medical Imaging参考文献 18被引用 9

一句话总结

该论文提出了一种轻量级3D全卷积网络用于腹部CT器官分割，通过GPU加速的数据增强、肺部和骨骼的无监督形态学标注，以及交并比（IOU）损失函数，实现了高精度。该方法在类别不平衡的分割任务中优于加权交叉熵损失，五个器官的Dice分数达到78–92%，在消费级GPU上的平均推理时间为5.98秒。

ABSTRACT

Fully-convolutional neural networks have achieved superior performance in a variety of image segmentation tasks. However, their training requires laborious manual annotation of large datasets, as well as acceleration by parallel processors with high-bandwidth memory, such as GPUs. We show that simple models can achieve competitive accuracy for organ segmentation on CT images when trained with extensive data augmentation, which leverages existing graphics hardware to quickly apply geometric and photometric transformations to 3D image data. On 3 mm^3 CT volumes, our GPU implementation is 2.6-8X faster than a widely-used CPU version, including communication overhead. We also show how to automatically generate training labels using rudimentary morphological operations, which are efficiently computed by 3D Fourier transforms. We combined fully-automatic labels for the lungs and bone with semi-automatic ones for the liver, kidneys and bladder, to create a dataset of 130 labeled CT scans. To achieve the best results from data augmentation, our model uses the intersection-over-union (IOU) loss function, a close relative of the Dice loss. We discuss its mathematical properties and explain why it outperforms the usual weighted cross-entropy loss for unbalanced segmentation tasks. We conclude that there is no unique IOU loss function, as the naive one belongs to a broad family of functions with the same essential properties. When combining data augmentation with the IOU loss, our model achieves a Dice score of 78-92% for each organ. The trained model, code and dataset will be made publicly available, to further medical imaging research.

研究动机与目标

• 通过结合人工标注与无监督形态学标注（如肺部和骨骼），解决医学影像标注数据稀缺的问题。
• 通过采用简单且内存高效的模型，而非复杂架构，降低3D医学图像分割中的计算需求。
• 通过用基于IOU的损失函数替代标准交叉熵损失，提升在类别不平衡数据集（医学影像中常见）中的分割精度。
• 利用GPU纹理采样硬件加速数据增强，实现在训练过程中实时进行几何和光照变换。
• 公开提供一个数据集和训练好的模型，以支持未来多器官CT分割研究。

提出的方法

• 利用CUDA和纹理采样技术，通过GPU加速实现低延迟处理的3D数据增强，包括几何和光照变换。
• 通过3D傅里叶变换加速的形态学操作，无监督地生成肺部和骨骼的训练标签，减少人工标注工作量。
• 将肝脏、肾脏和膀胱的半自动人工标注与无监督标签结合，构建了包含130个扫描的训练数据集。
• 采用参数极少的轻量级3D全卷积网络（FCN），以降低内存占用并提升训练效率。
• 采用交并比（IOU）损失函数，经过数学分析表明其在类别不平衡分割任务中能提供均衡的误差惩罚。
• 证明IOU损失属于具有等效优化特性的函数族，且在实践中优于加权交叉熵损失。

实验结果

研究问题

• RQ1 在使用大量数据增强和IOU损失训练的情况下，一个简单的3D FCN是否能在不依赖复杂架构的前提下实现具有竞争力的器官分割性能？
• RQ2 与基于CPU的方法相比，GPU加速的数据增强在3D CT分割任务中的速度和训练效率如何？
• RQ3 无监督形态学操作在多大程度上能生成可靠训练标签（如肺部和骨骼），其效果与人工标注相比如何？
• RQ4 IOU损失函数是否在类别不平衡的医学图像分割任务中始终优于加权交叉熵损失，特别是在小尺寸或形状不规则的器官上？
• RQ5 一个单一的、通用型模型是否能以低计算成本实现对多个器官的高精度同步分割？

主要发现

• IOU损失函数在所有器官上显著优于加权交叉熵损失，尤其在减少少数类结构（如膀胱）的假阳性方面表现突出。
• 数据增强提升了模型泛化能力，在与IOU损失结合时，Dice分数最高提升10%；而与交叉熵损失结合时则导致性能下降。
• 在仅使用3mm³各向同性分辨率的130个扫描数据集上，该模型在五个器官（肝脏、肾脏、膀胱、肺部、骨骼）上实现了78–92%的Dice分数。
• 在LiTS肝脏分割挑战赛测试集上，最佳模型的平均Dice分数达到90.5%，在100多个提交结果中排名第49位。
• 该模型在仅使用2429 MB GPU内存的消费级硬件上，平均每个3D CT体积的处理时间为5.98秒，展现出良好的效率。
• 作者证明IOU损失并非唯一，而是属于具有等效优化特性的函数族，且在类别不平衡分割中能提供均衡的误差惩罚。