[論文レビュー] MaskSAM: Towards Auto-prompt SAM with Mask Classification for Volumetric Medical Image Segmentation

Bin Xie, Hao Tang+4 more|arXiv (Cornell University)|Mar 21, 2024

Advanced Neural Network Applications被引用数 4

ひとこと要約

MaskSAMは、3Dボリュメトリック医療画像分類のためのプロンプトフリーでマスク分類認識型のSegment Anything Model（SAM）のアダプテーションを提案する。可学習プロンプト生成器を導入し、補助的なマスク、ボックス、分類器トークンを生成するとともに、SAMのトランスフォーマーブロックに3DディープコンボリューションおよびディープMLPアダプタを統合することで、AMOS2022で90.52%のDiceスコアを達成し、nnUNetを2.7%上回る最先端の性能を実現した。

ABSTRACT

Segment Anything Model (SAM), a prompt-driven foundation model for natural image segmentation, has demonstrated impressive zero-shot performance. However, SAM does not work when directly applied to medical image segmentation, since SAM lacks the ability to predict semantic labels, requires additional prompts, and presents suboptimal performance. Following the above issues, we propose MaskSAM, a novel mask classification prompt-free SAM adaptation framework for medical image segmentation. We design a prompt generator combined with the image encoder in SAM to generate a set of auxiliary classifier tokens, auxiliary binary masks, and auxiliary bounding boxes. Each pair of auxiliary mask and box prompts can solve the requirements of extra prompts. The semantic label prediction can be addressed by the sum of the auxiliary classifier tokens and the learnable global classifier tokens in the mask decoder of SAM. Meanwhile, we design a 3D depth-convolution adapter for image embeddings and a 3D depth-MLP adapter for prompt embeddings to efficiently fine-tune SAM. Our method achieves state-of-the-art performance on AMOS2022, 90.52% Dice, which improved by 2.7% compared to nnUNet. Our method surpasses nnUNet by 1.7% on ACDC and 1.0% on Synapse datasets.

研究の動機と目的

• SAMの医療画像分類における限界、特に意味的ラベル予測の欠如と手動プロンプト依存性を解消する。
• 事前学習済み2D SAMを用いて、ゼロショットでプロンプトフリーの3D医療画像分類を実現する。
• 2D自然画像と3D医療ボリュームの間のドメインギャップを埋めるために、SAMのアーキテクチャを3D空間理解に適応させる。
• ユーザーが提供するプロンプトなしに、正確なマスクとその対応する意味的ラベルを同時に予測する統合フレームワークを導入する。
• タスク固有の再学習なしに、主な3D医療画像分類ベンチマークで最先端の性能を達成する。

提案手法

• SAMの画像エンコーダー内のマルチレベル特徴マップから、補助的なバイナリマスク、バウンディングボックス、分類器トークンを生成するプロンプト生成器を設計する。
• マスクデコーダーに可学習グローバル分類器トークンを導入し、補助分類器トークンと合算することで、各マスクごとの意味的ラベルを予測する。
• 各トランスフォーマーブロックに3Dディープコンボリューションアダプタ（DConvAdapter）を統合し、3Dインダクティブバイアスを持つ3D画像埋め込みを処理する。
• マスクデコーダーに3DディープMLPアダプタ（DMLPAdapter）を導入し、3Dプロンプト埋め込みを処理する。スイープ接続を備えた逆バトルネック構造を用いて特徴の精錬を向上させる。
• 画像エンコーダーおよびマスクエンコーダーの両方で、3D空間モデリングを強化するため、深さ位置埋め込み（DPosEmbed）を適用する。
• 予測されたセグメンテーションと真値セグメンテーションの間でバイパーティットマッチングを適用し、損失を計算することで、可変長の多クラスマスクを用いたエンドツーエンド学習を可能にする。

実験結果

リサーチクエスチョン

• RQ1 プロンプトフリーなSAMのアダプテーションは、手動によるポイントやボックスのプロンプトを必要とせず、高精度な3D医療画像分類を達成できるか？
• RQ2 SAMは元々バイナリマスク予測を目的として設計されているが、どのようにして分類されたマスクの意味的ラベルを予測できるように拡張できるか？
• RQ3 2Dで事前学習されたSAMが、ボリュメトリック医療画像からの3D空間情報を効果的に抽出・利用できるようにするためには、どのようなアーキテクチャ的変更が必要か？
• RQ4 3D対応アダプタ（DConvAdapterとDMLPAdapter）は、標準アダプタと比較して、3D一般化性能および分類精度をどのように向上させるか？
• RQ5 可学習プロンプトおよび分類器トークンを備えた統合的エンドツーエンドフレームワークは、nnUNetなどの既存の最先端手法をどれほど上回るか？

主な発見

• MaskSAMはAMOS2022データセットで90.52%のDiceスコアを達成し、nnUNetを2.7%上回る。
• ACDCデータセットでは、DiceスコアがnnUNetを1.7%上回り、心筋MRI分類への強い一般化能力を示した。
• Synapseデータセットでは、nnUNetを1.0%上回る性能を発揮し、多様な解剖的構造にわたる頑健性を確認した。
• アブレーションスタディの結果、平均化されたバウンディングボックスを用いたプロンプト生成器（MaskAvgBBoxPG）が最良の性能を示し、マスクのみの生成器よりも1.9%向上した。
• DMLPAdapterとDConvAdapterの統合により、ベースラインから0.3%の性能向上が得られ、DMLPAdapterが代替の深さMLP配置よりも優れた特徴精錬を示した。
• 深さ位置埋め込み（DPosEmbed）の追加により、ベースライン比で0.15%以上の性能向上が確認され、3D空間モデリングにおけるその役割が裏付けられた。