利用概率稠密位移網絡消除深度配準與常規配準之間的差距pddNet

https://github.com/multimodallearning/pdd_net

Closing the Gap between Deep and Conventional Image Registration using Probabilistic Dense Displacement Networks

摘要

診斷任務，手術圖像引導，放射治療以及運動分析在很大程度上依賴於準確的患者內對齊。患者間的配準可以依賴基於atlas風格或者landmark定位或者形狀分析。當標籤少並且解剖差異較大，傳統的註冊方法往往仍然優於深度學習方法，到目前爲止，深度學習方法主要處理相對較小或較低複雜性的變形。我們通過利用概率密集置換優化的思想來解決這個缺點，這種思想在許多大變形的配準任務中表現出色。我們建議設計一個具有近似最小卷積和平均場推理的網絡，用於在離散弱監督配準設置下的微分位移正則化。通過使用這些有意義並理論上已證明的約束，我們可學習的配准算法包含很少的可訓練權值(主要用於特徵提取)，並且通過很少的標記掃描更容易訓練。它在訓練和推斷方面非常快，並且在腹部CT具有挑戰性的患者間配準方面達到了最先進的準確性，比以前的深度學習方法有15%的Dice重疊。

關鍵字：registration、deep learning、probabilistic、 abdominal

1、簡介和相關工作

傳統的醫學圖像配準主要依賴於迭代，最小化不同多尺度的翹曲waping moving相似度的和正則化懲罰來實現。基於深度學習的圖像配準(Deep learning based image registration, DLIR)旨在模擬這一過程，通過訓練一個卷積網絡，該網絡可以預測給定兩次新的未訓練掃描的非線性對齊函數。因此，一個單一的前饋傳遞函數必須使用許多卷積層而不是多個warping步驟。DLIR的監督可以基於自動或手動對應、語義標籤或內在成本函數。它在時間敏感的應用方面具有巨大的潛力，例如通過多atlas配準的圖像引導、融合、跟蹤和形狀分析。然而，由於大空間的潛在變形，可以映射兩個對應的解剖到另一個，這個問題比圖像分割的約束少得多，因此仍然是一個開放的挑戰。

腦部局部變形爲幾毫米，並有大量的標記數據集。對於腹部、前列腺或肺部的其他解剖，形狀變化只有幾釐米，DLIR主要應用於不太複雜的病例。對於吸氣-呼氣肺配準，DLIR的準確性~2.5mm仍低於常規方法<1mm。弱監督的DLIR方法Label-Reg在腹部CT上，病患建對齊平均的Dice爲42.7%比NiftyReg算法Dice56.1%要低。

我們的假設是，如果不借助複雜的多階段翹曲管道warping pipelines，很難用深度連續迴歸網絡對不同患者之間大的形變進行建模。相反，離散配準的使用，即同時探索一個大的量化位移空間，已經被證明可以更有效地捕捉腹部和胸部的變形[5,12,16]，並且可以通過很少或單個的扭曲步驟實現。不出所料，DLIR的2D視覺中已經探索了離散位移設置:即FlowNet-C[2]。提出了一種不包含可訓練權值的相關層(見[2]中的Eq. 1)，通過使用產生441通道聯合特徵圖的密集量化位移空間(21×21像素偏移量)移動移動圖像，計算兩幅圖像的特徵相似度。接下來，我們學習了一個非常大的441(+32)×256×3×3的核(然後進一步卷積)，它忽略了位移空間的顯式四維幾何。因此，大量的優化參數導致了對監督訓練數據的巨大需求。將這個想法擴展到3D是非常困難的，因爲在稠密相關之後，維數增加到6D，儘管它有很多好處，但是還沒有被考慮到。概率和不確定性建模已經在DLIR中進行了研究，參考文獻[9,17]，但沒有在離散的環境中進行研究。

我們提出了一種新的DLIR學習模型，該模型通過引入帶有可微約束的強正則化來更好地利用概率密集位移採樣probabilistic dense dispalcement sampling 的優勢問題的本質。因此，我們使用均值場mean-field推斷進行正則化[8,18]，並使用近似最小卷積[3]計算標籤間的兼容性，從而從空間變換的擬閤中分離卷積特徵。我們的特徵提取器使用三維可變形卷積3D defoemable convolutions[4]，是非常輕量級的。據我們所知，這是將離散DLIR與均值場正則化的可微應用相結合的方法首選。與之前的工作相比，我們的模型需要更少的可訓練重量，捕獲更大的變形，並且可以從少量的標記掃描訓練到高精度。我們同時提出了一種新的非局部標籤損失來代替更廣泛使用的基於空間變換spatial transformer的損失。

2、方法

我們旨在通過在和上約束並學習特徵映射，找到空間變換，對齊和。爲了訓練處一個合適的特徵提取，該特徵提取對噪聲和信息強度不敏感。在訓練期間我們提供了一個監督標記，.我們把空間座標定義爲連續變量，使用三線性插值對離散網格進行採樣。在粗糙的網格上用(幾千個)控制點來參數化。位移範圍被限制在一個離散的位移空間discrete displacement space，其線性間距爲，其中q是定義捕獲範圍的標量，在我們的例子中，是3375。網絡模型預測出位移概率displacement probabilities的6D tensor,每個控制點的維數4-6的和是1。概率與位移的內積得到這些概率估計數的加權平均值，從而得到在推理過程中的三維位移。

1）卷積特徵學習網絡:爲了學習一個有意義的非線性映射，從輸入強度到稠密的特徵量(通道，步長爲3)，我們採用Obelisk
方法,一個三維可變形卷積與可訓練的偏移量，然後一個簡單的1×1MLP非常有效地捕捉空間上下文。我們通過在Obelisk層之前添加一個具有4個通道的普通5×5×5卷積核來擴展作者的實現，從而學習邊緣特徵。該網絡有64個空間濾波器偏移量，共有120k個可訓練參數，用於固定和移動掃描，生成和。

2）緻密位移不相關的相關層：第一部分提供了特徵表示，我們旨在找到正則化爲一場displacement field，將矢量vector 分配給每個控制點來進行非線性變換最大限度的提高了fixed和warped moving scan的（label）相似性。正如傳統的離散配準，和相關層correlation layer中描述的那樣，我們對位移搜索空間上的相似性度量進行密集的評估。通過學習描述符特徵維度的負均方差（mean squared error, MSE）獲得不相似dissimilarities 6D tensor 。可以使用差異度量例如相關係數correlation coefficient。由於控制點稀疏性，實驗中對位移相似度displacement similarity的評價要求不差過2GFlops。位移捕捉範圍q設置爲0.4.

3）使用最小卷積和均值場推理正則化：由於非線性配準通常是不適定的ill-posed，因此採用附加先驗保持形變空間平滑。對比其他在DLIR上的工作，原則上學習一個一個無約束的形變，並僅強制空間平滑作爲損失項，我們建議將正則化約束建模爲網絡的一部分。基於位移平方差常用於Markov 隨機場配準(Markov random field, MRF)，例如通過循環信念傳播優化（loopy belief propagation,LBP）。[7]和[18]將圖模型中的平滑約束集成到端到端的學習分割網絡。由於LBP需要更多迭代才的帶最優結果，因此並不合適作爲展開的網絡層，我們使用了用於[8]中的離散優快速平均場推斷（兩次迭代），在[8中使用了5次迭代。它由兩個交替的步驟組成：標籤兼容變換（獨立地作用域空間控制點）和基於濾波器的消息傳遞（使用平均池化層實現，步長爲1）。

正如[3]中所指出的那樣，稠密位移空間的擴散正則化可以使用以3D位移偏移爲根的拋物線的較低包絡來計算，其高度等於不同項的和與之前的平均場推斷迭代。這個較低的包絡線不能直接微分，但是我們可以得到非常精確的近似值，首先使用最小池化（stride=1）來找到局部極小值然後使用兩個平均代價池操作average pooling operation(stride = 1)來提供二次平滑。正如圖1中藍色快所示，我們的方法新正則化部分包括左右與3個位移維度（min-convolution）的最小和平均池層，以及作用於3個空間維（mean-field inference）的平均濾波。每個操作、放縮和偏差因素a1-a6之前介紹和優化功能層在端到端一起訓練

概率變換損失和標籤監督：我們可以進一步利用位移採樣的概率性質，並根據非局部加權[11]的方法指定我們的監督監督標籤損失項。即我們首先否定正則部分的輸出（縮放）到僞概率使用softmax計算位移。然後，在相同空間移位位置對移動分割的one-hot表示進行採樣，並將這些向量乘以估計概率，一計算作爲MSE的真實值（one-hot）分割的標籤損失。3D位移場的連續值是將概率估計與位移標籤相乘的加權平均值，然後對圖像分辨率進行三線性插值得到。擴散正則懲罰在位移場的所有的3維空間梯度是用戶定義用於平衡平滑變換（較低的Jacobians的標準偏差）和準確結構調整。

3 實驗驗證

爲了證明我們的方法能夠在不同病人腹部捕捉到非常大的形變，我們對CIS-CERAL3的10個經過對比增強的3D CT掃描的訓練數據，進行了3倍交叉驗證實驗，每9個解刨結構手動分割：肝臟liver、脾臟spleen、胰腺pancreas、膽囊gallbladder、膀胱unary bladder，右腎right kidney, 左腎left kidney,右腰大肌 right psoas major muscle(psoas)和左腰大肌left psoas。將圖像重採樣到各向同性體素大小爲1.5mm^3，尺寸爲233x168x286個體素，不需要人工預對齊。

我們將概率密度置換網絡(probabilistic dense displacement network ,pdd-net)與兩種傳統算法NiftyReg和deeds進行比較，這兩種算法在[16]的患者間腹部CT配準研究表現最佳，但DLIR尚未處理這一任務。NiftyReg與互信息mutual information和5級多分辨率方案一起用於捕獲大的形變運行時間40~50s。Deed考慮單尺度密度的移位空間（大概需要4-6s），然後擴展到離散優化的三等級（25-35秒的運行時）。下一步，我們在數據上訓練了弱監督DLIR方法Label-Reg在我們數據上（每個摺疊在24小時以上）。爲了將內存需求降低到32G以下，分辨率降低到2.2mm基本通道數量減半到16。進一步踐行了小的調整，以優化病人之間的訓練。我們實現了FlowNet-c的三維擴展在帶有Obelisk特徵提取的pytorch中，一個密集相關層a dense correlation layer和一個具有輸入通道的正則化網絡由5個batch-norma層高的PReLU。它有200萬個權值和輸出一個(非概率)三維位移場。爲了獲取更有意義的結果，需要在Obelisk層的中間輸出加入語義分割損失。我們提出的方法使用相同的特徵學習部分（200k參數），但現在使用min-convolution，mean-field推理（沒有語義指導）和non-local label loss只增加6個可訓練的權重（不是2百萬個）。在消融研究中分析了這三種選擇的影響，並考慮用手工自相似上下文特徵代替Obelisk特徵學習。我們使用擴散的正則項權重用於控制32^3的網格大小和fixed scans的整個仿射增強和訓練我們Adma網絡（學習率0.01）在1500次迭代大約90分鐘GPU顯存16G。我們實現了一個實例梯度下降優化器，優化了前饋預測。[1]也使用了這個思想，但在我們的例子中，它快了100倍（0.24s比24s），因爲我們可以直接操作預先計算的位移概率，並且不需要通過網絡進行迭代反向傳播。

表1。定量比較交叉驗證的VISCERAL解刨3個數據集的10個掃描，基於24個組合的測試掃描沒有在訓練集上看到（數字表示Dice得分）。我們的pdd-net性能優於其他DLIR方法Label-Reg和FlowNet-C，比傳統方法（NiftyReg和deeds）的優勢在15%。我們剝離研究顯示1）用於學習Obelisk特徵與手動自相似上下文（SSC）描述符。2）使用平均場推斷。3）使用心得非局部標籤損失。另外一個快速實例級優化實現了pdd-net+inst。 FlowNet-C是我們[2]的3D擴展，具有Obelisk特徵和可訓練的正則化網絡。比較配準前的Dice平均爲30%。

4、結果和討論

pdd-net的推理時間只是0.57秒，產生的位移場似是而非，Jacobian行列式標準差爲0.4小於1%摺疊體素(負雅可比矩陣)。表1顯示了Dice平均分數交叉驗證了24組配準。其中爲使用任何標記的訓練掃描進行任何評估的測試配準。我們的方法比Label-Reg和FlowNet-C兩種DL方法，以大約15%個點，只有30%的初始對齊的病人間挑戰達到56.7%的Dice係數。這比一個網格級別的傳統離散配準deeds要好10%。我們的實例優化（每掃描對）需要0.24秒，較少摺疊（少於0.6%）並且進一步將準確率提高到了58.4%，這是高於傳統的多級配準deeds和NiftyReg。

比較deeds+SSC和一個網格級pdd+SSC，到我們pdd+SSC的變體。它使用相同的自相似性的特點，只有適應的正則部分的α參數，我們得到一個相思精確和形變複雜性。這表明所提出的最小卷積和兩個平均場推斷的正則化層幾乎可以匹配全序列的能力[5]中MRF優化。使用弱監督學習的特點，結果增加20%以上的Dice。非局部損失項和我們實例微調，分別貢獻了5%和2%進一步收益的Dice重疊。考慮到具有更多可訓練權值的無約束FlowNet-C的效果更差，或者我們的變體只是用小卷積而不是進行空間域濾波，均值域推斷的重要性是顯而易見。的。我們實現了一個比常規配準更魯棒的對齊質量（Dice更低的標準差）。可視化配準示例如圖2所示，並作爲補充材料中的表面呈現視頻文件。

5、結論

我們的新型pdd-net將概率稠密位移與可微分的平均場正則化相結合，在7個更大的解剖病人間腹部CT配準，達到了70%Dice以上的one-to-one精度。它比之前的基於深度學習的圖像配準方法（Label-Reg和FlowNet-C），要高出15%，並且可以通過少量的標記臊面進行魯棒性訓練。它彌補了DLIR（在小訓練數據集)與傳統方法的質量差距，例如NiftyReg和deeds，同時快了(0.5秒)。我們的概念提供了一個新的潛在方向，使得DLIR在圖像引導介入、診斷和基於atlas形狀分析中的運用，超越了目前使用缺乏幾何可解釋性的像素分割網絡。未來的工作可以通過使用多個校準階段和控制點的更自適應採樣來獲得進一步的收益。對具有附加評價指標(表面距離)的較大數據集進行更詳細的評價，可以對該方法的優缺點提供更多的見解。