[論文レビュー] 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

ひとこと要約

本論文は、GPUアクセラレートされたデータ拡張、肺および骨臓の自己教師付き形態的ラベリング、および交差乗算（IOU）損失関数を用いることで、高い精度を達成する軽量な3次元完全畳み込みネットワーク（FCN）を提案する。この手法は、特に不均衡なセグメンテーションタスクにおいて、重み付き交差エントロピー損失を上回り、5つの臓器で78–92%のDiceスコアを達成し、コンsumer 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.

研究の動機と目的

• 肺や骨のような臓器の手動ラベリングと、自己教師付き形態的ラベリングを組み合わせることで、医療画像データのラベル不足を緩和すること。
• 複雑なアーキテクチャの代わりに、単純でメモリ効率の良いモデルを用いることで、3次元医療画像セグメンテーションにおける計算負荷を低減すること。
• 医療画像で一般的な不均衡データセットにおいて、標準的な交差エントロピー損失の代わりにIOUに基づく損失関数を採用することで、セグメンテーション精度を向上させること。
• GPUテクスチャサンプリングハードウェアを活用したデータ拡張により、トレーニング中にリアルタイムで幾何的および光度的変換を実行できるように、データ拡張を高速化すること。
• 今後のマルチ臓器CTセグメンテーション研究を支援するため、公開可能なデータセットとトレーニング済みモデルを提供すること。

提案手法

• CUDAとテクスチャサンプリングを用いて、低遅延処理が可能なGPUアクセラレートされた3次元データ拡張（幾何的および光度的変換）を実装した。
• 3次元フーリエ変換を用いた高速化により、肺および骨の自己教師付きトレーニングラベルを形態的演算で生成し、手動ラベリングの負荷を低減した。
• 肝臓、腎臓、膀胱の半自動的手動ラベリングと、自己教師付きラベルを統合し、130スキャン分のデータセットを構築した。
• メモリ使用量を低減し、トレーニング効率を向上させるために、パrameter数が最小限の軽量な3次元完全畳み込みネットワーク（FCN）を採用した。
• 交差乗算（IOU）損失関数を用い、数学的に解析し、不均衡なセグメンテーションタスクにおいてバランスの取れた誤差ペナルティを提供することを示した。
• IOU損失が同様の最適化特性を持つ関数の一般化された家族に属しており、実際の応用において重み付き交差エントロピーを上回ることを実証した。

実験結果

リサーチクエスチョン

• RQ1 広範なデータ拡張とIOU損失関数を用いてトレーニングされた単純な3次元FCNは、複雑なアーキテクチャを用いなくても、競争力のある臓器セグメンテーション性能を達成できるか？
• RQ2 GPUアクセラレートされたデータ拡張は、CPUベースの方法と比較して、3次元CTセグメンテーションにおける速度とトレーニング効率においてどのように差を示すか？
• RQ3 自己教師付き形態的演算は、肺や骨のような臓器に対して信頼性の高いトレーニングラベルを生成できるか？また、手動ラベリングと比較してどの程度の性能を示すか？
• RQ4 IOU損失関数は、特に小形または不規則な形状の臓器において、不均衡な医療画像セグメンテーションタスクで一貫して重み付き交差エントロピーを上回るか？
• RQ5 1つの汎用モデルが、高い精度と低い計算コストで複数の臓器を同時にセグメンテーションできるか？

主な発見

• IOU損失関数は、膀胱のような少数クラス構造の誤検出（ファルスポジティブ）を著しく低減させ、全臓器において重み付き交差エントロピー損失を上回った。
• データ拡張によりモデルの一般化性能が向上し、IOU損失と組み合わせた場合、Diceスコアが最大10%向上したが、交差エントロピーと組み合わせた場合は性能が低下した。
• 3mm³の等方的解像度を用いて、130スキャン分のデータセットで、5つの臓器（肝臓、腎臓、膀胱、肺、骨）で78–92%のDiceスコアを達成した。
• LiTS肝臓セグメンテーションチャレンジのテストセットでは、最良のモデルが平均Diceスコア90.5%を達成し、100件以上の提出の中から49位となった。
• モデルは、コンsumerハードウェアで1つの3次元CTボリュームを平均5.98秒で処理し、GPUメモリ使用量は2429 MBにとどまり、効率性が確認された。
• 著者らは、IOU損失が一意的ではなく、同様の最適化特性を持つ関数の家族に属しており、不均衡なセグメンテーションにおいてバランスの取れた誤差ペナルティを提供することを実証した。