掃碼關注下方公眾號:"Python編程與深度學習",領取配套學習資源,並有不定時深度學習相關文章及代碼分享。
今天分享一篇發表在CVPR2020上的關於醫學圖像處理的論文:Structure Boundary Preserving Segmentation for Medical Image With Ambiguous Boundary (論文鏈接)。
1 研究背景
當前醫學影像分割任務存在兩個關鍵問題:
- 分割區域邊界的模糊性
- 缺乏專業領域知識的前提下分割區域的不確定性
這篇文章針對上述問題提出了解決方案,主要有如下三點貢獻:
- 提出一種新穎的邊界關鍵點選擇算法(boundary key point selection algorithm),使得能夠用關鍵點來表示分割目標域的結構邊界(解決上述第一個問題)
- 提出了一個新穎的結構邊界保持(structure boundary preserving segmentation)分割框架,並且使得模型能夠自動提取和學習結構關鍵點信息而不需要醫生的交互(解決上述第二個問題)
- 文中提出的分割框架可以應用到不同分割模型上,並且作者做了足夠的實驗進行驗證
文中的模型效果如下圖,相比於U-Net,作者提出的模型能夠更好的處理分割區域的邊界。
2 方法
2.1 整體流程
如上圖(Figure 2)所示,整個模型可以分解為三個主要部分:分割網絡、邊界保持模塊(Boundary Preserving Block, BPB)以及形狀邊界判別器(Shape Boundary-aware Evaluator, SBE)。
分割網絡: 可以由流行的語義分割網絡構成,例如U-Net等。
邊界保持模塊: 該模塊能夠嵌入到卷積網絡中,其思想是在卷積模塊中,根據其輸入特征產生對應的邊界點圖,而后將邊界點圖融入到下一個卷積模塊的特征輸入中,同時邊界點圖做一個監督學習。(感覺有點類似中間監督+注意力機制)
形狀邊界判別器: 模型最終產生的分割圖預測與中間卷積模塊中產生的邊界點圖可以通過形狀邊界判別器來約束它們之間表達的一致性。
這篇文章的創新主要在於邊界保持模塊和形狀邊界判別器,后面將主要對這兩個模塊進行介紹。
2.2 邊界關鍵點選擇算法(Boundary Key Point Selection Algorithm)
由於文中采用邊界關鍵點來表示目標邊界,因此在介紹邊界保持模塊之前,需要先了解怎么選擇邊界關鍵點。文中邊界關鍵點選擇算法的具體過程是:
- 采用常規邊界檢測算法[2]獲取標簽分割圖上目標區域的邊界
- 在目標區域的邊界上隨機選取$n$個點組成點集,用$P_n^t=\{(x_1^t,y_1^t),(x_2^t,y_2^t),\cdots,(x_n^t,y_n^t)\}$來表示,其中$t$表示迭代的輪數
- 然后將$P_n^t$中的點連接起來圍成邊界區域,用$S_n^t=c(P_n^t)$表示
- 而后選擇$t$次迭代中邊界區域$S_n^t$與標簽分割圖$S_{GT}$具有最大IOU(Intersection Over Union)的點集作為結構邊界關鍵點集,可以用下面公式進行表示:
$$\tilde{P}=P_n^{\tilde{t}}, where\ \tilde{t}=argmax_{t\in\{1,\cdots,T\}}IOU(S_n^t,S_{GT})$$ - 結構邊界關鍵點集$\tilde{P}$表示的是圖中的一個一個位置,為了能夠更加適合處理圖像的卷積網絡訓練,文中采用2D的邊界點圖(boundary point map)來表示這些點。具體地,邊界點圖中是以這些點為中心,將它們周圍半徑為$R$的區域內都置為1,其余位置置為0。(提示:邊界點圖的構造方式類似關鍵點檢測任務中heatmap的構造方法,這種做法的好處在於2D的表示更適合卷積網絡的訓練和預測)
上述的過程寫成算法形式就如Algorithm 1所示:
2.3 邊界保持模塊(Boundary Preserving Block, BPB)
上圖(Figure 3)展示了如何在卷積網絡中嵌入BPB模塊,以達到在預測分割圖的過程中將邊界關鍵點考慮進去。具體地,假設現在是第$i$個卷積模塊,模塊的輸入特征維度為$f^i\in \mathbf{R}^{w^i\times h^i\times c^i}$,BPB的過程為:
- 采用上述算法生成對應該模塊的邊界點圖$\hat{M}^i\in\mathbf{R}^{w^i\times h^i\times 1}$
- 然后將輸入特征$f^i$和$\hat{M}^i$做元素點乘($\hat{M}^i$在第三維度上需要廣播),而后將結果和輸入特征$f^i$進行相加,得到結構邊界信息保持特征$v^i\in \mathbf{R}^{w^i\times h^i\times c^i}$作為輸出特征。
這個過程可寫成如下公式:
$$v^i=f^i\oplus(f^i \otimes \hat{M}^i)$$
Figure 3中需要注意的還有$\hat{M}^i$的生成過程,將並行的多路空洞卷積(the dilated convolution)產生的特征連接起來,再進一步做邊界點圖的預測,用公式表達如下:
$$\hat{M}^i=\sigma\left (d_1^1\{[d_1^1(f^i),d_1^3(f^i),d_2^3(f^i),d_4^3(f^i),d_6^3(f^i)]\}\right )$$
其中$d_r^s(f^i)$表示空洞率(dilation rate)為$r$,卷積核大小為$s\times s$的空洞卷積函數(dilated convolution function)。
2.4 形狀邊界判別器(Shape Boundary-aware Evaluator, SBE)
采用上述的BPB模塊嵌入到卷積網絡中其實已經可以達到結構邊界保持的目的了,但可能存在問題:模型預測的邊界點圖和模型預測的分割區域可能不是很匹配,比如某一些點雖然在邊界點圖預測正確了,但最終預測卻缺失了。為了解決這個問題,文中提出了SBE,SBE是一個二分類器,其輸入是邊界點圖和分割圖,輸出為一致性分數,高分即一致,低分則不一致。我的理解是相當於在預測邊界點圖和預測分割圖之間加了約束,使得模型能夠產生更優質的預測結果。
2.5 損失函數
綜合上述的分析,文中提出的模型會產生3個部分的預測:分割圖預測、各個卷積模塊對應的邊界點圖預測、SBE的預測。因此損失函數由這三部分的損失函數組成,表示為:
$$L_{Total}=L_{Seg}+L_{BA}+\sum_{i=1}^lL_{Map}^i$$
其中$L_{Seg}$為預測的分割圖的損失函數:
$$L_{Seg}=-S_{GT}\cdot log\left (\hat{S}_{Pred}\right )-(1-S_{GT})\cdot log\left (1-\hat{S}_{Pred}\right )$$
$L_{Map}^i$表示預測的邊界點圖的損失函數:
$$L_{Map}^i=-M_{GT}^i\cdot log\hat{M}^i-(1-M_{GT}^i)\cdot log(1-\hat{M}^i)$$
$L_{BA}$為作用在SBE預測結果上的損失函數:
$$L_{BA}=-log\left(D(\hat{S}_{Pred};M_{GT}) \right)$$
其中$D(\cdot)$表示SBE接收分割圖和邊界點圖作為輸入並產生預測分數。
3 實驗結果
這里我只給出論文中的定量和定性分析的表格或者圖片,具體的結果分析可以看原文。
4 參考資料
[1] http://openaccess.thecvf.com/content_CVPR_2020/papers/Lee_Structure_Boundary_Preserving_Segmentation_for_Medical_Image_With_Ambiguous_Boundary_CVPR_2020_paper.pdf
[2] John Canny. A computational approach to edge detection. IEEE Transactions on pattern analysis and machine intelligence, (6):679–698, 1986. 4