【論文筆記】圖像分割和圖像配準聯合學習模型——DeepAtlas

本文是論文《DeepAtlas: Joint Semi-Supervised Learning of Image Registration and Segmentation》的閱讀筆記。

文章第一個提出了一個圖像配準和圖像分割聯合學習的網絡模型 DeepAtlas，該模型實現了弱監督的圖像配準和半監督的圖像分割。在圖像配準時使用圖像的分割標籤作爲監督數據，如果沒有分割標籤，則通過分割網絡產生；而經過配準後的圖像增加了在圖像分割時可利用的訓練數據的量，相當於是一種數據增強。該模型不僅在分割和配準的精度上有所提升，並且還可以在訓練數據有限的情況下實現較好的效果。

一、記號

• $I_m$：浮動圖像（moving image）
• $I_t$：目標圖像（target image）
• $\mathcal{F}_R$：配準網絡
• $\theta_r$：配準網絡的參數
• $\mathcal{F}_S$：分割網絡
• $\theta_s$：分割網絡的參數
• $u=\mathcal{F}_R(I_m,I_t;\theta_r)$：形變場
• $\phi^{-1}=u+id$：形變圖，其中 $id$ 是恆等變換
• $I_m^w=I_m\circ\phi^{-1}$：配準後的圖像
• $S_t$：目標圖像分割標籤
• $S_m^w=S_m\circ\phi^{-1}$：配準後圖像分割標籤

二、網絡結構

DeepAtlas 的目的是當數據集中只有少量的分割標籤可用時，通過聯合訓練來讓分割和配準實現較高的精度。

網絡的結構如上圖所示，藍色的實線表示弱監督的配準，黃色虛線表示半監督的分割。

文章在附件中給出了分割網絡和配準網絡的具體結構，如下圖左右兩圖所示：

1. 配準網絡

配準網絡的損失主要有三個損失函數組成：配準正則損失 $\mathcal{L}_r$，圖像相似度損失 $\mathcal{L}_i$ 和解剖損失（分割相似度損失） $\mathcal{L}_a$。配準正則損失 $\mathcal{L}_r$ 可以讓形變場 $\phi$ 變得光滑，圖像相似度損失 $\mathcal{L}_i$ 用來評價浮動圖像 $I_m$ 和配準後圖像 $I_m^w$ 之間的相似度，解剖損失（分割相似度損失） $\mathcal{L}_a$ 是目標圖像分割標籤 $S_t$ 和配準後圖像分割標籤 $S_m^w$ 之間的相似度損失。

如此一來，配準學習的過程可以由下式表示：
$\theta_{r}^{\star}=\underset{\theta_{r}}{\operatorname{argmin}}\left\{\mathcal{L}_{i}\left(I_{m} \circ \Phi^{-1}, I_{t}\right)+\lambda_{r} \mathcal{L}_{r}\left(\Phi^{-1}\right)+\lambda_{a} \mathcal{L}_{a}\left(S_{m} \circ \Phi^{-1}, S_{t}\right)\right\}$
其中 $\lambda_r,\lambda_a\geq0$。

2. 分割網絡

分割網絡的輸入是一張圖像 $I$，輸出相應的分割結果 $\hat{S}=\mathcal{F}_S(I;\theta_s)$，分割網絡的損失主要有兩個損失函數組成：解剖損失 $\mathcal{L}_a$ 和有監督分割損失 $\mathcal{L}_{sp}$。解剖損失和配準網絡中的相同，有監督的分割損失 $\mathcal{L}_{sp}(\hat{S},S)$ 是分割網絡的分割結果 $\hat{S}$ 和人工分割結果 $S$ 之間的相似度損失。但是浮動圖像 $I_m$ 和目標圖像 $I_t$ 的分割標籤的存在情況有多種可能，所以相應的損失函數也存在以下四種情況：
$\left\{\begin{array}{l} \mathcal{L}_{a}=\mathcal{L}_{a}\left(S_{m} \circ \Phi^{-1}, \mathcal{F}_{\mathcal{S}}\left(I_{t}\right)\right) \text { and } \mathcal{L}_{s p}=\mathcal{L}_{s p}\left(\mathcal{F}_{\mathcal{S}}\left(I_{m}\right), S_{m}\right), \text { if } I_{t} \text { is unlabeled; } \\ \mathcal{L}_{a}=\mathcal{L}_{a}\left(\mathcal{F}_{\mathcal{S}}\left(I_{m}\right) \circ \Phi^{-1}, S_{t}\right) \text { and } \mathcal{L}_{s p}=\mathcal{L}_{s p}\left(\mathcal{F}_{\mathcal{S}}\left(I_{t}\right), S_{t}\right), \text { if } I_{m} \text { is unlabeled; } \\ \mathcal{L}_{a}=\mathcal{L}_{a}\left(S_{m} \circ \Phi^{-1}, S_{t}\right) \text { and } \mathcal{L}_{s p}=\mathcal{L}_{s p}\left(\mathcal{F}_{\mathcal{S}}\left(I_{m}\right), S_{m}\right), \text { if } I_{m} \text { and } I_{t} \text { are labeled; } \\ \mathcal{L}_{a}=\mathcal{L}_{s p}=0, \text { if both } I_{t} \text { and } I_{m} \text { are unlabeled. } \end{array}\right.$
分割的學習過程可以由下式表示：
$\theta_{s}^{\star}=\underset{\theta_{s}}{\operatorname{argmin}}\left(\lambda_{a} \mathcal{L}_{a}+\lambda_{s p} \mathcal{L}_{s p}\right), \quad \lambda_{a}, \lambda_{s p} \geq 0$

三、實施細節

• 解剖相似度損失 $\mathcal{L}_{a}$ 和有監督的分割損失 $\mathcal{L}_{sp}$ 採用的是 soft multi-class Dice loss：

$\mathcal{L}_{\text {dice}}\left(S, S^{\star}\right)=1-\frac{1}{K} \sum_{k=1}^{K} \frac{\sum_{x} S_{k}(x) S_{k}^{\star}(x)}{\sum_{x} S_{k}(x)+\sum_{x} S_{k}^{\star}(x)}$

​ 其中 $k$ 表示分割標籤的下標，$x$ 是體素位置，$S$ 和 $S^*$ 是兩個要比較的分割標籤。

• 圖像相似度損失 $\mathcal{L}_i$ 採用的是正則化的互相關（NCC）：

$\mathcal{L}_{i}\left(I_{m}^{w}, I_{t}\right)=1-N C C\left(I_{m}^{w}, I_{t}\right)$

• 配準正則損失 $\mathcal{L}_r$ 採用的是彎曲能（bending energy）：

$\mathcal{L}_{r}(\mathbf{u})=\frac{1}{N} \sum_{\mathbf{x}} \sum_{i=1}^{d}\left\|H\left(u_{i}(\mathbf{x})\right)\right\|_{F}^{2}$

​ 其中 $||\cdot||_F$ 表示弗羅貝尼烏斯範數（Frobenius norm），$H(u_i(x))$ 是第 $i$ 個成分 $u(x)$ 的 Hessian 矩陣，$d$ 表示維度，$N$ 表示體素數。

在訓練時，會交替的訓練分割網絡和配準網絡，當一個網絡在訓練時，另一個網絡的參數保持不變，並且是每訓練配準網絡20次才訓練分割網絡1次，這是因爲分割網絡更容易收斂。

四、實驗結果