The Fully Convolutional Transformer for Medical Image Segmentation

論文：The Fully Convolutional Transformer for Medical Image Segmentation (thecvf.com)

代碼：Thanos-DB/FullyConvolutionalTransformer (github.com)

期刊/會議：WACV 2023

摘要

我們提出了一種新的transformer，能夠分割不同形態(tài)的醫(yī)學(xué)圖像。醫(yī)學(xué)圖像分析的細(xì)粒度特性所帶來的挑戰(zhàn)意味著transformer對其分析的適應(yīng)仍處于初級階段。UNet壓倒性的成功在于它能夠欣賞分割任務(wù)的細(xì)粒度性質(zhì)，這是現(xiàn)有的基于transformer的模型目前不具備的能力。為了解決這個(gè)缺點(diǎn)，我們提出了全卷積transformer(FCT)，它建立在卷積神經(jīng)網(wǎng)絡(luò)學(xué)習(xí)有效圖像表示的能力的基礎(chǔ)上，并將它們與transformer的能力相結(jié)合，有效地捕獲其輸入中的長期依賴關(guān)系。FCT是醫(yī)學(xué)影像文獻(xiàn)中第一個(gè)全卷積Transformer模型。它分兩個(gè)階段處理輸入，首先，它學(xué)習(xí)從輸入圖像中提取長期語義依賴關(guān)系，然后學(xué)習(xí)從特征中捕獲分層的全局屬性。FCT結(jié)構(gòu)緊湊、準(zhǔn)確、健壯。我們的結(jié)果表明，它在不需要任何預(yù)訓(xùn)練的情況下，在不同數(shù)據(jù)形態(tài)的多個(gè)醫(yī)學(xué)圖像分割數(shù)據(jù)集上，大大優(yōu)于所有現(xiàn)有的transformer架構(gòu)。FCT在ACDC數(shù)據(jù)集上比其直接模型高出1.3%，在Synapse數(shù)據(jù)集上高出4.4%，在Spleen數(shù)據(jù)集上高出1.2%，在ISIC 2017數(shù)據(jù)集上高出1.1%，在dice metric上的參數(shù)少了多達(dá)5倍。在ACDC Post-2017MICCAI-Challenge在線測試集上，我們的模型在未見過的MRI測試用例上設(shè)置了新的最先進(jìn)的技術(shù)，優(yōu)于大型集成模型以及參數(shù)更少的nnUNet。

1、簡介

醫(yī)學(xué)圖像分割是計(jì)算機(jī)輔助診斷的關(guān)鍵工具。它有助于檢測和定位圖像中病變的邊界，有助于快速識別腫瘤和癌變區(qū)域的潛在存在。這有可能加快診斷，提高檢測腫瘤的可能性，并使臨床醫(yī)生更有效地利用他們的時(shí)間，對患者的結(jié)果有利[15]。傳統(tǒng)上，現(xiàn)代醫(yī)學(xué)圖像分割算法構(gòu)建為對稱的自頂向下編碼器-解碼器結(jié)構(gòu)，首先將輸入圖像壓縮(編碼)到潛在空間，然后學(xué)習(xí)解碼圖像中感興趣區(qū)域的位置。將中間信號的水平傳播(跳越連接)添加到這個(gè)垂直信息流中，我們就得到了UNet架構(gòu)，這可以說是最近分割算法中最具影響力的飛躍。今天大多數(shù)現(xiàn)代分割系統(tǒng)都包括UNet或其變體。UNet成功的關(guān)鍵在于其全卷積的性質(zhì)。UNet在其結(jié)構(gòu)中不估計(jì)任何非卷積可訓(xùn)練參數(shù)。

基于卷積神經(jīng)網(wǎng)絡(luò)(CNN)的UNet模型在醫(yī)學(xué)圖像分割任務(wù)中的準(zhǔn)確性和性能方面取得了巨大的成功。然而，為了真正幫助臨床醫(yī)生進(jìn)行早期疾病診斷，它們?nèi)匀恍枰~外的改進(jìn)。卷積算子固有的局部特性是CNN的一個(gè)關(guān)鍵問題，因?yàn)樗柚沽怂鼈兝脕碜暂斎雸D像的長范圍語義依賴。人們提出了各種方法來為CNN添加全局上下文，最引人注目的是引入注意力機(jī)制，以及擴(kuò)大卷積核以增加核的感受野。然而，這些方法都有自己的缺點(diǎn)。Transformer在語言學(xué)習(xí)任務(wù)中取得了巨大的成功，因?yàn)樗鼈兡軌蛴行У靥幚矸浅ｉL范圍的序列依賴。這導(dǎo)致它們最近適應(yīng)了各種視覺任務(wù)[7,18,21,22]。最近提出的架構(gòu)，如ViT[7]，已經(jīng)超過了cnn在基準(zhǔn)成像任務(wù)上的性能，而最近對ViT的許多改進(jìn)，如CvT [36]， CCT[10]和Swin Transformer[25]，已經(jīng)表明transformer不需要龐大的數(shù)據(jù)消耗模型，甚至可以處理少量數(shù)據(jù)，從而超過CNN的性能。通常，ViT風(fēng)格的模型首先從圖像中提取離散的非重疊patch(在NLP中稱為token)。然后，他們通過位置編碼將空間定位注入到這些patch中，并將此表示通過標(biāo)準(zhǔn)transformer層來建模數(shù)據(jù)中的長期語義依賴關(guān)系。

考慮到CNN和Transformer的明顯優(yōu)點(diǎn)，我們認(rèn)為醫(yī)學(xué)圖像分割的下一步是一個(gè)完全卷積編碼器-解碼器深度學(xué)習(xí)模型，能夠有效地利用醫(yī)學(xué)圖像中的長期語義依賴。為了實(shí)現(xiàn)這一目標(biāo)，我們提出了第一個(gè)用于醫(yī)學(xué)圖像分割的全卷積Transformer。我們新穎的全卷積Transformer層構(gòu)成了我們模型的主要構(gòu)建塊。它包含兩個(gè)關(guān)鍵組件，一個(gè)卷積注意力模塊和一個(gè)全卷積Wide-Focus模塊(見第3節(jié))。我們將我們的貢獻(xiàn)形式化如下:

1. 我們提出了第一個(gè)用于醫(yī)學(xué)圖像分割的全卷積Transformer，它超越了所有現(xiàn)有的基于卷積和transformer的醫(yī)學(xué)圖像分割架構(gòu)的性能，用于多個(gè)二分類和語義分割數(shù)據(jù)集。
2. 我們提出了一種新型的全卷積transformer層，它使用卷積注意力模塊來學(xué)習(xí)長范圍的語義上下文，然后通過寬焦點(diǎn)模塊使用多分辨率空洞卷積創(chuàng)建分層的局部到全局上下文。
3. 通過廣泛的消融研究，我們展示了我們模型的各種構(gòu)建塊在其對模型性能影響的背景下的影響。

2、相關(guān)工作

早期的CNN和Attention模型：UNet[29]是第一個(gè)用于醫(yī)學(xué)圖像分割的CNN模型。最早將注意力模型引入醫(yī)學(xué)圖像分割的工作之一，是通過將門控函數(shù)應(yīng)用于UNet[26]的編碼器到解碼器的特征傳播。FocusNet[17]等方法采用雙編碼器-解碼器結(jié)構(gòu)，其中注意力門選學(xué)習(xí)將相關(guān)特征從一個(gè)UNet的解碼器傳播到下一個(gè)UNet的編碼器。FocesNet++[19]是在分組卷積的各種過濾器組中集成注意力機(jī)制的第一個(gè)作品之一。還有許多UNet的變體，它們使用不同的殘差塊來增強(qiáng)特征提取[32,28,33,20,16]。UNet++[43]在編碼器和解碼器之間創(chuàng)建了嵌套的分層密集跳過連接路徑，以減少它們之間學(xué)習(xí)特征的語義差距。作為最近最具影響力的UNet變體，nnUNet[14]自動調(diào)整自身來預(yù)處理數(shù)據(jù)，并選擇最適合任務(wù)的最佳網(wǎng)絡(luò)架構(gòu)，而不需要人工干預(yù)。

Transformer模型:最初的Transformer架構(gòu)[31]徹底改變了自然語言處理任務(wù)，并迅速成為視覺理解任務(wù)的模型[7]。Transformer在視覺方面工作得很好，因?yàn)樗鼈兡軌騽?chuàng)建長范圍的視覺環(huán)境，但存在固有的缺點(diǎn)，不能利用CNN等圖像中的空間環(huán)境。最近的工作轉(zhuǎn)向了克服這一缺陷的可能解決方案。CvT [36]， CCT[10]和Swin Transformer[25]都是在transformer中集成足夠的空間環(huán)境的嘗試。在醫(yī)學(xué)圖像分割中，大多數(shù)現(xiàn)有研究著眼于創(chuàng)建用于特征處理的Transformer-CNN混合模型。與Attention UNet[26]類似，UNet Transformer[27]增強(qiáng)了CNN，在跳過連接內(nèi)增加了多頭注意。TransUNet[5]是最早提出的用于醫(yī)學(xué)圖像分割的Transformer-CNN混合模型之一，它使用Transformer編碼器饋送到級聯(lián)卷積解碼器。與TransUNet類似，UNETR[12]和Swin UNETR[11]在編碼器上使用Transformer和卷積解碼器來構(gòu)造分割地圖。Transfuse[40]運(yùn)行雙分支編碼器，一個(gè)帶有卷積層，另一個(gè)帶有transformer層，并將其特征與新穎的BiFusion模塊結(jié)合起來。然而，這個(gè)模型的解碼器是卷積的。

當(dāng)前的工作：最近有一個(gè)轉(zhuǎn)變，從創(chuàng)建混合Transformer-CNN模型，到改進(jìn)transformer塊本身，以處理醫(yī)學(xué)圖像的細(xì)微差別。Swin UNet[3]是第一個(gè)提出用于處理醫(yī)學(xué)圖像的純transformer的架構(gòu)。這里的純指的是僅由transformer層提取和處理的圖像特征，而不需要預(yù)訓(xùn)練的骨干網(wǎng)絡(luò)架構(gòu)。DS-TransUNet[24]引入Transformer Interactive Fusion模塊，以獲得更好的表示全局依賴。這兩個(gè)模型的計(jì)算核心都是Swin Transformer塊。同時(shí)進(jìn)行的工作，如nnFormer[42]和DFormer[37]，試圖利用醫(yī)學(xué)圖像中的本地和全局上下文，通過特別制作的多頭自我關(guān)注塊來滿足這一任務(wù)。這些模型的主要缺點(diǎn)是它們固有的注意力投射和特征處理的線性性質(zhì)，F(xiàn)CT旨在緩解這一點(diǎn)。

現(xiàn)有的醫(yī)學(xué)影像分割模型目前至少存在以下三個(gè)局限性之一。它們要么基于CNN主干網(wǎng)絡(luò)，要么使用卷積層創(chuàng)建，因此限制了它們超越感受野以獲得圖像語義上下文的能力(早期CNN方法)。他們試圖將Transformer集成到他們的特征處理管道中，以利用它們創(chuàng)建長期語義上下文的能力，但反過來，使模型龐大且計(jì)算復(fù)雜(混合Transformer-CNN)。他們試圖通過創(chuàng)建用于分割的純Transformer模型來減少這種計(jì)算負(fù)擔(dān)，而不試圖在低級特征提取階段(并發(fā)工作)對局部空間上下文建模。與現(xiàn)有方法不同，我們的全卷積Transformer沒有這些缺點(diǎn)，同時(shí)仍然是一個(gè)純基于Transformer的醫(yī)療圖像分割架構(gòu)。補(bǔ)充資料中的表4額外總結(jié)了FCT與現(xiàn)有模型相比的主要差異。

3、全卷積Transformer

給定一個(gè)數(shù)據(jù)集 { X , Y } \{\mathbf{X}, \mathbf{Y}\} {X,Y}，其中，$\mathbf{X}$是我們模型的輸入圖像，$\mathbf{Y}$是相應(yīng)的語義或二分類分割映射。對于每個(gè)圖像${\mathbf{x}}_i\in {\mathbb{R}}^{H\times W\times C}$，其中$H$和$W$為圖像的空間分辨率， C = { 3 , … , N } C =\{3,\ldots, N\} C={3,…,N}為輸入通道數(shù)，我們的模型產(chǎn)生一個(gè)輸出分割映射${\mathbf{y}}_i\in {\mathbb{R}}^{H\times W\times K}$，其中， K ∈ { 1 , … , D } K∈\{1,\ldots,D\} K∈{1,…,D}。FCT的輸入是從輸入3D圖像的每個(gè)切片中采樣的2D patch。我們的模型遵循熟悉的UNet形狀，F(xiàn)CT層作為其基本構(gòu)建塊。與現(xiàn)有的方法不同，我們的模型既不是CNN-Transformer的混合，也不是Transformer-UNet的結(jié)構(gòu)，它使用現(xiàn)成的transformer層來編碼或細(xì)化輸入特征。它首先從圖像中提取重疊的patch，然后創(chuàng)建基于patch的掃描嵌入，然后在這些patch上應(yīng)用多頭自注意，從而構(gòu)建特征表示。然后通過我們的Wide-Focus模塊處理給定圖像的輸出投影，以從投影中提取細(xì)粒度信息。圖1顯示了我們的網(wǎng)絡(luò)體系結(jié)構(gòu)的概述。

3.1 FCT層

每個(gè)FCT層都從LayerNormalization-Conv-Conv-Maxpool操作開始。我們從經(jīng)驗(yàn)上注意到，與直接先創(chuàng)建圖像的patch-wise投影相比，在3×3內(nèi)核大小較小的patch上按順序應(yīng)用這些連續(xù)卷積有助于更好地編碼圖像信息。每個(gè)卷積層后面都有一個(gè)Gelu激活函數(shù)。我們的FCT塊與其他模型塊不同的第一個(gè)實(shí)例是通過它對醫(yī)學(xué)成像的卷積注意力應(yīng)用。

MaxPool的輸出被輸入到轉(zhuǎn)換函數(shù)$\mathbf{T}(?)$中，轉(zhuǎn)換函數(shù)$\mathbf{T}(?)$將其轉(zhuǎn)換為新的token映射。我們選擇的$\mathbf{T}(?)$是Depthwise-Convolution operator。我們選擇一個(gè)較小的內(nèi)核大小3×3, 步長為s×s和一個(gè)有效的填充，以確保:(1)與大多數(shù)現(xiàn)有工作不同，提取的patch是重疊的，并且(2)卷積操作不會始終改變輸出大小。接下來是LayerNormalization操作。得到的token映射$p_{i+1}\in {\mathbb{R}}^{W_t\times H_t\times C_t}$被平化為$W_t{H}_t\times C_t$，創(chuàng)建了我們的patch嵌入式輸入。下一個(gè)例子是，我們的FCT層不同于現(xiàn)有的基于transformer的醫(yī)學(xué)成像應(yīng)用方法，是通過它的注意力投影。所有現(xiàn)有模型都采用線性逐點(diǎn)線性映射來進(jìn)行多頭自我注意(MHSA)計(jì)算。這導(dǎo)致Transformer模型失去空間信息，這對成像應(yīng)用非常重要。現(xiàn)有的方法試圖通過卷積增強(qiáng)來緩解這個(gè)問題，使其適應(yīng)成像任務(wù)。然而，這為所提出的模型增加了額外的計(jì)算成本。受[36]中提出的方法的啟發(fā)，我們將MHSA塊中的逐點(diǎn)線性映射替換為Depthwise-Convolution，以降低計(jì)算成本，并從圖像中利用更好的空間上下文信息。patch嵌入和卷積注意力投影構(gòu)成了我們的卷積注意力的組成部分。與[36]不同的是，我們注意到用LayerNormalization替換BatchNormalization有助于提高性能。此外，刪除Point-wise Convolution會導(dǎo)致一個(gè)更簡單的模型，而不會損失任何性能。Depthwise-Convolution提供的空間上下文進(jìn)一步消除了對位置編碼的需求，位置編碼用于在輸入中插入空間信息，并順序跟蹤每個(gè)patch的位置，從而進(jìn)一步簡化了架構(gòu)設(shè)計(jì)。

一般的Transformer層遵循線性層MHSA塊，因此丟失了圖像中的所有空間上下文。直接用卷積替換這些線性層是一種相對簡單的方法，可以緩解這個(gè)問題并提高性能。然而，醫(yī)學(xué)圖像需要細(xì)粒度的信息處理。記住這一點(diǎn)，我們采用了一個(gè)多分支卷積層，其中一層對MHSA輸出應(yīng)用空間卷積，而其他層應(yīng)用空洞卷積，增加感受野，以獲得更好的空間上下文。然后，我們通過求和來融合這些特征，并將它們傳遞到特征聚合層。這種特征聚合是通過另一個(gè)空間卷積算子完成的。我們稱這個(gè)模塊為Wide-Focus。殘差連接用于增強(qiáng)整個(gè)層的特征傳播。最后的特征被重新塑造，并進(jìn)一步傳播到下一個(gè)FCT層。圖1(上)顯示了FCT層。

3.2 編碼器

我們的模型的編碼器包含四個(gè)FCT層，負(fù)責(zé)特征提取和傳播。對于第$l$個(gè)transformer層，卷積注意力模塊的輸出為，${\mathbf{z}}_l^{\prime }=\mathbf{MHSA}(z_{l?1})+{\mathbf{z}}_{l?1}^{q/k/v}$，其中，${\mathbf{z}}_{l?1}^{q/k/v}=\mathbf{Flatten}(\mathbf{DepthConv}(\mathbf{Reshape}(z_{l?1})))$。$\mathbf{MHSA}({\mathbf{z}}_{l?1})=softmax(Q{K}^T/\sqrt{d})V$。然后由Wide-Focus (WF)模塊處理$z_l^{\prime }$, ${\mathbf{z}}_l=\mathbf{WF}({\mathbf{z}}_l)+{\mathbf{z}}_l^{\prime }$。我們進(jìn)一步為編碼器注入金字塔風(fēng)格的圖像輸入，目的是在不同尺度上突出顯示不同類別和更小的ROI特征。值得注意的是，即使沒有這種多尺度圖像金字塔輸入，我們的模型也能夠獲得最先進(jìn)的結(jié)果。數(shù)據(jù)的(瓶頸)潛在編碼是使用另一個(gè)FCT層創(chuàng)建的。

3.3 解碼器

解碼器將瓶頸表示作為其輸入，并學(xué)習(xí)從該信息中重新采樣二分類或語義分割映射。為了在解碼器層中創(chuàng)建更好的上下文相關(guān)性，還使用從編碼器到解碼器的跳過連接，其中來自編碼器層的具有相同分辨率的特征映射與解碼器層連接。解碼器的形狀與編碼器對稱。解碼器中的層對應(yīng)于編碼器中的圖像金字塔層，輸出中間分割映射，提供額外的監(jiān)督并提高模型的預(yù)測能力。上下文相關(guān)性是通過首先對特征量進(jìn)行上采樣，然后將其傳遞到FCT層以了解其最佳可能表示來創(chuàng)建的。我們沒有在FCT的最低規(guī)模上采用深度監(jiān)管，因此我們的模型不是“完全深度監(jiān)管”。這是因?yàn)槲覀冇^察到，輸入圖像掃描中的感興趣區(qū)域(roi)有時(shí)太小，無法在最低尺度(28 × 28)進(jìn)行分割，這導(dǎo)致模型性能較差。這種低規(guī)模的輸出在模型中增加了強(qiáng)烈的偏差，以預(yù)測一些輸出roi作為背景類。

4、實(shí)驗(yàn)

數(shù)據(jù)集：(MRI) Automatic Cardiac Diagnosis Challenge (ACDC) [2], (CT) Synapse Multi-organ Segmentation Challenge1, (CT) Spleen Segmentation Dataset [1] and (Dermoscopy) ISIC 2017 [6] Skin Cancer Segmentation Challenge。

實(shí)驗(yàn)細(xì)節(jié)：模型輸入有兩種尺寸224 x 224，384 x 384。Adam優(yōu)化器，學(xué)習(xí)率1e-3。

5、結(jié)果

6、總結(jié)

我們提出了全卷積transformer，它能夠準(zhǔn)確地執(zhí)行二分類和語義分割任務(wù)，參數(shù)比現(xiàn)有模型更少。FCT在參數(shù)數(shù)量上比nnFormer小5倍以上，比TransUNet和LeViT-UNet小3倍以上。FCT層由兩個(gè)關(guān)鍵組件組成——卷積注意力和Wide-Focus。卷積注意力通過使用深度可分離卷積為模型創(chuàng)建重疊的patch，消除了在patch創(chuàng)建階段對位置編碼的需求。我們基于深度可分離卷積的MHSA塊集成了空間信息，首次在醫(yī)學(xué)成像背景下估計(jì)長距離語義依賴關(guān)系。從我們的消融實(shí)驗(yàn)中可以看出，Wide-Focus有助于利用醫(yī)學(xué)圖像中存在的細(xì)粒度特征信息，并且是提高transformer塊性能的重要因素。我們通過在多個(gè)高度競爭的不同模式和維度的細(xì)分?jǐn)?shù)據(jù)集上的最先進(jìn)的結(jié)果證明了我們模型的能力。我們的FCT塊是第一個(gè)為醫(yī)學(xué)成像應(yīng)用而提出的全卷積transformer塊，并且可以輕松擴(kuò)展到醫(yī)學(xué)成像的其他領(lǐng)域和應(yīng)用。我們相信我們的模型可以作為未來分割任務(wù)的有效骨干網(wǎng)絡(luò)，并為基于transformer的醫(yī)學(xué)圖像處理的創(chuàng)新鋪平道路。文章來源地址http://www.zghlxwxcb.cn/news/detail-691782.html

到了這里，關(guān)于FCT: The Fully Convolutional Transformer for Medical Image Segmentation 論文解讀的文章就介紹完了。如果您還想了解更多內(nèi)容，請?jiān)谟疑辖撬阉鱐OY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！