[論文レビュー] Uncertainty-based graph convolutional networks for organ segmentation refinement

Roger D. Soberanis-Mukul, Nassir Navab+1 more|arXiv (Cornell University)|Jun 5, 2019

Advanced X-ray and CT Imaging被引用数 14

ひとこと要約

本論文は、CTボリュームからのCNN出力における臓器セグメンテーションを精緻化するための、未知の不確実性を考慮したグラフ畳み込みニューラルネットワーク（GCN）フレームワークを提案する。モンテカルロドロップアウトを活用して高信頼度・低信頼度ボクセルを特定し、高信頼度予測をラベル付きノードとして使用する半教師ありグラフを構築する。その後、GCNが不確実な領域を精緻化することで、ベースラインのU-Net予測に比べて膵臓で1%、脾臓で2%のDiceスコア向上を達成した。

ABSTRACT

Organ segmentation in CT volumes is an important pre-processing step in many computer assisted intervention and diagnosis methods. In recent years, convolutional neural networks have dominated the state of the art in this task. However, since this problem presents a challenging environment due to high variability in the organ's shape and similarity between tissues, the generation of false negative and false positive regions in the output segmentation is a common issue. Recent works have shown that the uncertainty analysis of the model can provide us with useful information about potential errors in the segmentation. In this context, we proposed a segmentation refinement method based on uncertainty analysis and graph convolutional networks. We employ the uncertainty levels of the convolutional network in a particular input volume to formulate a semi-supervised graph learning problem that is solved by training a graph convolutional network. To test our method we refine the initial output of a 2D U-Net. We validate our framework with the NIH pancreas dataset and the spleen dataset of the medical segmentation decathlon. We show that our method outperforms the state-of-the art CRF refinement method by improving the dice score by 1% for the pancreas and 2% for spleen, with respect to the original U-Net's prediction. Finally, we discuss the results and current limitations of the model for future work in this research direction. For reproducibility purposes, we make our code publicly available.

研究の動機と目的

• 医用CTスキャンにおける形状の多様性や組織の類似性に起因する臓器セグメンテーションにおける誤検出（FP）および誤検出（FN）を是正すること。
• モデルの不確実性を活用して、CNNベースのセグメンテーションの信頼性を向上させること。
• ボリュームセグメンテーションの後処理戦略として、半教師ありグラフ学習を検討すること。
• ベースCNNの再トレーニングなしに推論時のみで動作するフレームワークを構築すること。
• 不確実性が極めて重要となる低データ環境における手法の性能を評価すること。

提案手法

• 2次元U-Netにモンテカルロドロップアウト（MCDO）を適用し、モデルの不確実性を推定し、予測信頼度マップを生成する。
• 不確実性のしきい値に基づいて、ボクセルを高信頼度の前景／背景および低信頼度（潜在的FP/FN）領域に分類する。
• 不確実性に依存する空間的制約を用いて、臓器周囲に3次元形状適合型の領域 of interest（ROI）を定義する。
• スパarsな、半ラベル付きのグラフを構築する：高信頼度ボクセルはラベル付きノード、低信頼度ボクセルはラベルなしノードとする。
• ラベル付きノードを用いて半教師ありの方法で学習されたグラフ畳み込みネットワーク（GCN）を用い、ラベルなしノードのラベルを予測する。
• 訓練済みGCNを用いて全グラフを評価し、ノード表現に空間的・強度的特徴および期待値特徴を統合することで、精緻化されたセグメンテーションを取得する。

実験結果

リサーチクエスチョン

• RQ1 CNNからの不確実性推定が、臓器セグメンテーションにおける誤検出および誤検出領域の特定を改善できるか？
• RQ2 高信頼度予測を教師として用いることで、半教師ありグラフ学習が低信頼度セグメンテーション領域を効果的に精緻化できるか？
• RQ3 本手法のGCN精緻化法は、CRFのようなベースライン精緻化手法に比べて、セグメンテーション精度で優れているか？
• RQ4 低データ環境において、GCN精緻化戦略はCNN自身の期待値（平均予測）を上回る性能を示すか？
• RQ5 本フレームワークは、ベースCNNの再トレーニングなしにプラグイン型の後処理モジュールとして適用可能か？

主な発見

• 提案されたGCN精緻化手法は、元のU-Net予測に比べて膵臓セグメンテーションで1%、脾臓セグメンテーションで2%のDiceスコア向上を達成した。
• NIH膵臓データセットにおいて、本手法は最先端のCRFベースの精緻化手法をDiceスコアの観点で上回った。
• 低データ環境（例：10例のトレーニングボリューム）において、GCN精緻化はCNN自身の期待値よりも高いDiceスコアを達成しており、不確実性推定の信頼性が低い場合でも頑健であることが示された。
• 相対的改善度分析から、20例中の13例の膵臓ボリュームでGCN精緻化によりDiceスコアが期待値と同等以上に向上した。特にモデルが不十分に学習された場合に顕著であった。
• GCNはCNNの期待値に起因するアーチファクトに起因する誤検出領域を低減でき、モデルの誤りを是正できる能力を示した。
• 本手法は2次元および3次元CNNと互換性があり、ベースネットワークの再トレーニングなしに推論時のみで動作する後処理モジュールとして利用可能である。