[論文レビュー] DiffuseExpand: Expanding dataset for 2D medical image segmentation using diffusion models

Shitong Shao, Xiaohan Yuan+4 more|arXiv (Cornell University)|Apr 26, 2023

Generative Adversarial Networks and Image Synthesis被引用数 8

ひとこと要約

DiffuseExpandは拡散確率モデルを用いて、2D医用画像分割の高品質な対ペアの画像・マスクサンプルを生成し、データ拡張を効率化し分割性能を向上させる。マスク優先の合成、分類器ガイダンス、加速サンプリング、最終的な品質フィルタの4段階パイプラインを導入。

ABSTRACT

Dataset expansion can effectively alleviate the problem of data scarcity for medical image segmentation, due to privacy concerns and labeling difficulties. However, existing expansion algorithms still face great challenges due to their inability of guaranteeing the diversity of synthesized images with paired segmentation masks. In recent years, Diffusion Probabilistic Models (DPMs) have shown powerful image synthesis performance, even better than Generative Adversarial Networks. Based on this insight, we propose an approach called DiffuseExpand for expanding datasets for 2D medical image segmentation using DPM, which first samples a variety of masks from Gaussian noise to ensure the diversity, and then synthesizes images to ensure the alignment of images and masks. After that, DiffuseExpand chooses high-quality samples to further enhance the effectiveness of data expansion. Our comparison and ablation experiments on COVID-19 and CGMH Pelvis datasets demonstrate the effectiveness of DiffuseExpand. Our code is released at https://github.com/shaoshitong/DiffuseExpand.

研究の動機と目的

• 医療画像分割におけるデータ不足の課題（プライバシーとラベリング制約による）を解決する。
• 多様性と整合性の高い対となる画像とマスクサンプルを生成する拡散モデルベースのパイプラインを開発する。
• 段階的条件付け、分類器ガイダンス、サンプルスクリーニングを通じてデータ拡張品質を向上させる。
• COVID-19胸部X線およびCGMH Pelvisデータセットで有効性を実証する。
• オープンソースコードを提供し、アブレーションと比較によって改善を検証する。

提案手法

• ステージI: 事前学習済み拡散確率モデル（DPM）を分類器ガイダンスなしで微調整し、分布を医用画像へ写像する。
• ステージII: セグメンテーショナーを訓練して分類器ガイダンスを可能にし、DiceおよびBCE損失を用いて画像とマスクの整合性を評価する。
• ステージIII: DPM Solver++を用いて複数の条件下で画像-マスク対を合成し、分類器ガイダンスに基づく条件付きサンプリングを可能にする。
• ステージIV: 学習済みネットワークとDice閾値を用いて低品質サンプルを破棄し、高品質な拡張データを確保する。
• 条件付きサンプリングの連続時間確率流ODE定式化を導入し、tauという温度パラメータで勾配の大きさを制御する分類器ガイダンスを活用する。
• 分類器なしサンプリングと分類器ありサンプリングを比較するアブレーション、tau/sの多様性とDice性能への影響の評価を提供する。

実験結果

リサーチクエスチョン

• RQ1 拡散モデルは2D医用画像分割のために多様でよく整列した画像-マスクペアを生成できるか？
• RQ2 段階的条件付けと分類器ガイダンスは拡張分割データセットの品質と有用性を改善するか？
• RQ3 サンプルスクリーニング（ステージIV）はDiceスコアとサンプル多様性にどのような影響を及ぼすか？
• RQ4 医用画像領域における拡散ベースの拡張手法はGANベースやデータ拡張ベースの手法と比較してどうか？

主な発見

• DiffuseExpandは高忠実度かつ多様なImage-Maskペアを生み出し、試験データセットでDiceスコアを向上させる。
• 訓練済みセグメンテータと分類器ガイダンスを組み合わせた分類器ガイダンスと、ステージIVのスクリーニングが画像とマスクの整合性を高める。
• ガイダンスのためのtauベース条件付けは、同等の設定で勾配スケール値sを用いる場合よりもDiceスコアを高くする。
• 拡張はDiceベースの分割性能においてSynthmedおよびXLsorを上回ることが、試験設定下で示された。
• ステージIVのスクリーニングはサンプル多様性を低減するが全体的な品質を高め、下流の分割に寄与する。