[論文レビュー] PC-SwinMorph: Patch Representation for Unsupervised Medical Image Registration and Segmentation

Lihao Liu, Zhening Huang+3 more|arXiv (Cornell University)|Mar 10, 2022

Advanced Neural Network Applications被引用数 13

ひとこと要約

PC-SwinMorphは、2つの新しいピクセルベースの戦略を用いて、医療画像の登録とセグメンテーションのための統合的で教師なしフレームワークを提案する。1つは局所的およびグローバルな特徴表現を強化するピクセル単位の対照的学習法であり、もう1つはピクセル分割のアーティファクトを排除する3Dウィンドウ／シフトドウィンドウマルチヘッド自己注意モジュールである。本手法は最先端の性能を達成し、CANDIおよびLPBA40データセットにおいてベースラインモデルよりDiceスコアを最大4%向上させた。

ABSTRACT

Medical image registration and segmentation are critical tasks for several clinical procedures. Manual realisation of those tasks is time-consuming and the quality is highly dependent on the level of expertise of the physician. To mitigate that laborious task, automatic tools have been developed where the majority of solutions are supervised techniques. However, in medical domain, the strong assumption of having a well-representative ground truth is far from being realistic. To overcome this challenge, unsupervised techniques have been investigated. However, they are still limited in performance and they fail to produce plausible results. In this work, we propose a novel unified unsupervised framework for image registration and segmentation that we called PC-SwinMorph. The core of our framework is two patch-based strategies, where we demonstrate that patch representation is key for performance gain. We first introduce a patch-based contrastive strategy that enforces locality conditions and richer feature representation. Secondly, we utilise a 3D window/shifted-window multi-head self-attention module as a patch stitching strategy to eliminate artifacts from the patch splitting. We demonstrate, through a set of numerical and visual results, that our technique outperforms current state-of-the-art unsupervised techniques.

研究の動機と目的

• 教師なし医療画像登録およびセグメンテーションの限界を解消するため、特に監視が不足しているため微細な解剖的詳細を捉えられない点に焦点を当てる。
• 局所的およびグローバルな解剖的モデリングを向上させるために、ピクセルベースの空間構造を活用した特徴表現の向上を図る。
• 3D医療画像におけるピクセルベース表現に起因するアーティファクトを軽減するため、3D W-MSAおよびSW-MSAモジュールを用いた新しいピクセルステッチ機構を提案する。
• 真の変形場やセグメンテーションマスクを必要とせずに、登録およびセグメンテーションの両面で最先端の性能を達成する。
• ピクセルベースの表現が、複雑な医療画像タスクにおいてグローバルトークン埋め込みよりも意味を持つことを示す。

提案手法

• フレームワークはVoxelMorphをバックボーンとして用い、エンコーダーをピクセルベース表現戦略に置き換えることで特徴学習を強化する。
• ピクセル単位の対照的学習戦略を導入し、画像間のピクセルを比較することで局所的整合性とより豊かな特徴表現を強制する。
• エンコーダー／デコーダーの後段に3Dウィンドウ／シフトドウィンドウマルチヘッド自己注意（3D W-MSAおよび3D SW-MSA）モジュールを適用し、ピクセルをステッチして長距離依存性を回復する。
• 対照的損失とピクセルステッチを統合したフレームワークを構築し、同時に登録およびセグメンテーション性能を向上させる。
• 3D SW-MSAモジュールをメインネットワークの後段に1層のみ適用することで、計算コストを最小限に抑えつつ特徴の融合を最大化する。
• 真の変形場やセグメンテーションマスクを必要とせず、教師なし損失関数を用いてエンドツーエンドでモデルを訓練する。

実験結果

リサーチクエスチョン

• RQ1 グローバルトークンベースの手法と比較して、ピクセルベース表現は教師なし医療画像登録およびセグメンテーションを顕著に改善できるか？
• RQ2 ピクセル単位の対照的学習戦略を導入することで、特徴表現および微細な解剖的アライメントが向上するか？
• RQ3 3Dウィンドウ／シフトドウィンドウマルチヘッド自己注意機構は、ピクセル分割に起因するアーティファクトを効果的に低減し、ピクセル間の特徴交換を改善できるか？
• RQ4 ピクセル対照的学習とピクセルステッチの組み合わせは、実際の医療画像データセットにおける登録およびセグメンテーション性能をどの程度向上させるか？
• RQ5 本手法は、最先端の性能を達成しつつ、計算コストも効率的であると言えるか？

主な発見

• 提案されたPC-SwinMorphフレームワークは、CANDIおよびLPBA40データセットの両方でベースラインモデルより4%のDiceスコア向上を達成した。
• ピクセル単位の対照的学習戦略のみで、対照的学習ベースラインと比較してDiceスコアが2%向上し、特徴の豊かさ向上の有効性が示された。
• 3D W-MSAおよび3D SW-MSAモジュールの統合により、CANDIデータセットでは登録性能が3.06%、セグメンテーション性能が2.65%向上した。
• CANDIデータセットでは、セグメンテーションのDiceスコアが0.817、登録のDiceスコアが0.812に達し、すべての最先端の教師なし手法を上回った。
• 計算コストの増加は最小限であり、わずか0.6 GFlops（416.04から416.10 GFlops）の増加にとどまり、本手法の効率性を裏付けた。
• 可視化および数値的結果から、ピクセルベースアプローチがグローバル表現手法よりも微細な解剖的詳細をより強固に捉えられることを確認した。