[CorrMLP] Correlation-aware Coarse-to-fine MLPs for Deformable Medical Image Registration

第61回 コンピュータビジョン勉強会＠関東（前編） 2024/07/07
https://kantocv.connpass.com/event/321174/
#cvsaisentan #CVPR2024

での発表資料です。

Transcript

1. Correlation-aware Coarse-to-fine MLPs for Deformable Medical Image Registration 第61回 コンピュータビジョン勉強会@関東 2024/07/07（⽇） 飯⽥啄⺒
2. ⾃⼰紹介 2020.4 - 2024.3 中国系AIベンチャーで⾃動運転のCV 2024.4 - 現在 医療AIベンチャーでCV 飯⽥啄⺒＠京都
3. 紹介する論⽂ 2枚のフレームをレジストレーション（位置合わせ）する論⽂ 略称 ︓CorrMLP 発⾳ ︓こらぁえむえるぴー 動画 ︓https://www.youtube.com/watch?v=Rq2EAVr_KQM コード︓https://github.com/MungoMeng/Registration-CorrMLP ランク︓Oral & Award Candidate この画像に合わせる オリジナル 変形
4. 医療画像レジストレーション（＝位置合わせ） レジストレーションってどんなタスク︖できると何が嬉しい︖ Ø 標準的なレイアウトの⽅が医師の診断がやりやすい Ø 撮影条件によるズレの修正 Ø 動作処理 何もしないと⼼臓、呼吸などにより形状が⼤きく変わってしまう Ø 患者の変化を捉えられる Ø 画像ガイド⼿術 Ø モーダルの異なるデータを統合して、医師の診断精度を⾼められる 出典︓Hammoudeh, Ahmad, and Stéphane Dupont. “Deep Learning in Medical Image Registration: Introduction and Survey.” arXiv, January 10, 2024. https://doi.org/10.48550/arXiv.2309.00727. Moving Image Fixed Image Moving Image Fixed Image
5. やり⽅ Deformable Registration Rigid Registration 教師なし学習 剛体変換では、変換⾏列のパラメータを推定。 ⾮剛体変換では、変位場（Displacement Field）を推定。 出典︓ Chen, Junyu, Yihao Liu, Shuwen Wei, Zhangxing Bian, Shalini Subramanian, Aaron Carass, Jerry L. Prince, and Yong Du. “A Survey on Deep Learning in Medical Image Registration: New Technologies, Uncertainty, Evaluation Metrics, and Beyond.” arXiv, September 12, 2023. https://doi.org/10.48550/arXiv.2307.15615. ① 剛体変換 ⼤きな動きをまずあわせる ② ⾮剛体変換 細かいズレをあわせる ⼈体は⾮剛体なので、基本的には変位場を学習/推定する。 実際には…2つとも使ってる︕（現状） ⼤きな位置ずれは、アフィン変換でざっくりと位置合わせし、 細かいずれは、⾮剛体変換を⾏う 出典︓著者スライド 今回紹介するCorrMLPは、あらかじめFreeSurfer/ FLIRTというソフトで アフィン変換済みのデータに対して、レジストレーションを⾏っています。
6. レジストレーション⼿法の変遷 2020 2008 2018 2024 2022 SyN (2008) 2枚のペア画像から最適化問題を解く 両側から変形させて、⼀貫性をもたせる。 さらに、逆にも辿れるようにする。 最適化の時代 ⽋点︓その場で最適化するので遅い︕ Avants, B. B., C. L. Epstein, M. Grossman, and J. C. Gee. “Symmetric Diffeomorphic Image Registration with Cross-Correlation: Evaluating Automated Labeling of Elderly and Neurodegenerative Brain.” Medical Image Analysis 12, no. 1 (February 2008): 26–41. https://doi.org/10.1016/j.media.2007.06.004. Balakrishnan, Guha, Amy Zhao, Mert R. Sabuncu, John Guttag, and Adrian V. Dalca. “VoxelMorph: A Learning Framework for Deformable Medical Image Registration,” September 1, 2019. https://doi.org/10.1109/TMI.2019.2897538. DNNの時代 DNNを取り込んだモチベーション - データ⾊々使える - 事前学習重みを利⽤できる - 推論時の最適化を避けられる 2014 2012 2016 2010 VoxelMorph (2018) CNNを使ったレジストレーション⼿法の先駆け U-Netで変位場を出⼒し、STNsで変形。 問題を簡単にしたり⾊々やってる U-Net, STNs アフィン変換 メッシュグリッドで ⾮剛体変換も可能 STNs (2015) 空間的な変換パラメータを予測
7. DNNベース以降の本論⽂の位置づけ 2020 2018 2024 2022 VoxelMorph (CVPR2018) CNN 局所的な変位場 しか作れない Transformer 計算量が多く フル解像度で計算困難 TransMorph (MedIA2021) MLP CorrMLP (CVPR2024) NICE-Net (MICCAI2022) Swin-VoxelMorph (MICCAI2022) NICE-Trans (MICCAI2023) Transformer化 Transformer化 TransMorph ⾼解像度ではTransformerを外している RCN (ICCV2021) 反復化 RCN MLP化 反復除去 Chen, Junyu, Eric C. Frey, Yufan He, William P. Segars, Ye Li, and Yong Du. “TransMorph: Transformer for Unsupervised Medical Image Registration.” Medical Image Analysis 82 (November 2022): 102615. https://doi.org/10.1016/j.media.2022.102615. Zhao, Shengyu, Yue Dong, Eric I.-Chao Chang, and Yan Xu. “Recursive Cascaded Networks for Unsupervised Medical Image Registration,” 10600–610, 2019. https://openaccess.thecvf.com/content_ICCV_2019/html/Zhao_Recursive_Cascaded_Networks_for_Unsupervised_Medical_Image_Registration_ICCV_2019_paper.html. 紹介する論⽂
8. CorrMLPの超概略 Conv Conv Conv CMW-MLP CMW-MLP CMW-MLP 𝑖𝑚𝑔!"# 𝑖𝑚𝑔!"#$% 中間特徴量 MaxPooling x1/2 MaxPooling x1/2 ステップレベルの結果改善 エンコーダの特徴をレジストレーション 中間特徴量 変位場 変位場 概略 Ø MLPを利⽤したレジストレーション⼿法CorrMLPを提案 Transformerベースのモデルより、計算効率がよく、⻑距離依存性を捉えることができる Ø Coarse-to-Fine戦略により、複数解像度の特徴マップの局所的な相関を捉えられる CorrMLPのイメージ図 ※ 発表者が作成
9. 全体のアーキテクチャ 全体的にはU-Net的な構造になっている。 エンコーダでマルチスケールの特徴を抽出し、デコーダでは逆順に受け取りながら解像度を上げていく。 2x2x2 Max Pooling 3x3x3 Conv x 2 ... 変位場
10. 全体のアーキテクチャ 全体的にはU-Net的な構造になっている。 エンコーダでマルチスケールの特徴を抽出し、デコーダでは逆順に受け取りながら解像度を上げていく。 2x2x2 Max Pooling 3x3x3 Conv x 2 ... 変位場 2つの⼊⼒の空間的な対応関係を探索 （あとで解説）
11. 全体のアーキテクチャ 全体的にはU-Net的な構造になっている。 エンコーダでマルチスケールの特徴を抽出し、デコーダでは逆順に受け取りながら解像度を上げていく。 2x2x2 Max Pooling 3x3x3 Conv x 2 ... 前回の変位場を受取り、 2つのCMW-MLPブロックを通して、改善していく 変位場 2つの⼊⼒の空間的な対応関係を探索 （あとで解説）
12. 全体のアーキテクチャ 全体的にはU-Net的な構造になっている。 エンコーダでマルチスケールの特徴を抽出し、デコーダでは逆順に受け取りながら解像度を上げていく。 2x2x2 Max Pooling 3x3x3 Conv x 2 ... 𝝍を使って𝑭𝒎 を動かして⼊⼒ (STNsでやってる） 前回の変位場を受取り、 2つのCMW-MLPブロックを通して、改善していく 変位場 2つの⼊⼒の空間的な対応関係を探索 （あとで解説）
13. 全体のアーキテクチャ 全体的にはU-Net的な構造になっている。 エンコーダでマルチスケールの特徴を抽出し、デコーダでは逆順に受け取りながら解像度を上げていく。 2x2x2 Max Pooling 3x3x3 Conv x 2 ... 𝝍を使って𝑭𝒎 を動かして⼊⼒ (STNsでやってる） 画像レベル 2つの⼊⼒の相関を再計算 ステップレベル 前ステップと現在ステップの特徴間 前回の変位場を受取り、 2つのCMW-MLPブロックを通して、改善していく 変位場 2つの⼊⼒の空間的な対応関係を探索 （あとで解説）
14. 全体のアーキテクチャ 全体的にはU-Net的な構造になっている。 エンコーダでマルチスケールの特徴を抽出し、デコーダでは逆順に受け取りながら解像度を上げていく。 2x2x2 Max Pooling 3x3x3 Conv x 2 ... 𝝍を使って𝑭𝒎 を動かして⼊⼒ (STNsでやってる） 画像レベル 2つの⼊⼒の相関を再計算 ステップレベル 前ステップと現在ステップの特徴間 前回の変位場を受取り、 2つのCMW-MLPブロックを通して、改善していく 変位場 残差変位場 CMW-MLPブロックでは、前ステップの変位場の修正量を学習 2つの⼊⼒の空間的な対応関係を探索 （あとで解説）
15. 全体のアーキテクチャ 全体的にはU-Net的な構造になっている。 エンコーダでマルチスケールの特徴を抽出し、デコーダでは逆順に受け取りながら解像度を上げていく。 2x2x2 Max Pooling 3x3x3 Conv x 2 ... 𝝍を使って𝑭𝒎 を動かして⼊⼒ (STNsでやってる） 画像レベル 2つの⼊⼒の相関を再計算 ステップレベル 前ステップと現在ステップの特徴間 前回の変位場を受取り、 2つのCMW-MLPブロックを通して、改善していく 変位場 残差変位場 CMW-MLPブロックでは、前ステップの変位場の修正量を学習 Coarse-to-Fineなアーキテクチャになっており、画像間、ステップ間の相関を活⽤したレジストレーションが反復なしで可能 2つの⼊⼒の空間的な対応関係を探索 （あとで解説）
16. CMW-MLP block (Correlation-aware Multi-Window MLP block) Multi-window MLP 重要なチャネルをハイライト LayerNorm → Conv → LeakyReLU → SE Channel Attention CMW-MLPブロックの利点 1. 局所的に⾮線形な変形ができる (gMLPの利点） ペアとなるピクセル同⼠は近傍に現れるはず 2. 複数の異なる変形量を表現できる マルチウィンドウのおかげ ココ 3D Correlation Layer カーネル3Dカーネル内の値の積和。 パターンが同じであれば⾼くなる。 出典︓Kang, Miao, Xiaojun Hu, Weilin Huang, Matthew R. Scott, and Mauricio Reyes. “Dual-Stream Pyramid Registration Network.” Medical Image Analysis 78 (May 1, 2022): 102379. https://doi.org/10.1016/j.media.2022.102379. GAPとSoftmax関数による チャネルの重み付け 複数のウィンドウサイズで gMLPブロックに通す
17. 損失関数 類似損失 ℒ!"# ターゲット画像𝐼$ と変形後の画像𝐼#∘& の類似度を測る損失 正規化相互相関（NCC, Normalized Cross-Correlation） お気持ち︓正規化することで、パターンに着⽬した類似度が測れる 正則化項 ℒ'() 変位場 𝜓 が滑らかになるようなペナルティ関数 CorrMLPの学習は次の損失関数で教師なし学習を⾏う。 𝑛!ボクセル内の平均 お気持ち︓勾配の急な変化を抑える→滑らかになる 平均を引く 標準偏差で割る 全体損失
18. 3D Brain MRI データセット これに 合わせる 合わされる 画像 各⼿法で合わせた画像たち Fixed image との差分 スプリット 枚数 データセット名 備考 train 2,656 ADNI ABIDE ADHD IXI val/test 100 Mindboggle 50 MRI for val, 50 MRI for test Buckner 40 MRI for test 前処理︓ Ø 強度正規化 Ø FreeSurfer/ FLIRT（アフィンレジストレーション） 1mm*の等⽅性ボクセルを持つMNI-152脳テンプレートに合わせる Ø 144x192x160のボクセルにクロップ 定性結果
19. 定量評価と評価指標︓3D Brain MRI Dice Similarity Coefficients (DSC)＝位置合わせ精度 Negative Jacobian Determinants (NJD)＝変換の滑らかさ トレードオフになりがち Transformerベースの⽅が変換CNNよりも変形させ やすい。受容野が広いためと考えられる。 coarse-to-fine戦略の⼿法はDSCが総じて⾼い CorrMLPはDSCもNJDも⾼い性能を達成 各ペアで毎回最適化問題を反復的に解く ※ ⼊⼒サイズ＝144x192x160のボクセル 1024x1024x4よりちょっと⼤きいくらいのサイズ感 DNN-based RegistrationはCPUでも意外に速い。 GPUも使えると10倍程度速い（⼊⼒依存）。
20. 定量評価︓4D cardiac cine-MRI MLPベースの⽅が Transformerベースよりも 変形に強い スプリット 枚数 データセット名 備考 train - ACDC cine-MRI, 150 patients 左⼼室、右⼼室、⼼筋のセグメンテーションラベル ペア画像には拡張末期と収縮末期のフレームを使う train/val/test=90/10/50 End-Diastole︓⼼臓が⼀番⼤きいとき、End-Systole︓⼼臓が⼀番⼩さいとき データセット 前処理︓ Ø 1.5x1.5x3.15mm*のボクセル間隔で再サンプリング Ø 中⼼の128x128x32のボクセルをクロップ Ø [0, 1]にMinMaxScale
21. Ablation Study アーキテクチャデザイン MLPブロックデザイン U-Net like MLP ︓ベースライン MLPMorph ︓VoxelMorph, TransMorphのようなアーキテクチャ Encoder=MLP, Decoder=CNN, Multi-window MLP (MW-MLP) block w/o Correlation-Layer CMW-MLP Blockを別物に置き換え。 Ø MLPMorphの時点でVoxelMorph, TransMorphを圧倒。 CorrMLPだともっと良い。 Ø VoxelMorph, TransMorphはConvでFull/Half解像度画像を処理するので、 ⾼解像度で広範囲に細かい特徴を捉えることが難しい。 Ø CorrMLPで各レベルの相関計算を除去すると、性能劣化 ※ NJDは同じ正則化設定だと同様のNJDになったので省略 Ø マルチウィンドウ（MW）の効果がある。 MW-MLPは相関レイヤがなくてもマルチウィンドウ化で他より も良い Ø CMW-MLPでブランチを削除すると性能劣化する 3D Correlation Layer がない
22. まとめ まとめ Ø MLPを利⽤したことで、計算効率が良く、⾼解像度でも広範囲なマッチングができる Ø 画像レベルとステップレベルでの相関計算を⾏ったCoarse-to-Fine戦略を提案。 Ø マルチウィンドウを利⽤した相関計算により、マルチレンジで相関計算を⾏える Award Candidateに選出されている理由（これ以外思いつきませんでした…） Ø フル解像度で⻑距離レンジの相関計算ができる Ø マルチウィンドウ化により、場所によって変形量が異なる問題にも対処している Ø DSC、NJDで両⽅同時にSoTAを達成している Ø ツッコミどころのほとんどない⼿法を提案している ウィンドウサイズをある程度決めなければいけないところは、好みではないが、妥当だと思う。
23. 参考⽂献 1. [CorrMLP] Meng, Mingyuan, Dagan Feng, Lei Bi, and Jinman Kim. “Correlation-Aware Coarse-to-Fine MLPs for Deformable Medical Image Registration,” 9645–54, 2024. https://openaccess.thecvf.com/content/CVPR2024/html/Meng_Correlation-aware_Coarse-to-fine_MLPs_for_Deformable_Medical_Image_Registration_CVPR_2024_paper.html. 2. [SyN] Avants, B. B., C. L. Epstein, M. Grossman, and J. C. Gee. “Symmetric Diffeomorphic Image Registration with Cross-Correlation: Evaluating Automated Labeling of Elderly and Neurodegenerative Brain.” Medical Image Analysis 12, no. 1 (February 2008): 26–41. https://doi.org/10.1016/j.media.2007.06.004. 3. [STNs] Jaderberg, Max, Karen Simonyan, Andrew Zisserman, and Koray Kavukcuoglu. “Spatial Transformer Networks.” arXiv, February 4, 2016. https://doi.org/10.48550/arXiv.1506.02025. 4. [TransMorph] Chen, Junyu, Eric C. Frey, Yufan He, William P. Segars, Ye Li, and Yong Du. “TransMorph: Transformer for Unsupervised Medical Image Registration.” Medical Image Analysis 82 (November 2022): 102615. https://doi.org/10.1016/j.media.2022.102615. 5. [RCN] Zhao, Shengyu, Yue Dong, Eric I.-Chao Chang, and Yan Xu. “Recursive Cascaded Networks for Unsupervised Medical Image Registration,” 10600–610, 2019. https://openaccess.thecvf.com/content_ICCV_2019/html/Zhao_Recursive_Cascaded_Networks_for_Unsupervised_Medical_Image_Registration_ICCV_2019_paper.html. 6. Kang, Miao, Xiaojun Hu, Weilin Huang, Matthew R. Scott, and Mauricio Reyes. “Dual-Stream Pyramid Registration Network.” Medical Image Analysis 78 (May 1, 2022): 102379. https://doi.org/10.1016/j.media.2022.102379. 7. Hammoudeh, Ahmad, and Stéphane Dupont. “Deep Learning in Medical Image Registration: Introduction and Survey.” arXiv, January 10, 2024. https://doi.org/10.48550/arXiv.2309.00727. 8. Chen, Junyu, Yihao Liu, Shuwen Wei, Zhangxing Bian, Shalini Subramanian, Aaron Carass, Jerry L. Prince, and Yong Du. “A Survey on Deep Learning in Medical Image Registration: New Technologies, Uncertainty, Evaluation Metrics, and Beyond.” arXiv, September 12, 2023. https://doi.org/10.48550/arXiv.2307.15615.