GRAB-Net: Graph-Based Boundary-Aware Network for Medical Point Cloud Segmentation

Authors: Yifan Liu, Wuyang Li, Jie Liu, Hui Chen, and Yixuan Yuan, Member, IEEE

Keywords:?Point cloud segmentation, graph-based framework, boundary-aware segmentation.

Abstract

點(diǎn)云分割在許多醫(yī)學(xué)應(yīng)用中都是基礎(chǔ)性的,例如動(dòng)脈瘤夾閉和矯正手術(shù)規(guī)劃。最近的方法主要集中在設(shè)計(jì)強(qiáng)大的局部特征提取器,但普遍忽視了物體邊界周圍的分割,這對臨床實(shí)踐極為有害,并降低了整體分割性能。為了解決這個(gè)問題,提出了一種基于圖的邊界感知網(wǎng)絡(luò)(GRAB-Net),包括三個(gè)范式:基于圖的邊界感知模塊(GBM)、外部邊界上下文分配模塊(OCM)和內(nèi)部邊界特征校正模塊(IFM),用于醫(yī)學(xué)點(diǎn)云分割。為了提高邊界周圍的分割性能,GBM被設(shè)計(jì)用于檢測邊界,并在圖域中交換語義和邊界特征之間的補(bǔ)充信息,其中語義-邊界相關(guān)性在全局建模,并通過圖推理交換信息線索。此外,為了減少降低邊界外分割性能的上下文混淆,提出了OCM來構(gòu)建上下文圖,在這里,不同類別的點(diǎn)被賦予不同的上下文,由幾何標(biāo)志引導(dǎo)。另外,提出了IFM來以對比方式區(qū)分邊界內(nèi)部模糊的特征,其中提出了邊界感知對比策略來促進(jìn)判別式表示學(xué)習(xí)。在兩個(gè)公開數(shù)據(jù)集IntrA和3DTeethSeg上進(jìn)行了大量實(shí)驗(yàn),證明了方法對于現(xiàn)有技術(shù)的優(yōu)越性。

點(diǎn)云分割是醫(yī)學(xué)應(yīng)用領(lǐng)域的一種基礎(chǔ)技術(shù)。例如,對3D掃描的牙科模型進(jìn)行分割,有助于牙醫(yī)模擬牙齒拔除、刪除和重排,從而更好地預(yù)測治療效果。另一個(gè)例子是對3D血管表面進(jìn)行顱內(nèi)分割,為動(dòng)脈瘤夾閉手術(shù)過程提供有益的邊界線索。在臨床實(shí)踐中,一種可行的方式是手動(dòng)對物體進(jìn)行分割,但這種方式費(fèi)時(shí)費(fèi)力,且容易出現(xiàn)觀察者之間的差異。因此,迫切需要準(zhǔn)確可靠的自動(dòng)點(diǎn)云分割方法來獲得定量評估。

近年來,提出了許多點(diǎn)云分割方法,可以根據(jù)設(shè)計(jì)理念分為三類。提取器方法詳盡設(shè)計(jì)了各種局部特征提取器來提取信息豐富的表示。語義增強(qiáng)方法則嘗試引入額外信息來增強(qiáng)原始語義特征。盡管取得了一些性能提升,但它們在物體邊界附近的表現(xiàn)難以令人滿意,這對于臨床實(shí)踐是有害的,因?yàn)樵S多操作(如動(dòng)脈瘤夾閉和牙齒拔除)都是沿著邊界線進(jìn)行的。為解決這一問題,提出了邊界感知方法,采用邊界感知和邊界感知對比策略。前一種策略通常是使用額外的分支預(yù)測邊界掩碼,并將其并入語義特征中;而后一種策略()則旨在以對比的方式區(qū)分邊界附近模糊的特征。

盡管在邊界區(qū)域取得了一些改進(jìn),但現(xiàn)有邊界感知框架仍然存在兩個(gè)挑戰(zhàn)。首先,當(dāng)前方法直接組合語義和邊界特征,并通過局部特征提取對它們的關(guān)系進(jìn)行建模,這忽視了全局語義和形狀線索,導(dǎo)致對分割任務(wù)的約束不足。更具體地說,語義和形狀線索隱藏在各自的特征中,我們希望網(wǎng)絡(luò)能夠全局感知這兩種信息,即整個(gè)語義分布和完整的邊界形狀,從而為網(wǎng)絡(luò)生成適當(dāng)?shù)姆指钐卣魈峁┳銐虻募s束。然而,由于特征提取過程引入的局部性,現(xiàn)有工作只能感知部分隱藏信息。另外,當(dāng)前方法采用的直接組合兩種特征的做法相當(dāng)粗糙,這增加了網(wǎng)絡(luò)捕捉它們相互關(guān)聯(lián)的難度。為解決這些限制,我們使用圖技術(shù)全局細(xì)致地對語義-邊界相關(guān)性進(jìn)行建模。在圖域中,對偶對應(yīng)關(guān)系可以由動(dòng)態(tài)構(gòu)造的圖連接來表示,從而捕捉全局范圍依賴關(guān)系。在此基礎(chǔ)上,執(zhí)行圖卷積以自適應(yīng)交換所需信息,為生成具有更精細(xì)細(xì)節(jié)的特征提供充分約束。 ?

第二個(gè)挑戰(zhàn)在于特征提取過程中邊界區(qū)域存在的上下文混淆問題,進(jìn)一步導(dǎo)致了模糊的表示,即特征模糊性問題。例如,如圖1所示,點(diǎn)A和B在坐標(biāo)空間中很接近,因此它們的上下文(即特征提取的參與者)具有高度相似性。這種上下文混淆會(huì)導(dǎo)致網(wǎng)絡(luò)產(chǎn)生相似的特征,進(jìn)而做出相同的類別預(yù)測,而這兩個(gè)點(diǎn)實(shí)際上應(yīng)該屬于不同的類別。為解決這一問題,對于預(yù)測邊界之外的特征,我們致力于為不同類別的點(diǎn)分配不相似的上下文?？紤]到邊界是不同類別的自然幾何分界線,我們可以將上下文重新表示為圖連接,并切斷穿過預(yù)測邊界的連接,這樣可以確保預(yù)測邊界之外的點(diǎn)只與同一側(cè)(即同一類別)的點(diǎn)相連,從而區(qū)分不同類別的上下文。

Methodology

本文提出了一種用于醫(yī)學(xué)點(diǎn)云分割的GRAB-Net方法,包含三個(gè)關(guān)鍵組件:基于圖的邊界感知模塊(GBM)、外邊界上下文分配模塊(OCM)和內(nèi)邊界特征矯正模塊(IFM)。整體框架如圖2所示。給定輸入點(diǎn)云 $I \in \mathbb{R}^{N \times C}$,其中包含坐標(biāo) $P \in \mathbb{R}^{N \times 3}$ 和其他屬性如法向量,$N$為點(diǎn)數(shù),$C$為特征維度。首先利用共享編碼器的雙分支骨干網(wǎng)絡(luò)提取語義特征$X_s \in \mathbb{R}^{N \times D}$和邊界特征$X_b \in \mathbb{R}^{N \times D}$,其中$D$為特征維度。

圖 2. 所提出框架的圖示，由 (a) 基于圖的邊界感知模塊 (GBM)、(b) 外邊界上下文分配模塊 (OCM) 和 (c) 內(nèi)邊界特征校正模塊組成 （IFM）。 輸入點(diǎn)云被輸入雙分支網(wǎng)絡(luò)以分別提取語義特征和邊界特征。 GBM 模型以及通過對偶圖推理 (DGR) 點(diǎn)和邊界特征之間的相互交換相關(guān)性。 最后，IFM以對比的方式進(jìn)一步減少了邊界內(nèi)的特征模糊性。

GBM模塊將點(diǎn)級特征$X_s$和$X_b$投影到語義節(jié)點(diǎn)$V_s$和邊界節(jié)點(diǎn)$V_b$,在圖域中全局建模$V_s$和$V_b$之間的關(guān)系并交換對偶信息,得到更新后的圖節(jié)點(diǎn)$\hat{V}_s$和$\hat{V}_b$,進(jìn)一步重投影為點(diǎn)級特征$\hat{X}_s$和$\hat{X}_b$。具體包括以下步驟:

1. 相干節(jié)點(diǎn)投影:將骨干網(wǎng)絡(luò)提取的點(diǎn)級特征$X_s, X_b \in \mathbb{R}^{N \times D}$投影到更緊湊的特征$V_s, V_b \in \mathbb{R}^{M \times D} (M<N)$,分別作為語義和邊界圖節(jié)點(diǎn)。

   從原始點(diǎn)云$P \in \mathbb{R}^{N \times 3}$中均勻采樣得到$M$個(gè)點(diǎn)的子集$P_{sub} \in \mathbb{R}^{M \times 3}$(圖3中紅色點(diǎn))。對每個(gè)采樣點(diǎn)$P^{sub}_i$,在$P$中搜索$k_1$個(gè)最近鄰(包括$P^{sub}_i$)構(gòu)成鄰域$\mathcal{N}(P^{sub}_i)$。最后以排列不變方式聚合這些鄰域點(diǎn)的特征:

   $V_i = \gamma(\text{MAX}(h(X_{P_j}))), P_j \in \mathcal{N}_{k_1}(P^{sub}_i) \tag{1}$

   其中$X_{P_j} \in X$為點(diǎn)$P_j$對應(yīng)的特征,$\gamma, h$為多層感知機(jī)(MLP),$\text{MAX}$在點(diǎn)維度進(jìn)行,$V_i \in \mathbb{R}^D$為采樣點(diǎn)$P^{sub}_i$聚合后的圖節(jié)點(diǎn)。聚合每個(gè)采樣點(diǎn)$P^{sub}_i$的鄰域語義和邊界特征$X_s$和$X_b$,可以分別得到語義和邊界圖節(jié)點(diǎn)$V_s$和$V_b$。

2. 對偶圖推理(DGR):給定投影的語義節(jié)點(diǎn)$V_s$和邊界節(jié)點(diǎn)$V_b$,可構(gòu)建對偶圖$\mathcal{G}_d = (\mathcal{V},\mathcal{A})$來促進(jìn)信息交換。

圖 3. 相干節(jié)點(diǎn)投影過程的圖示。

B. 外邊界上下文賦值模塊(OCM)

該模塊旨在解決特征提取過程中邊界附近的上下文混淆問題,這可能會(huì)降低醫(yī)學(xué)點(diǎn)云分割在邊界周圍的性能。提出的OCM通過利用圖技術(shù)為邊界附近的每個(gè)點(diǎn)賦予上下文。

1) 上下文圖構(gòu)建:
- 將每個(gè)點(diǎn)視為圖節(jié)點(diǎn),其上下文由與其他節(jié)點(diǎn)的圖連接表示。
- 構(gòu)建上下文圖$G_c = (V_c, E_c)$,其中$V_c = \mathbf{X}_s$(待細(xì)化的語義特征),而$E_c \in \{0, 1\}^{|V_c| \times |V_c|}$是描述節(jié)點(diǎn)連接的鄰接矩陣。
- 通過切斷預(yù)測邊界點(diǎn)$B$之間的連接(從邊界logits $\mathbf{Y}_{bb}$獲得)來分離不同類別點(diǎn)的上下文。
- 將$B$中的點(diǎn)視為內(nèi)邊界點(diǎn)$P_{in}$,其余點(diǎn)視為外邊界點(diǎn)$P_{out}$。
- 第$i$個(gè)和第$j$個(gè)點(diǎn)之間的連接$E_{c_{i,j}}$定義為:

$$
E_{c_{i,j}} = \begin{cases}
1, & (\text{若 } P_j \in \mathcal{N}_{k_2}(P_i)) \land (d(P_i, P_j) < \min\limits_{P_k \in P_{in}} d(P_i, P_k)), \\
0, & \text{否則},
\end{cases}
$$

其中$\land$是邏輯與運(yùn)算符。這確?？缭竭吔绲倪B接被禁止,而外邊界點(diǎn)$P_{out}$上的上下文混淆點(diǎn)被賦予適當(dāng)?shù)纳舷挛摹?/p>

2) 上下文圖推理:
- 利用構(gòu)建的上下文圖$G_c$,執(zhí)行圖推理以產(chǎn)生具有清晰上下文的特征。
- 對鄰接矩陣$E_c$進(jìn)行自環(huán)和度數(shù)歸一化,得到$\mathbf{A} = \mathbf{D}^{-\frac{1}{2}}(\mathbf{E}_c + \mathbf{I})\mathbf{D}^{-\frac{1}{2}}$,其中$\mathbf{I}$是單位矩陣,而$\mathbf{D}_{i,i} = \sum_j \mathbf{E}_{i,j}$, $\mathbf{D}_{i,j \neq i} = 0$。
- 基于歸一化鄰接矩陣$\mathbf{A}$執(zhí)行圖推理:

$$
\tilde{\mathbf{X}}_s = \text{ReLU}(\mathbf{A}\mathbf{V}_c\mathbf{W}),
$$

其中$\mathbf{W} \in \mathbb{R}^{D \times D}$是可訓(xùn)練參數(shù),$\tilde{\mathbf{X}}_s$是推理后的語義特征。
- 將清晰的特征$\tilde{\mathbf{X}}_s$傳遞到MLP層,得到的logits $\mathbf{Y}_{es}$由交叉熵?fù)p失$\mathcal{L}_c = -\frac{1}{N}\sum_{i=1}^N \mathcal{CE}(\mathbf{Y}_{es_i}, \mathbf{Y}_{sgt_i})$監(jiān)督,其中$\mathbf{Y}_{sgt_i}$是第$i$個(gè)

點(diǎn)的類別標(biāo)注。

總之,OCM采用了一種新穎的上下文圖構(gòu)建策略,根據(jù)預(yù)測的邊界為不同類別的點(diǎn)賦予不同的上下文。然后執(zhí)行圖推理以產(chǎn)生更少模糊的表示。

圖 4. 上下文圖構(gòu)造和推理的圖示。

C. 內(nèi)邊界特征校正模塊(IFM)

OCM可以解決外邊界點(diǎn)的上下文混淆問題,但由于缺乏幾何標(biāo)志,內(nèi)邊界點(diǎn)的特征仍然存在模糊,導(dǎo)致分割預(yù)測不準(zhǔn)確。為了解決這一瓶頸,我們設(shè)計(jì)了IFM,通過精心設(shè)計(jì)的樣本內(nèi)對比和樣本間邊界感知對比來區(qū)分內(nèi)邊界處的混淆特征。

1) 樣本內(nèi)邊界感知對比:
- 對于樣本內(nèi)對比,我們的目標(biāo)是使內(nèi)邊界點(diǎn)處的類別特定嵌入與同類別的外邊界點(diǎn)嵌入相似,但與不同類別的嵌入不同。
- 這樣可以有效減少特征模糊,因?yàn)榛煜奶卣鞅粌?yōu)化朝著特征空間中正確的方向。
- 首先將特征通過兩個(gè)后續(xù)的MLP層投影到特征嵌入E,并在嵌入空間中進(jìn)行對比,遵循[34]。
- 對于屬于內(nèi)邊界點(diǎn)的類別特定的錨嵌入Ei∈{Eu|Pu∈Pin},將屬于k3最近非邊界鄰居且同類別的嵌入E+i={Eu|Ygtu=Ygti,Pu∈NPoutk3(Pi)}視為正嵌入,將屬于k3最近非邊界鄰居但不同類別的嵌入E-i={Eu|Ygtu≠Ygti,Pu∈NPoutk3(Pi)}視為負(fù)嵌入。
- 選擇k3最近鄰居而不是所有非邊界點(diǎn)的原因是,這些幾何相鄰的嵌入是困難樣本,而困難的正/負(fù)樣本已被證明比容易樣本更有利于對比學(xué)習(xí)[34]。
- 為了約束錨嵌入與正/負(fù)嵌入的相似性,樣本內(nèi)邊界感知對比損失Lintra定義為:

$$
L_{intra} = -\frac{1}{|B|}\sum_{E_i}\frac{1}{|K_i^+|}\sum_{E_j\in K_i^+}\log\frac{h_\theta(E_i, E_j)}{\sum_{E_k\in K_i^-}h_\theta(E_i, E_k)},
$$

其中$h_\theta(\cdot)$是相似度函數(shù),我們采用指數(shù)余弦相似度:$h_\theta(p, q) = \exp(\frac{p\cdot q}{|p||q|}\cdot\frac{1}{\tau})$,其中$\tau$是溫度因子。通過優(yōu)化這個(gè)損失,模型可以區(qū)分內(nèi)邊界處的模糊特征。

2) 樣本間邊界感知對比:
- 最新研究表明,大量有效負(fù)樣本對對比表示學(xué)習(xí)至關(guān)重要。而在樣本內(nèi)對比中,有效負(fù)嵌入的數(shù)量受當(dāng)前點(diǎn)云樣本大小的限制。
- 因此,我們提出樣本間邊界感知對比,以納入其他樣本的嵌入。
- 為此,我們在訓(xùn)練過程中維護(hù)一個(gè)外部記憶庫M∈RL×C×D,其中L是訓(xùn)練樣本長

度,C是類別數(shù),D是特征維度。

- 在訓(xùn)練過程中,我們首先預(yù)熱網(wǎng)絡(luò)T個(gè)epoch以生成合理的嵌入,然后初始化記憶庫MT?？紤]到保存所有點(diǎn)級嵌入需要過多內(nèi)存,我們只保存每個(gè)類別的平均嵌入(中值嵌入也是一種可行的健壯替代方案)。
- 然后在epoch t=T+1,...,對于每個(gè)樣本,我們可以類似于等式8計(jì)算樣本間對比損失Linter,但不同的是,用記憶庫Mt-1中保存的嵌入替換正/負(fù)嵌入:E+i={Eu|Ygtu=Ygti,Eu∈Mt-1}, E-i={Eu|Ygtu≠Ygti,Eu∈Mt-1}。
- 計(jì)算Linter后,以動(dòng)量方式更新記憶庫:

$$
M_i^{t,c} = \alpha M_i^{t-1,c} + (1-\alpha)E_i^c
$$

其中$\alpha$是前一個(gè)和當(dāng)前嵌入的平衡權(quán)重,而$E_i^c$是第i個(gè)訓(xùn)練樣本中類別c的平均嵌入。這樣,記憶庫可以更新以保留適當(dāng)?shù)暮蜻x嵌入,與錨嵌入進(jìn)行有效對比,促使混淆的錨特征朝著特征空間中正確的方向優(yōu)化。

D. 優(yōu)化

在訓(xùn)練過程中,我們聯(lián)合優(yōu)化GBM、OCM和IFM中的損失。首先在GBM中,語義的交叉熵?fù)p失Ls和邊界的交叉熵?fù)p失Lb用于指導(dǎo)網(wǎng)絡(luò)生成適當(dāng)?shù)谋硎?。然后在OCM中,交叉熵?fù)p失Lc用于監(jiān)督外邊界預(yù)測特征的細(xì)化。最后,在IFM中,提出樣本內(nèi)損失Lintra和樣本間損失Linter來區(qū)分內(nèi)邊界預(yù)測的模糊特征。

總之,所提GRAB-Net的總體目標(biāo)函數(shù)為:

$$
L_{overall} = \lambda_1(L_s + L_b) + \lambda_2L_c + \lambda_3(L_{intra} + L_{inter})
$$

其中$\lambda_1$、$\lambda_2$和$\lambda_3$是平衡每一項(xiàng)貢獻(xiàn)的權(quán)重因子。

圖5 顱內(nèi)分割結(jié)果的幾個(gè)典型例子。 每行代表 (a) PointNet++ [8]、(b) PointConv [11]、(c) PCT [16]、(d) PAConv [17] (e) PointTransformer [18]、(f) CBL [ 25]，（g）Ours，（h）Ground Truth。 請注意，右下角的標(biāo)量代表 B-IoU 量，用于衡量邊界區(qū)域周圍的分割性能。

圖6.牙齒分割結(jié)果的兩個(gè)典型示例。 每行代表 (a) PCT [16]、(b) PAConv [17]、(c) PointTransformer [18]、(d) CBL [25]、(e) Ours 和 (f) ground Truth 的結(jié)果。

1. 消融實(shí)驗(yàn):
- 單獨(dú)添加GBM、OCM和IFM模塊都能提高基線的性能,表明全局建模語義-邊界關(guān)聯(lián)、糾正混淆的上下文以及區(qū)分模糊特征的優(yōu)勢。
- 結(jié)合任意兩個(gè)或所有三個(gè)模塊的性能優(yōu)于單獨(dú)使用一個(gè)模塊,證實(shí)了它們的兼容性。

2. GBM中的投影節(jié)點(diǎn)數(shù)量:
- 節(jié)點(diǎn)數(shù)過少(如32)會(huì)因投影過程中的信息損失而導(dǎo)致性能不佳。
- 節(jié)點(diǎn)數(shù)增多會(huì)提高性能,但過多(如超過128)會(huì)飽和并下降。
- 64個(gè)節(jié)點(diǎn)在信息保留和澄清之間達(dá)到最佳平衡。

3. OCM中的圖推理層數(shù):
- 2層性能最佳。1層不足以傳播信息,而3層以上可能過擬合。

4. 最近鄰居數(shù)k1, k2, k3:
- 這些影響一致性節(jié)點(diǎn)投影、上下文圖構(gòu)建和邊界感知對比。
- 分析表明性能對k值的選擇很敏感。

5. 標(biāo)簽不確定性的影響:
- 在損失函數(shù)中考慮標(biāo)簽不確定性可提高分割性能。

6. 在IntrA和3DTeethSeg數(shù)據(jù)集上的結(jié)果:
- 所提出的方法在兩個(gè)數(shù)據(jù)集上都優(yōu)于現(xiàn)有的三維分割方法,尤其在邊界區(qū)域。
- 定性結(jié)果直觀地展示了該方法在各種形狀的血管瘤和牙齒分割中的優(yōu)越性。

Reference

[1]?Liu, Y., Li, W., Liu, J., Chen, H., & Yuan, Y. (2023). GRAB-Net: Graph-based boundary-aware network for medical point cloud segmentation.?IEEE Transactions on Medical Imaging.文章來源地址http://www.zghlxwxcb.cn/news/detail-847323.html

到了這里，關(guān)于文獻(xiàn)學(xué)習(xí)-30-GRAB-Net: 用于醫(yī)學(xué)點(diǎn)云分割的圖的邊界感知網(wǎng)絡(luò)的文章就介紹完了。如果您還想了解更多內(nèi)容，請?jiān)谟疑辖撬阉鱐OY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！