光流法修正的時序圖像語義分割模型*

邱曉夢,王 琳,谷文俊,宋 偉,田浩來,胡 譽

(1.鄭州大學河南省大數據研究院,河南 鄭州 450052;2.鄭州大學計算機與人工智能學院,河南 鄭州 450001; 3.北京唯邁醫療設備有限公司,北京 100000;4.中國科學院高能物理研究所,北京 100049)

1 引言

醫學圖像處理非常重要且應用性極強,其中對醫學圖像中的器官、病灶和感興趣區域進行分割是醫療診斷和手術計劃等任務的重要輔助手段。醫學圖像分割在醫學研究中的作用與價值主要體現在以下幾個方面:(1)提取感興趣的區域,從而忽略其他區域的干擾;(2)用于人體器官或病灶尺寸的測量,有助于醫生診斷或修改病人的治療方案;(3)獲取解剖圖譜信息,為醫學圖像的三維重建和可視化等提供原始數據。然而,受成像設備以及患者體位等因素的影響,醫學圖像不可避免地會出現偽影和噪聲,且在傳輸過程中圖像的質量也會受到不同程度的影響。以上這些給圖像分割及診斷工作造成了一定的困擾和挑戰,故本文對醫學圖像分割方法進行研究,以尋求更佳的分割效果。

傳統圖像分割方法包括閾值法、區域生長法和邊緣檢測法[1]等。其中,閾值法只考慮像素點灰度值本身的特征,不考慮空間特征;區域生長法需要人為選取種子,往往會導致區域內產生空洞,并且這2種方法對噪聲比較敏感。邊緣檢測法不能保證邊緣的連續性和封閉性,并且在高細節區存在大量碎邊緣。為了解決以上問題,研究人員將深度學習方法應用于圖像分割任務中,利用相關網絡的學習功能弱化噪聲對分割的影響,從而改善分割效果。語義分割作為圖像分割的一個分支,它為輸入圖像的每個像素分配一個語義類別,以得到像素化的密集分類[2]。當前語義分割被廣泛應用于地理信息系統、無人駕駛、醫學影像分析和機器人等多個領域,并取得了較好的效果。

Long等[3]在2015年提出了全卷積神經網絡FCN(Fully Convolutional Network),首次將深度學習應用于圖像語義分割領域,成為語義分割的開山之作,但FCN在對各個像素進行分類時沒有充分考慮到像素與像素之間的關系。Ronneberger等[4]提出了U-Net(U-shape Network),其編碼、解碼和跳躍連接結構充分地融合了不同尺度之間的信息,得到了更具魯棒性的分割結果。U-Net++[5]通過不同深度的U-Net有效集來降低未知的網絡深度,它們可以部分共享一個編碼器,通過深度監督同時進行學習。而后TransUNet[6]使用了一個結合卷積神經網絡CNN(Convolutional Neural Network)和Transformer[7]的結構作為編碼器,同時采用一個級聯的上采樣來確保預測的準確率,解決了Transformer低級細節不足的問題。DS-TransUNet(Dual Swin Transformer U-Net)[8]網絡構建了一個雙分支的Transformer結構,緩解了分割時像素級信息丟失的問題。

上述圖像語義分割主要針對單幀圖像,對帶有時序特征的圖像,即視頻流,可以充分利用相鄰幀之間的關系,將上一幀圖像的分割結果作為先驗知識傳入到下一幀圖像中,或者利用光流傳遞時序信息,進行消息傳遞,從而獲取更多的分割信息。已有相關研究中,STFCN(Spatio-Temporal FCN)[9]模型將長短期記憶網絡和FCN相結合,構成了一種端到端的時空卷積神經網絡。Netwarp[10]模型利用光流信息,把經過卷積層的上一幀特征傳遞到當前幀的對應位置上。DFF(Deep Feature Flow)[11]模型指定關鍵幀,計算其他幀到該幀之間的光流信息,減少了計算開銷。DVSNet(Dynamic Video Segmentation Network)[12]使用輕量級決策網絡來評估每個區域的置信度得分,在分割效率和質量之間取得了平衡。TDNet(Temporally Distributed Network)[13]是一種時間分布的視頻語義分割網絡,在每個時間步中,只需執行輕量級計算就可從單個子網絡中提取子特征組,然后用一種新的注意傳播模塊來補償幀間的幾何形變,最終收獲了更快的速度和更短的延遲。CSANet(Cross and Self-Attention Network)[14]是一種將自注意力和交互注意力并行的孿生神經網絡,它不僅能在相鄰幀之間傳播時域空間特征,還能聚合當前幀內的空間語義信息。

然而,醫學圖像的血管比較細,分割部分往往與周圍組織的對比度低,因此,在對醫學圖像分割的過程中往往存在邊緣提取效果不佳,且容易出現血管斷裂的問題。在單幀語義分割中,U-Net以其獨特的網絡結構較好地解決了邊緣提取問題,但并未充分利用時序信息中的先驗知識,故無法進一步提高分割的精確度。同時,現有的視頻語義分割模型雖利用了時序信息,但在邊緣提取上有所欠缺。本文結合了光流、U-Net和能獲取更多信息并進行特征提取的Inception[15]結構,提出了一種既考慮醫學圖像特征又兼顧時序信息的網絡模型。該模型首先采用光流模塊獲取相鄰2幀之間的運動信息,然后使用U-Net網絡和Inception結構分別對當前幀和光流信息進行特征提取,最后使用修正模塊對當前幀和光流信息的特征進行權重分配,實現利用光流對當前幀的修正作用,從而改善圖像分割的效果。在相關數據集上的實驗結果表明,本文模型取得了優于對比模型的預測結果。

2 模型結構

本文提出的模型結構如圖1所示。模型包括數據預處理和模型訓練2部分。數據預處理部分對應光流模塊,用來獲取相鄰2幀之間的運動信息。模型訓練部分包括特征提取模塊和修正模塊,特征提取模塊使用2種方式對當前幀和光流信息進行特征提取,得到初步的分割圖像。修正模塊進行權重分配,使用光流信息對當前幀進行修正,從而使得分割結果通過獲取上一幀的信息,得到更多的分割細節。

Figure 1 Structure of the proposed model圖1 本文所提模型結構

2.1 光流模塊

模型預處理部分使用光流來獲取圖像的運動信息。光流是用來描述場景中的物體運動在連續2幀間產生動態變化的方法,其本質是一個二維向量場[16]。模型使用光流法有3個假設前提:(1)前后幀光照能量保持不變;(2)相鄰幀之間同一像素點運動較小;(3)相鄰像素點運動相似。

設某一點在第1幀中的光照能量表示為f(x,y,t),該點經過時間為dt,運動的距離為(dx,dy),得到式(1):

f(x,y,t)=f(x+dx,y+dy,t+dt)

(1)

其中,x、y為像素坐標,t為時間。

為求解dx和dy,對式(1)右邊進行一階泰勒展開并去除余項,再對兩邊同除dt得到光流方程,如式(2)所示:

fxu+fyv+ft=0

(2)

其中,fx和fy為圖像的梯度,ft為沿時間的梯度,且fx=?f/?x,fy=?f/?y,u=dx/dt,v=dy/dt,(u,v)為光流矢量,即(x,y)的瞬時速度,光流方程就是要求解這2個未知數。

由于存在不適定問題(具體體現在孔徑問題上),僅通過光流方程很難得到準確的光流場,故本文使用TV-L1[17]求解光流信息。假設相鄰2幀圖像為I0和I′,則TV-L1模型的能量函數如式(3)所示:

(3)

其中,Ω為有界區域,λ為權值常數,u(x)為點x的二維光流場,|?u|為圖像的二維梯度。前一項是數據約束項,表示相鄰2幀圖像在同一像素點上的灰度值差,后一項是運動正則化約束,即假設運動是連續的。

在獲得光流信息后,就可以對當前光流使用warp函數來獲得下一幀圖像。假設輸入幀為I1∈R3×H×W和I2∈R3×H×W,H為圖像的高,W為圖像的寬,I1和I2分別為t1和t2時刻的圖像,則這2幀之間的前向光流為F1→2∈R2×H×W,后向光流為F1←2∈R2×H×W。此時向后和向前的warp操作分別如式(4)和式(5)所示:

I1=warp(I2,F1→2)

(4)

I2=warp(I1,F1←2)

(5)

為了更直觀地表現光流,對光流進行可視化。其中最簡單的可視化方式就是用箭頭表征光流,如圖2c表示圖2a和圖2b之間的光流。圖中箭頭的方向和長度分別代表光流矢量的方向和大小。

Figure 2 Optical flow visualization圖2 光流可視化

2.2 特征提取模塊

對當前幀的分割使用的是U-Net模型結構,本文使用的U-Net結構及其參數如圖3所示,包含編碼器、解碼器和跳躍連接結構。在編碼部分進行4次下采樣操作,實現特征提取;解碼部分進行4次上采樣操作,用于恢復圖像的大小并且定位分割信息的位置。每次卷積之后都使用ReLU激活函數,它不僅可以使網絡訓練更快并防止梯度消失,還能降低過擬合的風險。在每次上采樣之后與特征提取部分對應通道數相同尺度的層進行融合。通過這種跳躍連接的結構,使得網絡在每個階段都能學習到模型在特征提取池化過程中丟失的相關特征,從而保證分割效果。

Figure 3 Structure of U-Net圖3 U-Net結構

對光流提取的運動信息,本文采用Inception結構進行特征提取。該結構能在每一層感受野不變的情況下加深網絡深度,使得網絡的精度更高,并且這種結構可以有效地減輕網絡的權重。Inception結構及參數如圖4所示。

Figure 4 Structure of Inception圖4 Inception結構

Inception結構將不同的卷積層通過并聯的方式結合在一起,在深度上拼接不同卷積層的結果矩陣,從而獲得一個更深的矩陣。這種網絡結構不僅可以防止過擬合問題,還可以降維處理尺寸較大的矩陣,聚合多尺度信息。

為了獲取更多的信息,本文使用了2層Inception結構,中間使用2個卷積核為1的卷積層進行連接。這種卷積層不僅能實現通道數的降維,還能實現跨通道的信息交互和整合,此外,它還能大大減少參數量。

2.3 修正模塊

對當前幀和光流信息分別使用U-Net和Inception結構進行分割之后,為實現利用光流修正分割結果的目的,本文對2個分割結果進行權重分配。

在神經網絡中,線性層的輸入一般是二維張量矩陣,但在實際應用中并沒有限制。本文將當前幀和光流信息的分割結果以最后一個維度進行拼接,并按照每個通道展平后輸入到線性層。此時將輸出特征設置為2,則可以得到一個表示權重的矩陣。線性層的計算公式如式(6)所示:

(6)

其中,W1和W1表示權重矩陣,b為常數。

緊接著對線性層的輸出進行Softmax操作,實現歸一化,將線性模型輸出的實數域映射到[0,1],表示概率分布的有效實數空間,從而得到2個部分各自所占的權重。歸一化的核心思想就是將每個轉換后的結果除以轉化結果的總和,即轉化后的結果在總數中所占的百分比即概率值。

最后將2個分割結果矩陣各自乘以自己所占的權重矩陣并進行拼接得到最終的分割結果。修正模塊的結構如圖5所示,其中Prod表示點乘操作。

Figure 5 Structure of correction module圖5 修正模塊結構

時序圖像語義分割模型的流程如算法1所示。

算法1 時序圖像語義分割算法輸入:2幅相鄰圖像I1、I2∈R3×H×W。輸出:語義分割圖像O∈R1×H×W。步驟1 計算2幅圖像之間的運動信息,即光流信息,F1→2=TV-L1(I1,I2),F1→2∈R2×H×W。步驟2 使用U-Net結構對I1進行分割,得到特征圖F1∈R3×H×W。步驟3 使用Inception結構對F1→2進行分割,得到特征圖F2∈R3×H×W。步驟4 使用修正模塊進行權重分配,得到最終結果O=Softmax(Linear(F1,F2))。

2.4 損失函數

損失函數是一種用來衡量錯誤和損失程度的函數,它可以反映出圖像和標簽之間的差異,本文使用的損失函數如式(7)所示:

Losstotal=0.5×Lossdice+0.5×Lossbce

(7)

骰子損失(Dice Loss)是醫學影像分割中最常使用的損失函數之一,用來評估預測圖像和標簽圖像之間相似度的一種度量損失,其計算方法如式(8)所示:

(8)

其中,X表示預測值集合,Y表示真實值集合,|X∩Y|表示2個集合對應元素的點乘。

二分類交叉熵損失(BCE Loss)是常用的二分類損失,用于評估數據在分割過程中對每個像素點進行分類時所產生的損失,可以衡量同一個隨機變量中的2個不同概率分布的差異程度,其計算方法如式(9)所示:

(9)

其中,n表示類別數,xi表示預測值,yi表示真實值。

本文同時使用Dice Loss和BCE Loss這2個損失函數,Dice Loss損失函數聚焦于相似性,可對分割細節進行優化,提高分割精度;BCE Loss則可以使得像素保持平滑的梯度。

3 實驗與結果分析

3.1 數據集

為了更加全面地測試本文模型的性能,選取不同類型的3個代表性數據集進行實驗,在預處理階段將圖像大小設置為512×512,并采用兩兩分組的方式,將每對相鄰幀分為一組。

(1)冠狀動脈造影圖。該數據集為某醫療公司的真實數據,每幅圖像都有對應的標簽,黑色為背景,白色為血管。訓練數據中共有1 200幅圖像,將其按照8∶2的比例劃分為訓練集和驗證集;測試集中包含38幅圖像。

(2)果蠅電鏡圖。該數據集為ISBI(International Symposium on Biomedical Imaging)挑戰賽提供的公開數據集,其數據是30組來自果蠅幼蟲腹側神經索的連續切片投射電鏡的圖像,每組由若干幅圖像組成。每幅圖像都帶有一個對應的標注分割圖,其中白色為細胞,黑色為膜。將該數據集按照8∶2的比例劃分為訓練集和測試集。

(3)健康腹部綜合器官圖。健康腹部綜合器官數據集是一個由腹部造影CT和腹部MR造影圖像組成的公開數據集。本文實驗僅使用其中的CT圖像,數據格式為DICOM。其中訓練集包含2 050個數據,驗證集包含266個數據,測試集包含558個數據。

3個數據集均為時序數據集,其中冠狀動脈造影數據集和健康腹部綜合器官數據集為人體醫學圖像,果蠅電鏡圖為生物醫學圖像。

3.2 實驗環境

本文實驗基于PyTorch框架,硬件設備是顯存為12 GB的英偉達GPU TitanV,編程語言為Python 3.7。在實驗中使用的優化器為Adam,學習率為10-4,批次大小為4,輪次為500。在實驗中使用4塊GPU和數據并行進行模型訓練。

3.3 實驗結果及分析

3.3.1 性能評價指標

評價指標主要是用來評估模型的性能優劣,判斷當前模型是否穩定且獲得的結果是否精確。本文選取了相似性系數Dice(Dice similarity)、像素準確率PA(Pixel Accuracy)、交并比IoU(Intersection over Union)3個評價指標。

在計算這3個指標時,將問題轉化為二分類問題,將要分割的2個部分分別看作正例和反例。其中,Dice系數是一種集合相似度度量函數,用于計算2個樣本的相似度;PA表示被分為正例的實例中實際為正例的比例;IoU表示某個類別預測結果與真實標簽之間交集與并集之間的比值。具體計算公式分別如式(10)～式(12)所示:

(10)

(11)

(12)

其中,TP代表真正例,即模型預測類別為正例,真實類別也為正例;FP代表假正例,即預測類別為正例,真實類別為反例;FN代表假反例,即預測類別為反例,真實類別為正例;TN代表真反例,即預測類別為反例,真實類別也為反例。這3個指標的取值均在[0,1],值越大表示分割效果越好。

3.3.2 冠狀動脈造影圖分割結果

基于該數據集的任務是分割出血管信息。為了驗證本文模型的可靠性,在相同的實驗環境下選取7個經典模型和本文模型進行對比。同時,隨機選取3個樣本,對它們在U-Net、U-Net++、Attention U-Net[18]、RefineNet[19]和本文模型上的預測結果進行展示,預測結果如圖6所示。其中圖6a表示當前幀,圖6b表示標簽,圖6c～圖6f為對比模型的結果,圖6g為本文模型的結果。

從圖6中的標注框可以看出,對比模型存在血管斷裂的情況,相對于對比模型,本文模型出現斷裂的次數更少,圖像的連通性更好;同時,在對比模型中極易產生噪聲,降低了分割的準確率,而本文模型基本未出現噪聲;此外,從第2幅圖的分割結果中可以看出,對比模型獲得的細節信息較少,不如本文模型提取的信息豐富。綜上,在冠狀動脈造影圖上本文模型取得了最佳的分割效果。

Figure 6 Segmentation results on coronary angiogram圖6 冠狀動脈造影圖上分割結果

3.3.3 果蠅電鏡圖分割結果

為了進一步驗證本文模型的有效性,使用相同的實驗環境在果蠅電鏡圖數據集上進行實驗。同樣隨機選取3幅圖像,將它們在U-Net、U-Net++、PSPNet(Pyramid Scence Parsing Network)[20]、SegNet[21]以及本文模型上的預測結果展示在圖7中。

從圖7中的標注框可以看出,U-Net和U-Net++的分割有較大可能出現斷裂,導致細胞的連通性降低;同時,對比模型出現噪聲的可能性更大,尤其是U-Net++;在第3幅圖中,SegNet預測結果中沒有斷裂也未產生噪聲,但無法識別出細胞原本的鋸形特征;而本文模型得到的分割信息在各方面都表現良好,分割的準確率高。

3.3.4 腹部綜合器官圖分割結果

為了驗證本文模型和對比模型在腹部綜合器官分割任務上的準確性和有效性,本文以相同實驗環境在腹部綜合器官數據集上進行分割實驗。同樣隨機選取3幅圖像,將它們在U-Net、U-Net++、Attention U-Net、RefineNet和本文模型上的預測結果展示在圖8中。

從第1幅預測圖中可以看出,U-Net++和RefineNet在凸出的部位上預測并不準確,捕獲的信息不充足。在第2幅圖中,Attention U-Net在預測圖中產生了小黑洞。在第3幅圖中,對比模型的預測都產生了噪聲。可見,本文模型不僅在凸出部分的分割效果良好,并且基本不產生噪聲,分割的準確率高。

Figure 7 Segmentation results on drosophila electron micrograph圖7 果蠅電鏡圖分割結果

Figure 8 Segmentation results on combined healthy abdominal organ segmentation圖8 腹部綜合器官圖分割結果

由3個數據集上的分割結果可知,本文模型均取得了較好的效果,為了更加直觀地表現模型的分割效果,本文計算了各個模型在這3個數據集上的評價指標Dice、PA和IoU,結果如表1所示。從表1可以看出,在冠狀動脈造影圖數據集上,本文模型較次優模型(U-Net)在3個指標上分別提高了0.6%,0.13%和2.13%;在果蠅電鏡圖數據集上分別提高了0.42%,0.67%和0.83%;在腹部綜合器官圖數據集上分別提高了0.88%,0.02%和1.68%。由此可見,本文模型具有較好的有效性和泛化性。同時,還可以看出,U-Net及其變體往往在分割結果上呈現出相對更好的效果,因此選取U-Net作為骨干網絡是非常合適的選擇。

3.3.5 參數分析及消融實驗

模型的重要參數都會對訓練結果產生一定的影響,比如學習率和優化器,同時Inception結構的數量也會對本文模型的結果有影響,故對這3個參數進行對比實驗。如表2所示為它們在腹部綜合器官圖數據集上的實驗結果。

(1)學習率。學習率是最影響性能的超參數之一,可以通過損失函數的梯度調整網絡權重,不同的學習率會對收斂產生影響。本文在訓練中的學習率為10-4,選擇10-2,10-3,10-5和10-6進行對比。從表2可以看出,評價指標都有所下降,說明過大或過小的學習率都會降低模型的預測結果。

(2)優化器。在訓練模型時,可以使用不同的優化器來最小化損失函數。本文選取的對比優化器是SGD(Stochastic Grodient Descent)和Adagrad,其中,SGD又稱隨機梯度下降,每次只根據一個樣本計算梯度,速度快;Adagrad優化器可以自動調節參數的學習率,降低跳過最優點的概率。本文使用的優化器是Adam,它使用梯度的指數加權平均和梯度平方的指數加權平均來動態地調整每個參數的學習率。從表2可以看出,Adam優化器更能改善分割效果。

(3)Inception模塊數量。本文模型在對光流信息進行操作時使用了2個Inception結構,為了驗證該模塊數量的影響,分別選擇1個和3個模塊進行對比實驗,從表2可以看出,較淺的網絡結構或過深的網絡結構都會導致信息有一定程度的損失。

Table 1 Results of comparison experiments表1 對比實驗結果

同時,為了查看不同模塊對模型性能的影響,本文在各個數據集上進行了一系列的消融實驗。實驗主要考慮了2點:(1)Inception結構的影響;(2)修正模塊的影響。結果如表2中消融實驗部分所示。

本文使用Inception結構對光流信息進行特征提取,Inception結構不僅可以減少模型參數還能獲得更多的語義信息,為驗證其效果在消融實驗中去除Inception結構。從表2中可以看出,在沒有Inception結構的情況下,3個評價指標的結果均有所降低。

考慮到2個分割結果對最終的輸出結果有著不同程度的影響,故在本文模型中對2個分割結果采用修正模塊進行權重分配。在消融實驗中,直接對2個分割結果進行拼接,即讓它們以相同的權重對結果產生影響。從表2可以看出,修正模塊可以提升模型的效果。

Table 2 Analysis of important parameters and ablation experiments表2 重要參數分析及消融實驗

4 結束語

本文提出了一種新穎的時序圖像語義分割模型。該模型以U-Net為骨干,充分利用其在醫學圖像分割上的優勢,輔之以光流,實現消息傳遞。通過在相關數據集上的圖像分割以及消融實驗,驗證了本文模型的有效性和泛化性。實驗結果表明,本文模型與對比模型相比獲得了更好的分割結果,進一步提升了在臨床診斷上的參考價值。

為了進一步優化分割的效果,本文仍然需要進一步思考2個問題:(1)如何更好地實現對邊緣的預測,使得分割結果能夠更好地貼合醫療靶區;(2)時序圖開始時可能沒有出現病灶信息,此時各種模型難免會產生一定的噪聲,如何減少這種噪聲的產生。