[論文レビュー] Chest X-ray lung and heart segmentation based on minimal training sets

Balázs Maga|arXiv (Cornell University)|Jan 20, 2021

COVID-19 diagnosis using AI参考文献 16被引用数 4

ひとこと要約

本稿では、最小限の学習データを用いて、胸部レントゲン画像における肺および心臓の正確なセグメンテーションを実現する、双方向ConvLSTMとアテンションゲートを統合したU-Netベースのアーキテクチャ、Attention BCDU-Netを提案する。JSRTデータセットにおいて、98.1%のDiceスコアと95.2%のIoUを達成し、SOTA性能を発揮した。10枚の学習サンプルでの実験でも97.0%のDiceと92.2%のIoUを達成し、mixupデータ拡張により顕著に性能が向上した。

ABSTRACT

As the COVID-19 pandemic aggravated the excessive workload of doctors globally, the demand for computer aided methods in medical imaging analysis increased even further. Such tools can result in more robust diagnostic pipelines which are less prone to human errors. In our paper, we present a deep neural network to which we refer to as Attention BCDU-Net, and apply it to the task of lung and heart segmentation from chest X-ray (CXR) images, a basic but ardous step in the diagnostic pipeline, for instance for the detection of cardiomegaly. We show that the fine-tuned model exceeds previous state-of-the-art results, reaching $98.1\pm 0.1\%$ Dice score and $95.2\pm 0.1\%$ IoU score on the dataset of Japanese Society of Radiological Technology (JSRT). Besides that, we demonstrate the relative simplicity of the task by attaining surprisingly strong results with training sets of size 10 and 20: in terms of Dice score, $97.0\pm 0.8\%$ and $97.3\pm 0.5$, respectively, while in terms of IoU score, $92.2\pm 1.2\%$ and $93.3\pm 0.4\%$, respectively. To achieve these scores, we capitalize on the mixup augmentation technique, which yields a remarkable gain above $4\%$ IoU score in the size 10 setup.

研究の動機と目的

• 限られた学習データでも高い性能を示す、胸部レントゲン画像における自動肺および心臓セグメンテーションのための深層学習モデルの開発。
• アテンション機構と双方向特徴統合を組み合わせた新規なネットワークアーキテクチャを用いて、既存のSOTAモデルを上回るセグメンテーション精度の向上。
• 医療画像セグメンテーションにおける低データ環境下でのmixupデータ拡張の有効性の調査。
• 特に、訓練損失が安定しているにもかかわらず、検証指標の改善が遅れるという、小規模データ設定で観察される特異な訓練ダイナミクスの分析。

提案手法

• 双方向ConvLSTM層を統合した、U-Netの変種であるAttention BCDU-Netを提案し、スカイプ接続による特徴統合を強化する。
• アップサンプリング段階での関心を特徴に集中させるために、アテンションゲートを組み込むことで、局所化精度を向上させる。
• 入力画像とマスクの線形補間を用いて、ハイパーパラメータδを用いて合成学習サンプルを生成するmixupデータ拡張を採用する。
• ネットワークの深さと損失関数のハイパーパラメータチューニングによりモデル性能を最適化し、評価指標としてDiceスコアとIoUを用いる。
• 検証損失に基づく早期停止を適用したが、訓練損失は安定しているにもかかわらず、検証指標に遅延した改善が観察された。
• JSRTデータセットを用い、10、20、209枚の学習セットサイズでモデルを訓練および評価した。

実験結果

リサーチクエスチョン

• RQ1 10～20枚の学習サンプルでのみ、深層学習モデルが胸部レントゲン画像でSOTAのセグメンテーション性能を達成できるか？
• RQ2 訓練データが極めて限られた状況下で、mixup拡張はセグメンテーション性能の向上にどの程度効果的か？
• RQ3 小規模データ設定において、訓練損失が安定しているにもかかわらず、検証Diceスコアが遅れて向上する理由は何か？
• RQ4 双方向ConvLSTMとアテンションゲートの統合は、医療画像セグメンテーションにおける特徴統合をどのように向上させるか？
• RQ5 訓練セットサイズが増加するに従い、医療画像のマスク同士の類似性が高いため、拡張の必要性が減少するのか？

主な発見

• Attention BCDU-Netは、JSRTデータセット全体を用いた場合、Diceスコア98.1±0.1%、IoU 95.2±0.1%を達成し、以前のSOTA結果を上回った。
• 10枚の学習サンプルでのみ、Diceスコア97.0±0.8%、IoU 92.2±1.2%を達成し、極小データでも優れた性能を示した。
• 20枚の学習サンプルでは、Diceスコア97.3±0.5%、IoU 93.3±0.4%を達成し、より少ないデータセットに比べて一貫した改善が見られた。
• 10枚のサンプル設定において、δ=0.2でmixup拡張を適用した場合、no mixupに比べてIoUスコアが4%以上向上し、小規模データにおける一般化性能向上に不可欠な役割を果たしていることが示された。
• 10枚の学習サンプル設定において、エポック400以降で検証Diceスコアが急激に上昇したが、訓練スコアは安定しており、予期しない一般化ダイナミクスが観察された。
• 訓練セットサイズが増加するに従い、mixupの恩恵が薄れる傾向にあり、これは、より大きなデータセットではマスクの類似性が高いため、拡張の必要性が自然に低下する傾向にあることを示唆している。