[論文レビュー] Anatomically-Controllable Medical Image Generation with Segmentation-Guided Diffusion Models

Nicholas Konz, Yuwen Chen+2 more|arXiv (Cornell University)|Feb 7, 2024

Medical Image Segmentation Techniques被引用数 4

ひとこと要約

本稿では、サンプリング中に部分的または完全な解剖的セグメンテーションマスクを条件として用いることで、解剖学的に現実的な医用画像生成を実現するセグメンテーションガイドド拡散モデルを提案する。新規のマスクアブレーショントレーニング戦略を用いることで、一部の構造がマスクされた状態でも、高精細で解剖学的に妥当な画像を生成する能力を学習し、特に乳腺MRIおよびCTデータセットにおいて、ControlNetなどの最先端モデルを上回る解剖学的リアリズムとマスク忠実度を達成した。

ABSTRACT

Diffusion models have enabled remarkably high-quality medical image generation, yet it is challenging to enforce anatomical constraints in generated images. To this end, we propose a diffusion model-based method that supports anatomically-controllable medical image generation, by following a multi-class anatomical segmentation mask at each sampling step. We additionally introduce a random mask ablation training algorithm to enable conditioning on a selected combination of anatomical constraints while allowing flexibility in other anatomical areas. We compare our method ("SegGuidedDiff") to existing methods on breast MRI and abdominal/neck-to-pelvis CT datasets with a wide range of anatomical objects. Results show that our method reaches a new state-of-the-art in the faithfulness of generated images to input anatomical masks on both datasets, and is on par for general anatomical realism. Finally, our model also enjoys the extra benefit of being able to adjust the anatomical similarity of generated images to real images of choice through interpolation in its latent space. SegGuidedDiff has many applications, including cross-modality translation, and the generation of paired or counterfactual data. Our code is available at https://github.com/mazurowski-lab/segmentation-guided-diffusion.

研究の動機と目的

• 拡散ベースの医用画像生成における劣悪なグローバルな解剖学的リアリズムの課題に対処すること。
• 画像生成中に部分的または不完全な解剖的セグメンテーションマスクを柔軟に条件として用いること。
• トレーニング中にマスクアブレーションを用いる自己教師付き学習を活用することで、無条件画像生成を改善すること。
• 事前に登録された画像合成や対照的分析などの応用を促進すること。
• 既存のマスク条件付きおよび無条件生成モデルを上回る解剖学的忠実度と入力マスクへの適合性を達成すること。

提案手法

• モデルは、各ノイズ除去ステップでセグメンテーションマスクに条件付けられた画像空間における拡散プロセスを用いる。
• 入力マスク内の欠落した解剖的クラスのすべての組み合わせにさらされる、新しいマスクアブレーショントレーニング戦略を採用する。
• このトレーニングアプローチにより、マスクされていない領域に現実的な解剖学的構造を推論・生成しつつ、提供された制約を忠実に従う能力が得られる。
• モデルは、実際の医用画像とその対応するセグメンテーションマスク上でエンドツーエンドにトレーニングされる。
• 生成画像のマスク忠実度を評価するために、補助的なセグメンテーションネットワークが用いられる。
• フレームワークは、部分的マスクを用いた条件付き生成と、空のマスクを用いた無条件生成の両方をサポートし、自己教師付き学習のおかげで後者においてもよりリアルな再現性が向上している。

Figure 1: Standard diffusion models (right) can fail to create realistic tissue, even if the overall image appears high-quality. From left to right: (1) randomly sampled real breast MRI images, (2) synthetic images created by our segmentation-guided diffusion model, and (3) by the same model but wit

実験結果

リサーチクエスチョン

• RQ1 各ノイズ除去ステップでセグメンテーションマスクのガイドを組み込むことで、拡散モデルがより解剖学的に現実的な医用画像を生成できるか？
• RQ2 不完全またはアブレーションされたマスクが提供された場合、モデルはマスクされていない領域に妥当な解剖学的構造を生成できるか？
• RQ3 マスクアブレーショントレーニングは、空のマスク（無条件）設定でも、モデルの現実的画像生成能力を向上させるか？
• RQ4 ControlNetなどの最先端モデルと比較して、モデルの解剖学的リアリズムとマスク忠実度のパフォーマンスはどの程度か？
• RQ5 モデルを用いて、後続の応用のためのペアドで事前に登録された医用画像を合成できるか？

主な発見

• 同じデータセットにおいて、本モデルは乳腺MRIでマスク忠実度のDiceスコアが0.9027、CT臓器データセットで0.8980を達成し、ControlNetの0.3636および0.1132を著しく上回った。
• 同様のデータセットにおいて、生成画像のセグメンテーションと実際の画像予測との重複率（Diceスコア）は0.8593および0.8797を記録し、高い解剖学的忠実度を示した。
• 空のマスク（無条件生成）でも、マスクアブレーショントレーニングによる自己教師付き学習のおかげで、標準的な無条件拡散モデルよりもより解剖学的に現実的な画像を生成した。
• 本モデルが生成した合成データでトレーニングされたセグメンテーションネットワークは、実データでトレーニングされたものとほぼ同等の性能を示し、Dice差が0.04未満であった。これは、生成画像の高品質なリアリズムを裏付けた。
• 一部の構造（例：血管および線維嚢胞組織）に制約を加えた場合でも、モデルは多様で解剖学的に妥当な画像を生成でき、一方で乳房形状など他の構造については変動を許容した。
• 著者らは、コミュニティ研究を支援するため、コードベースとペアドで事前に登録された生成乳腺MRIデータセットを公開した。

Figure 3: Breast MRI (top) and CT organ (bottom) images generated with different segmentation-conditional and unconditional models. Samples from conditional models ( blue ) are below their input mask, and unconditional model ( red ) samples are random. For breast MRI, the breast, blood vessel, and f