BBTUNet:基于上下文Transformer的肝臟腫瘤分割算法研究

宋長(zhǎng)明，宋 蒙，肖 露，梁朝陽(yáng)，彩 朔

（中原工學(xué)院理學(xué)院，河南鄭州 451191）

肝癌是世界上最普遍的癌癥之一，也是影響我國(guó)居民健康的重大公共衛(wèi)生問(wèn)題。CT 掃描是肝臟腫瘤診斷普遍采用的方式，從CT 圖像中快速準(zhǔn)確地分割出腫瘤對(duì)肝癌的臨床診斷具有重要意義。

隨著深度學(xué)習(xí)的迅猛發(fā)展，基于UNet[1-4]的網(wǎng)絡(luò)結(jié)構(gòu)被廣泛應(yīng)用于醫(yī)學(xué)圖像分割領(lǐng)域，如UNet++[5-8]、Attention UNet++[9]、KiUNet[10]、MUNet[11]等，在肝臟腫瘤分割任務(wù)中取得了很大的成就，文獻(xiàn)[12]在UNet的基礎(chǔ)上提出一種用于肝臟及肝腫瘤分割的BSUNet，有效提升了分割精度。近些年，Transformer[13-14]憑借其較好的長(zhǎng)依賴性關(guān)系提取性能，在視覺(jué)領(lǐng)域得到了諸多應(yīng)用[15-19]。文獻(xiàn)[20]提出了一個(gè)端到端的分割網(wǎng)絡(luò)RTNet，有效提升了糖尿病視網(wǎng)膜多病灶的分割精度。

該文提出一種新的模型BBTUNet 用于肝臟腫瘤分割。通過(guò)構(gòu)建Transformer 上下文Bridge，有效增強(qiáng)多尺度特征間的上下文依賴性關(guān)系；針對(duì)肝臟腫瘤CT 圖像中存在的對(duì)比度低、腫瘤多尺度和邊界模糊問(wèn)題，引入深度邊界細(xì)化模塊，重新設(shè)計(jì)Transformer 的前饋神經(jīng)網(wǎng)絡(luò)，細(xì)化分割邊界，獲取不同的感受野。實(shí)驗(yàn)證明，該文方法可以有效提升肝臟腫瘤的分割精度。

1 BBTUNet模型

1.1 BBTUNet 模型的整體結(jié)構(gòu)

BBTUNet結(jié)構(gòu)如圖1 所示，首先將輸入圖像送入Encoder 中提取特征，生成多尺度特征圖F1、F2、F3和F4，接著將不同尺度的特征圖均劃分為4×4 的圖像塊(為了保持卷積的連續(xù)性，劃分的圖像塊之間有重疊)，并將所有尺度的圖像塊展平成二維序列。為了保持通道的一致性，需要對(duì)各層特征圖的圖像塊的二維序列進(jìn)行重塑，得到tokenF1、tokenF2、tokenF3和tokenF4。然后通過(guò)Concat 操作將tokenFi拼接得到CToken，作為Transformer 的上下文Bridge 的輸入，經(jīng)過(guò)上下文Bridge 進(jìn)一步提取更豐富的全局上下文關(guān)系和局部的邊界特征，并將這些具有豐富信息的特征圖和Decoder中的上采樣特征進(jìn)行融合。實(shí)驗(yàn)結(jié)果表明，BBTUNet具有顯著優(yōu)勢(shì)，實(shí)現(xiàn)了更精確的分割。

圖1 BBTUNet整體結(jié)構(gòu)

1.2 上下文Bridge

為充分提取多尺度特征層間的上下文關(guān)系，彌補(bǔ)局部注意力，獲得更為精細(xì)的邊界，文中基于改進(jìn)的Transformer 機(jī)制構(gòu)建上下文Bridge，重新設(shè)計(jì)跳躍連接結(jié)構(gòu)，結(jié)構(gòu)如圖2 所示。

圖2 上下文Bridge結(jié)構(gòu)

首先完成圖像嵌入，將Encoder 輸出的多尺度特征 圖Fi∈ΦH×W×C重構(gòu)為二維|j=1,…,K}，其中每 個(gè)patch 的大小 為p×p，共個(gè)patch，得到特征圖展平的二維序列:

接著將序列化的特征圖投影到N維空間，為了保持patch 的位置信息，再把特定的位置嵌入到patch 中，嵌入公式如下：

其 中，R∈Φ(p2·C)×N表示patch 在N維空間的投影，Rpos∈ΦK×N是嵌入的位置信息。將嵌入的patch輸入到Transformer 模塊中，得到：

1.3 BFFN結(jié)構(gòu)

Transformer 具有很強(qiáng)的上下文長(zhǎng)依賴性表達(dá)能力，但對(duì)局部特征的學(xué)習(xí)不足。受文獻(xiàn)[21]啟發(fā)，該文在Transformer 模塊的前饋神經(jīng)網(wǎng)絡(luò)的全連接層之間嵌入Depth wise conv 層來(lái)解決局部信息提取不充分的問(wèn)題。考慮到肝臟腫瘤區(qū)域和周圍正常組織邊界情況復(fù)雜，且連續(xù)的下采樣造成了大量的邊界細(xì)節(jié)信息損耗，如果僅使用Depth wise conv，卷積核比較單一，會(huì)導(dǎo)致肝臟腫瘤邊界模糊，分割結(jié)果不夠精細(xì)。因此，通過(guò)引入帶有不同空洞率的深度邊界細(xì)化模塊改進(jìn)Transformer 模塊的前饋神經(jīng)網(wǎng)絡(luò)，一方面可以擴(kuò)大特征提取的感受野，獲得多尺度的上下文信息；另一方面可以細(xì)化肝臟腫瘤的邊界，提高分割精度。提出的改進(jìn)前饋神經(jīng)網(wǎng)絡(luò)BFFN 結(jié)構(gòu)如圖3 所示。

圖3 BFFN結(jié)構(gòu)

將上下文Bridge 處理后不同尺度的特征圖通過(guò)conv3×3提取各通道的基本特征:

其中，Dconv表示空洞卷積，r表示空洞率，Concat表示通道拼接，輸出:

將Fiout與Decoder 的上采樣特征進(jìn)行融合，獲取更加豐富的邊界特征。

1.4 損失函數(shù)

由于肝臟腫瘤具有復(fù)雜多樣的形態(tài)表現(xiàn)，腫瘤病灶區(qū)域和周圍其他組織之間的面積相差較大，所以模型的損失函數(shù)采用二元交叉熵?fù)p失函數(shù)(Binary Cross Entropy，BCE)。

BCE損失函數(shù)是二元分割任務(wù)中使用最廣泛的損失函數(shù)，通過(guò)計(jì)算各像素所屬類別的概率，將前景和背景中所有像素交叉熵的加權(quán)平均，損失函數(shù)定義為：

其中，yi表示ground truth 圖像中第i個(gè)像素的標(biāo)簽，pi表示預(yù)測(cè)圖像中第i個(gè)像素屬于前景的概率，N表示圖像中的像素點(diǎn)數(shù)目。

2 實(shí)驗(yàn)結(jié)果與分析

2.1 實(shí)驗(yàn)數(shù)據(jù)集

該文以肝臟腫瘤分割挑戰(zhàn)賽的公開(kāi)數(shù)據(jù)集3DIRCADB 醫(yī)學(xué)分割數(shù)據(jù)集為訓(xùn)練數(shù)據(jù)集和測(cè)試數(shù)據(jù)集。3DIRCADB 數(shù)據(jù)集由幾組患者的匿名醫(yī)學(xué)CT 圖像組成，由臨床專家對(duì)各種感興趣的結(jié)構(gòu)進(jìn)行人工分割，包含3DIRCADB-01 和3DIRCADB-02。其中，3DIRCADB-01 數(shù)據(jù)集由來(lái)自不同歐洲醫(yī)院的10名女性和10名男性的CT圖像組成。3DIRCADB-02數(shù)據(jù)集由兩個(gè)匿名患者的三維CT 掃描組成，圖像分辨率為512×512 像素。

2.2 實(shí)驗(yàn)環(huán)境和數(shù)據(jù)預(yù)處理

該文的實(shí)驗(yàn)基于Python 語(yǔ)言的Pytorch 框架對(duì)整體代碼進(jìn)行編碼運(yùn)行，硬件配置包括AMD 3700X處理器，32 GB 內(nèi)存和單張RTX 2070S 顯卡。

由于實(shí)驗(yàn)數(shù)據(jù)集來(lái)源于CT掃描，數(shù)據(jù)會(huì)受到掃描設(shè)備型號(hào)、制造商以及采集序列的影響，為了更好地訓(xùn)練模型，在實(shí)驗(yàn)的開(kāi)始階段對(duì)輸入圖像進(jìn)行Z-Score像素標(biāo)準(zhǔn)化處理，并將圖像的大小設(shè)置為256×256。

2.3 評(píng)價(jià)指標(biāo)

為了定量評(píng)估模型的分割性能，采用相關(guān)性系數(shù)(Dice)、交并比(IOU)、靈敏度(Sensitivity,Sen)、特異性(Specify,Spe)和準(zhǔn)確率(Accuracy,Acc)五個(gè)指標(biāo)作為實(shí)驗(yàn)的評(píng)估標(biāo)準(zhǔn)。

其中，TP、TN、FP、FN 分別表示真陽(yáng)性、真陰性、假陽(yáng)性和假陰性。

2.4 模型的訓(xùn)練

模型的訓(xùn)練階段采用Adam 優(yōu)化器，初始學(xué)習(xí)率為0.001，epoch 和batchsize 分別設(shè)置為100 和16。訓(xùn)練過(guò)程中的Dice 系數(shù)和Loss 隨epoach 的變化曲線如圖4 所示。

圖4 模型損失函數(shù)和Dice系數(shù)曲線圖

從圖4 可以看出，Dice 系數(shù)值隨著epoch 的增加逐漸提高，當(dāng)epoch 大于100 時(shí)穩(wěn)定在0.82 左右，而訓(xùn)練集和驗(yàn)證集上的Loss 值均在快速下降，并在epoch 大于100 時(shí)收斂，逐漸趨于0。

2.5 實(shí)驗(yàn)結(jié)果與分析

2.5.1 模型性能的客觀評(píng)估

為了驗(yàn)證模型的有效性，在3DIRCADB 數(shù)據(jù)集上進(jìn)行實(shí)驗(yàn)，與UNet、UNet++、AttentionUNet 和Inf-Net幾種網(wǎng)絡(luò)對(duì)比，進(jìn)行客觀的性能評(píng)估，結(jié)果如表1所示。

表1 不同方法的分割性能對(duì)比

由表1 可以看出，提出的BBTUNet 網(wǎng)絡(luò)相對(duì)于傳統(tǒng)的UNet 網(wǎng)絡(luò)效果提升顯著，其中Dice 值從71.2%提升到了82.1%，相比于表現(xiàn)較好的Inf-Net網(wǎng)絡(luò)依然有1.8%的提升；且在其他各項(xiàng)評(píng)價(jià)指標(biāo)上的表現(xiàn)同樣優(yōu)于對(duì)比網(wǎng)絡(luò)。由此可見(jiàn)，提出的基于Transformer 的上下文Bridge 比原始跳躍連接的性能更優(yōu)，可以有效提升肝臟腫瘤的分割精度。

2.5.2 模型性能的主觀評(píng)估

為了更加直觀地評(píng)估模型性能，將該文方法與上述四種網(wǎng)絡(luò)的分割結(jié)果進(jìn)行可視化，肝臟腫瘤的分割結(jié)果如圖5 所示。

圖5 不同網(wǎng)絡(luò)對(duì)肝臟腫瘤的分割結(jié)果

由圖5 可以看出，UNet、UNet++、AttentionUNet和Inf-Net 均可大致分割出肝臟腫瘤的病灶區(qū)域，但對(duì)于邊界模糊、小尺寸腫瘤的分割還存在不足，如AttentionUNet 雖然憑借Attention Gate 模塊，在復(fù)雜邊界的腫瘤分割中有較好的性能表現(xiàn)，卻存在不完整分割問(wèn)題，且對(duì)于多尺度腫瘤的分割容易出現(xiàn)嚴(yán)重的漏檢現(xiàn)象（見(jiàn)圖5 中第一、三行）；Inf-Net 通過(guò)反向注意力和邊緣注意力模塊在多尺度和邊界模糊的腫瘤分割中都有更好的表現(xiàn)，但對(duì)于多尺度小目標(biāo)的復(fù)雜分割，會(huì)出現(xiàn)分割不足現(xiàn)象（見(jiàn)圖5 中第三行）。提出的BBTUNet 網(wǎng)絡(luò)的分割結(jié)果更接近于專家的手工標(biāo)注結(jié)果，相較于其他網(wǎng)絡(luò)，可以獲得更清晰的分割邊界，整體性能表現(xiàn)優(yōu)異，且在復(fù)雜的分割情形中也能取得更精確的分割結(jié)果。

3 結(jié)論

針對(duì)目前肝臟腫瘤分割任務(wù)中的難點(diǎn)和現(xiàn)有醫(yī)學(xué)圖像分割方法的不足，該文提出了一種基于Transformer 的分割網(wǎng)絡(luò)BBTUNet。為了提升肝臟腫瘤的分割精度，將原UNet 網(wǎng)絡(luò)的跳躍連接結(jié)構(gòu)中加入基于Transformer 的上下文Bridge，獲取不同尺度特征層的上下文依賴性關(guān)系，彌補(bǔ)傳統(tǒng)CNN 局部特征的不足，并對(duì)傳統(tǒng)Transformer 的前饋神經(jīng)網(wǎng)絡(luò)改進(jìn)，通過(guò)由可分離空洞卷積構(gòu)成的邊界細(xì)化模塊細(xì)化Encoder 所提取的邊緣，進(jìn)一步學(xué)習(xí)邊界特征。實(shí)驗(yàn)結(jié)果表明，在肝臟腫瘤分割任務(wù)中，提出的BBTUNet 模型能夠增強(qiáng)腫瘤的邊界特征，較為準(zhǔn)確地分割出肝臟腫瘤的邊界，針對(duì)不同尺度腫瘤的分割也表現(xiàn)出了一定的優(yōu)勢(shì)，分割精確度達(dá)到82.1%，顯著高于其他分割模型。在臨床醫(yī)學(xué)中，準(zhǔn)確獲取腫瘤的位置、數(shù)量、形狀等信息對(duì)疾病診斷、后續(xù)治療方案的制定有著重要作用，在今后的工作中，可將文中提出的分割模型推廣到醫(yī)學(xué)分割領(lǐng)域的其他分割任務(wù)中，提升計(jì)算機(jī)輔助診療在臨床應(yīng)用中的精度，使之在未來(lái)得到更好的發(fā)展。