論文分享 3D RoI-aware U-Net for Accurate and Efficient Colorectal Tumor Segmentation

摘要

從三維磁共振（MR）圖像中分割結直腸癌區域是放射治療的一個關鍵步驟，傳統的放射治療方法要求精確地描繪腫瘤邊界，但要耗費大量的人力、時間和可重複性。雖然基於深度學習的方法在三維圖像分割任務中提供了良好的基線，但小的適用切片尺寸限制了有效的感受域，降低了分割性能。此外，大體積三維圖像的感興趣區域（RoIs）定位作爲一種先行操作，在速度、目標完整性、減少誤報等方面帶來了多重好處。與基於滑動窗口或非聯合定位分割的模型不同，我們提出了一種新的多任務框架3D-RoI-aware U-Net（3D-RU-Net），用於RoI定位和區域內分割，兩個任務共享一個骨幹編碼器網絡。利用編碼器得到的RoI，在編碼器的區域特徵中裁剪出多個層級的RoI，形成一個GPU存儲效率高的解碼器進行細節保護的分割，從而擴大了適用體積尺寸和有效的感受域。 爲了有效地訓練該模型，我們爲全局到局部的多任務學習過程設計了一個Dice形式的損失函數。基於所提出的方法所證明的效率增益，我們繼續集成具有不同接收場的模型，以獲得更高的性能，只需花費少量額外的計算費用。隨後對64例癌症患者進行了廣泛的實驗，並進行了四次交叉驗證，結果顯示在準確性和效率方面明顯優於傳統的先進框架。綜上所述，該方法由於其固有的可推廣性，在醫學圖像三維目標分割中具有巨大的推廣潛力。建議方法的代碼是公開的。

貢獻

1. 提出了一種三維聯合定位分割框架，包含共享的基於感興趣區域進行全局理解的Global Image Encoder，和利用設計的區域特徵金字塔進行RoI分割的Local Region Decoder。與競爭對手相比，該設計能夠充分利用大感受域，實現快速且高效的保留細節的整個體積的分割。
2. 考慮到類的自動重新平衡和更好的邊界識別，我們提出了一種基於Dice的全局到局部多任務混合損失（MHL）函數，以進一步提高精度。此外，加速框架鼓勵我們採用多感受域模型集成策略來抑制誤報，並以可接受的速度爲代價細化邊界細節。
3. 對採集到的數據集進行了大量的實驗，證明了我們提出的框架的有效性。此外，我們的方法本質上是通用的，可以應用於其他類似的應用。

網絡結構

3D RU-Net構建

We input whole image volumes to Global Image Encoder for multi-level feature encoding, employ an
encoder-only RoI locator for RoI localization, crop in-region feature tensors from multi-scale feature maps using RoI Pyramid Layer, and design a Local Region Decoder sub-network to perform multi-level feature fusion for high-resolution cancerous tissue segmentation.

1. Global Image Encoder

與構建一個完整的3D版的編碼器-解碼器架構（如3D FPN）或直接將流行的主幹網[49–51]擴展到3D不同，一個名爲全局圖像編碼器（Global Image encoder）的僅編碼的緊湊型網絡被構建來處理整個圖像，而不是像通常的做法那樣處理上下文受限的小部分圖像。

具體地說，編碼器使用ResBlocks[50]和MaxPooling層的堆棧來編碼整個捲圖像。每個Residual Block具有三個卷積層、三個Instance Normalization Layers[52]、三個ReLU層和一個Skip Connection，以獲得更好的梯度流動。在三維分割任務中，如果batchsize=1，則使用Instance Normalization Layers來提高魯棒性。

2. RoI Locator

RoI定位器是一個模板，可以是任何只包含編碼器的進行目標檢測的主幹網絡。由於大量訓練樣本的長寬比差異，學習精確的邊界框迴歸是很困難的。對於這個特定的三維語義分割任務，我們建議充分利用可用的體素級掩模，如下所述，以實現簡單和更穩健的邊界框預測。

具體來說，我們避免將體素級標籤降級爲對象級標籤來學習錨點擬合。該定位器以特徵映射$F_{III}$爲輸入，由核大小爲1的卷積層和Sigmoid激活函數組成。該模塊用於從全局圖像中預測下采樣的分割掩模。爲了解決極不平衡的前景與背景比率問題，定位器不是部分採樣，即對前景與背景以固定比例採樣或使用OHEM[53]，而是針對Dice loss進行訓練，這將在第2.2小節中介紹。然後，我們執行一個快速的3D連接性分析來計算所需的邊界框，公式爲$Bbox^{III}=(z^3，y^3，x^3，d^3，h^3，w^3)$，其中$(z^3，y^3，x^3)$表示起始座標，$(d^3，h^3，w^3)$表示$Bbox^{III}$在特徵地圖$F_{III}$中的深度、高度和寬度。

3. RoI Pyramid Layer

爲了提取檢測目標的RoI區域張量金字塔，我們首先從一個給定的邊界框$Bbox^{III}=(z^3，y^3，x^3，d^3，h^3，w^3)$構造一個邊界框金字塔$(Bbox^I，Bbox^{II}，Bbox^{III})$。具體地，邊界框金字塔由如下邊界框縮放準則迭代計算：$Bbox^{i-1}＝(z^i×s^i_z，y^i×s^i_y，x^i×s^i_x，d^i×s^i_z，h^i×p^i_y，w^i×s^i_x）$（其中$s^i_z，s^i_y，s^i_x$）表示$MaxPooling^i$的stride參數。在給定邊界框金字塔$(Bbox^I，Bbox^{II}，Bbox^{III})$的情況下，我們從整個圖像的特徵圖$F_{I}，F_{II}，F_{III}$中裁剪出原始的RoI張量金字塔$（f^I，f^{II}，f^{III}）$，without applying any bin-fitting operation，併爲後面的Local Region Decoder分支形成RoI張量金字塔。

4. Local Region Decoder

在給定RoI區域張量金字塔的基礎上，利用成功的多級特徵融合機制，構造了一個區域內分割子網絡，稱爲Local Region Decoder。解碼器的結構與具有跳躍連接的編碼器部分大致對稱，以融合相應尺度的特徵映射，而有益的區別在於解碼器分支的特徵張量的大小要小得多。由於RoI區域張量金字塔不包含形狀失真或尺度歸一化，因此該模塊可以在不丟失細節的情況下恢復RoI區域的空間維度。如果定位了多個RoI，則使用同一組解碼器權重迭代處理不同的RoI。

損失函數

Dice-based Multi-task Hybrid Loss Function

在多任務學習實踐中，每一個任務都面臨着不同的挑戰。在我們的例子中，全局圖像編碼器主要存在類不平衡的問題，而局部區域解碼器則必須關注目標區域的精確邊界。因此，我們提出了一個基於Dice的多任務loss（MHL）函數來有效地學習這些任務。

1. Dice Loss

where the sums are computed over the N voxels of the predicted volume $p_i∈ P$ and the ground truth volume $g_i∈ G$,$\epsilon$ is a minimal smoothness term that avoids division by 0 and is set as $10^{−4}$.

2. Dice Loss for Global Localization

爲了解決全局圖像RoI定位任務的類不平衡問題，我們採用了上述的Dice Loss：

where $P_{global}$ and $G_{global}$ denotes predictions of the localization top and down-sampled annotations.

3. Dice-based Contour-aware Loss for Local Segmentation

與定位任務相比，區域內分割分支需要多個約束才能獲得更好的邊界敏感分割結果。在語義分割實踐中，模糊邊界是最難學習的，但卻沒有引起足夠的重視。借鑑前面[43]增加輔助輪廓感知任務的探索，我們進一步利用Dice Loss來構造輔助任務，以幫助解決三維空間中輪廓標籤的極端稀疏性。實際操作中，我們在分割分支的輸出端增加一個由Sigmoid函數激活的1×1×1卷積層來預測輪廓體素，並與區域分割任務並行訓練。考慮到輔助任務，分割分支$L_{local}$的損失函數通過彙總加權損失表示如下：

where $λc= 0.5$, denoting the auxiliary task weight to ensure that the region segmentation task dominates while other tasks take effects.

最後，總體損失函數爲：

where $β = 10^{−4}$ denotes the balance of weight decay term and W denotes the parameters of the whole network.

Multiple Receptive Field Model Ensemble

由於單個模型的精度有限，多模型集成被認爲是一種有效的魯棒推理方法，在實際應用中得到了廣泛的應用，但代價是計算量大。在這種顯著加速的框架下，本文提出採用多感受域模型集成策略，融合結構相同但感受域設置不同的模型。這是對文獻[48]中提出的多分辨率策略的一種推廣，該策略將相同的感受域應用於具有不同空間分辨率的圖像，實際上是形成不同的空間感受域。這樣的推廣消除了細節丟失採樣，使每個模型對邊界細節的貢獻相等。具體來說，如表1所示，我們首先構造了一個感受域26×64×64的原始3D RU-Net，命名爲3D RU-Net-RF64。接下來，我們將ResBlock3的dilation rate調整爲2，將感受野擴展爲26×88×88，並建立了3D RU-Net-RF88；進一步將ResBlock2、ResBlock3和ResBlock4的dilation rate調整爲2，構建了一個26×112×112的3D RU-Net，命名爲3D RU-Net-RF112。在推理階段，如圖3所示，對三個網絡的輸出進行平均以生成最終預測。主要投票產生相似的分數，因此不討論。