掃碼關注下方公眾號："Python編程與深度學習"，領取配套學習資源，並有不定時深度學習相關文章及代碼分享。
今天分享一篇發表在CVPR2020上的關於醫學圖像處理的論文：Deep Distance Transform for Tubular Structure Segmentation in CT Scans (原文鏈接：[1])。
1 研究背景
醫學圖像中的管狀結構分割（如CT掃描中的血管分割）是使用計算機輔助早期篩查相關疾病的重要步驟。但是目前CT掃描中管狀結構的自動分割由於存在對比度差、噪聲大、背景復雜等問題而仍然存在很大的挑戰。同時，如下圖（Figure 1）所示，管狀結構其實可以由一系列圓心和半徑不斷變化的球體組成的。受此啟發，這篇文章嘗試將這一幾何特點融入到管狀結構的分割任務中，以提高其分割結果的准確性。
這篇文章主要有如下貢獻：
- 提出了一種幾何感知的管狀結構分割方法——深度距離變換(Deep Distance Transform, DDT)，其融合了骨架化的經典距離變換(the classical distance transform for skeletonization)方法和現代深度分割網絡
- 在6個醫學影像數據集上測試了DDT的性能
2 方法
2.1 整體流程
如上圖（Figure 2）所示，DDT在訓練階段和測試階段有些許差異。
訓練階段：如上圖綠色框所示，DDT的輸入是3D的CT掃描圖$X$，產生兩個預測結果：分割概率圖$P$和半徑分類圖$\hat{Z}$，並且分別根據標簽$Y$和$Z$進行監督學習。
測試階段： 如上圖藍色框所示，DDT的輸入是3D的CT掃描圖$X$，同樣產生$P$和$\hat{Z}$兩個預測結果，之后通過$P$中產生偽骨架$S$，再結合$\hat{Z}$上對應位置的預測值作為半徑，解析出最終的預測結果$\hat{Y}$
上述分析顯示，DDT的訓練需要有兩個標簽數據：分割圖標簽$Y$和半徑分類圖標簽$Z$，后面兩小節將對標簽的構造進行詳細介紹。
2.2 分割圖標簽$\hat{Y}$的構造
輸入數據為3D的CT掃描圖$X$ ($L\times W\times H$)，其對應的坐標集合為$V=\{v|v\in N_L\times N_W\times N_H\}$。定義$X$上位置$v$的值為$x_v=X(v)$，則$Y$中落在管狀結構中的位置$v_{in}$置為$y_{v_{in}}=1$，其余位置$v_{others}$置為$y_{v_{others}}=0$，因此有$y_v\in\{0,1\}$，是一個二類分割任務。
2.3 半徑分類圖標簽$\hat{Z}$的構造
上述的分割圖標簽構造和其它的3D分割方法是一致的，這篇文章的創新在於融入了半徑分類圖的預測，其作用是為每一個$y_v=1$的位置$v$預測一個半徑，從而可以根據Figure 1的幾何結構產生分割結果，再和模型預測的分割圖綜合起來得到最終的預測結果。
以位置$v$為中心的球體對應的半徑定義為$v$到最近的分割表面的距離$d_v$。先定義分割表面：
$$C_v=\{v|y_v=1,\exists u\in \mathcal{N}(v),y_u=0\}$$
其中$\mathcal{N}(v)$表示位置$v$的6個鄰居位置（3D空間中1個位置有6個鄰居），如果鄰居位置中存在$y_u=0$的，則說明該位置$v$在分割表面上。則$d_v$可以定義為：
$$d_v=\left\{\begin{matrix}\min_{u\in C_v}\left \|v-u\right \|_2,if\ y_v=1\\ 0, if\ y_v=0\end{matrix}\right.$$
上述公式表示：對於分割區域內的每一個位置$v$，計算其與分割表面$C_v$中距離最近的點作為其在半徑分割圖上$Z(v)$的值，對於分割區域外的位置沒有半徑表示，$Z(v)=0$。
但是注意到，這樣計算出來$d_v$是連續的實數，而[3]中提到，醫學影像數據的標注中普遍存在標注誤差，如果將$d_v$的學習視為回歸問題來訓練網絡，將由於標注存在的誤差而導致網絡難以收斂或者產生不穩定的預測結果，如[4]中所提到的。因此文中將$d_v$取整得到離散的數值$z_v\in \{0,\cdots,K\}$，從而將半徑分類圖標簽$Z$的學習是視為一個$K$分類問題，即有$Z(v)=z_v$，使得模型更容易訓練也具有更好的性能。
2.4 損失函數
如2.1小節中所介紹的，DDT的預測有兩部分：分割概率圖$P$和半徑分類概率圖$\hat{Z}$，因此在訓練中需要分別對這兩部分計算損失值，再通過反向傳播進行訓練。
對於分割圖部分，采用加權交叉熵損失函數 (weighted cross-entropy loss function)進行計算：
$$\mathcal{L}_{cls}=-\sum_{v\in V}\left(\beta_py_vlog\ p_v(W, w_{cls})+\beta_n(1-y_v)log\ (1-p_v(W, w_{cls}))\right )$$
其中$W$表示骨干網絡 (network backbone)的參數，$w_{cls}$表示分割圖預測分支的參數，$p_v(W,w_{cls})$表示位置$v$被預測為管狀結構體素的概率值。此外有$\beta_p=\frac{0.5}{\sum_v y_v},\beta_n=\frac{0.5}{\sum_v(1-y_v)}.$
對於半徑分類圖部分，其損失函數表示為：
$$\mathcal{L}_{dis}=-\beta_p\sum_{v\in V}\sum^K_{k=1}\left (\mathbf{1}(z_v=k)\left (log\ g_v^k(W,w_{dis})+\lambda\omega_vlog\ (1-\max_lg_v^l(W,w_{dis}))\right )\right )$$
其中$W$表示骨干網絡 (network backbone)的參數，$w_{dis}$表示半徑分類圖預測分支的參數。$\mathbf{1}(\cdot)$是指示函數，當$z_v=k$的時候返回1，否則返回0。$\lambda$是平衡兩個函數項的參數，在文中直接設置$\lambda=1$，$g_v^k(W,w_{dis})$表示位置$v$的半徑被預測為屬於第$k$個類別的概率，$\omega_v$是標准化權重，定義為$\omega_v=\frac{|arg\ max_lg_v^l(W,w_{dis})-z_v|}{K}$
最終的損失函數表示為：
$$\mathcal{L}=\mathcal{L}_{cls}+\mathcal{L}_{dis}$$
訓練的目標表示為：
$$(W^*,w_{cls}^*,w_{dis}^*)=arg\ min_{W,w_{cls},w_{dis}}\mathcal{L}$$
2.5 幾何學層面的精調 (Geometry-aware Refinement)
受[5]的啟發，文中提出在幾何學層面的精調方法 (gometry-aware refinement, GAR)，具體步驟如下：
1. 偽骨架生成 (Pseudo skeleton generation)：預測分割圖$P$通過閾值$T^p$過濾出偽骨架$S$，具體地有：$p_v>T^p,s_v=1;otherwise,s_v=0$
2. 形狀重建 (Shape reconstruction)：如Figure 1所示，對於每一個位置$v$，其對應的半徑為$\hat{z}_v=arg\ max_kg_v^k$，文中通過高斯分布生成對應的球體，多個球體合起來即可得到軟化后的結構形狀$\tilde{Y}^s$ (soft reconstructed shape)：$\tilde{y}_v^s=\sum_{u\in\{u'|s_{u'}>0\}}c_u\Phi(v;u,\Sigma_u)$，其中$\Phi(\cdot)$表示多元高斯分布的密度函數，$u$表示均值，$\Sigma_u=\left (\frac{\hat{z}_u}{3}\right )^2I$表示協方差矩陣。此外，采用標准化因子$c_u=\sqrt{(2\pi)^3det(\Sigma_u)}$放縮$\Phi(\cdot)$的結果。這一步可以得到基於$P$和$\hat{Z}$得到的形狀重建結果$\tilde{Y}$
3. 精調分割結果 (Segmentation refinement)：對於模型預測的分割概率圖$P$，可以采用上一步得到的結果$\tilde{Y}$進行調整：$\tilde{y}_v^r=\sum_{u\in \{u'|s_{u'}>0\}}p_u c_u\Phi(v;u,\Sigma_u)$，然后采用閾值過濾$\tilde{y}^r_v>T^r,\hat{y}_v=1$,$otherwise,\hat{y}_v=0$得到最終預測結果$\hat{Y}$
3 實驗結果
這里我只給出論文中的部分實驗結果，具體的實驗結果分析以及實驗和參數的設置請看原文。
4 參考資料
[1] http://openaccess.thecvf.com/content_CVPR_2020/papers/Wang_Deep_Distance_Transform_for_Tubular_Structure_Segmentation_in_CT_Scans_CVPR_2020_paper.pdf
[2] Fethallah Benmansour and Laurent D. Cohen. Tubular structure segmentation based on minimal path method and anisotropic enhancement. International Journal of Computer Vision, 92(2):192–210, 2011.
[3] Yan Wang, Yuyin Zhou, Peng Tang, Wei Shen, Elliot K. Fishman, and Alan L. Yuille. Training multi-organ segmentation networks with sample selection by relaxed upper confident bound. In Proc. MICCAI, 2018.
[4] Rasmus Rothe, Radu Timofte, and Luc Van Gool. Deep expectation of real and apparent age from a single image without facial landmarks. International Journal of Computer Vision, 126(2-4):144–157, 2018.
[5] Wei Shen, Kai Zhao, Yuan Jiang, Yan Wang, Xiang Bai, and Alan L. Yuille. Deepskeleton: Learning multi-task scale-associated deep side outputs for object skeleton extraction in natural images. IEEE Trans. Image Processing, 26(11):5298–5311, 2017.