關(guān)鍵詞:海洋微藻;圖像分割;UNet網(wǎng)絡(luò);UCTransNet網(wǎng)絡(luò);注意力機(jī)制
0 引言(Introduction)
在赤潮發(fā)生的早期,快速鑒定赤潮藻類對防治藻華有著十分重要的意義。圖像分割算法是將圖像在像素級別進(jìn)行分類的算法。早期的數(shù)字圖像處理分割主要采用閾值分割、區(qū)域分割、聚類等方法[1]。近年來,深度學(xué)習(xí)技術(shù)開始應(yīng)用到圖像分割領(lǐng)域。SHELHAMER等[2]提出了FCN(Fully Convolutional Networks)網(wǎng)絡(luò)模型,實現(xiàn)了圖片端到端的分割。RONNEBERGER等[3]提出了UNet網(wǎng)絡(luò)模型,U形分割模型開始興起。ZHOU等[4]提出了UNet++,使用密集型跳躍連接加強(qiáng)編碼譯碼器之間的聯(lián)系。OKTAY等[5]提出了Attention-UNet,將注意力機(jī)制引入UNet模型之中。
近年來,多頭注意力機(jī)制興起,CHEN等[6]提出了TransUNet網(wǎng)絡(luò)模型,使用Transformer作為譯碼器。VALANARASU 等[7]提出了Medical Transformer,在Transformer的基礎(chǔ)上加入了可學(xué)習(xí)的門控機(jī)制。JI等[8]將多尺度的語義特征整合到一個框架中,提出了MCTrans模型。GAO 等[9]把Transformer模塊同時加入UNet的編碼與譯碼模塊,提出了UTNet網(wǎng)絡(luò)。JI等[10]使用漸進(jìn)歸一化自注意力網(wǎng)絡(luò)PNS-Net來獲取更大的感受野。ZHANG等[11]將多個網(wǎng)絡(luò)融合,提出了MBT-Net網(wǎng)絡(luò)模型。WANG等[12]提出了使用Transformer代替跳躍連接的模型———UCTransNet,同時在譯碼階段加入了通道注意力模塊[13]。但是,UCTransNet過于重視特征圖通道之間的聯(lián)系而忽略了其空間的特征。本文在UCTransNet的基礎(chǔ)上引入空間注意力機(jī)制[14],得到Channel-Space Attention Module(CSAM)模塊,進(jìn)一步提高了分割效果。
1 算法改進(jìn)(Algorithm improvement)
1.1 UCTransNet網(wǎng)絡(luò)模型
UCTransNet網(wǎng)絡(luò)模型整體的編碼譯碼器依然采用UNet模型架構(gòu)。取消了原始的跳躍連接,將各級編碼完成的特征圖Ei 輸入帶有Transformer的通道交通融合模塊(CCT)中,然后輸出到譯碼器中。在譯碼階段采用通道交叉注意力模塊Channel wise Cross-Attention(CCA)將CCT的輸出和上采樣的輸出進(jìn)行融合。
原始UCTransNet模型架構(gòu)如圖1所示,在CCT模塊中各級特征圖首先經(jīng)過層歸一化(LN),使輸入數(shù)據(jù)的分布更規(guī)范、差異性更小。在歸一化結(jié)束后,將各級數(shù)據(jù)Concat作為多頭交叉注意力機(jī)制(MCA)的V 和K 輸入MCA中,同時將層歸一化之后的各級Ei 作為MCA的Q 輸入。
CCT模塊架構(gòu)如圖2所示,與原始多頭注意力機(jī)制不同的是,MCA是沿著通道軸進(jìn)行注意力操作,可以加強(qiáng)通道聯(lián)系,節(jié)省計算量,MCA架構(gòu)如圖3所示。經(jīng)過MCA后,各級特征圖進(jìn)行歸一化和全連接層MLP(Multi-Layer Perceptron),然后輸出到譯碼器上。通過CCT模塊的各級數(shù)據(jù)加強(qiáng)了通道上的語義聯(lián)系,為后續(xù)譯碼奠定了基礎(chǔ)。在譯碼階段,CCT輸出的Oi 與譯碼器上采樣的特征圖Di 一起輸入通道交叉注意力模塊(CCA),經(jīng)過一層全局池化和全連接層與Oi 相乘,形成融合編碼和譯碼數(shù)據(jù)的特征圖Oi*,然后經(jīng)上采樣輸出。
CCA模型架構(gòu)如圖4所示,CCA加強(qiáng)了Di 與Oi 在通道上的語義聯(lián)系,并且提高了較為重要的通道的權(quán)重。譯碼器將經(jīng)過4次CCA模塊與上采樣的特征圖,輸入1×1的卷積層與S型函數(shù)(Sigmoid)中,最終輸出與原始圖片尺寸一致的分割結(jié)果。
1.2CSAM 網(wǎng)絡(luò)架構(gòu)
由于顯微照片下的藻類細(xì)胞過小,為了進(jìn)一步突出特征圖的空間與通道特征,同時消除編碼與譯碼器之間的語義鴻溝。本研究在UCTransNet的基礎(chǔ)上進(jìn)一步改進(jìn)了CCA模塊,提出了通道空間融合注意力機(jī)制(CSAM),CSAM 模型架構(gòu)如圖5所示。
不同于CCA模塊比較重視通道之間的特征融合,CSAM在保留CCA通道能力的同時,加強(qiáng)了特征圖空間上的權(quán)重,同時在空間和通道多維度上融合編碼與譯碼的數(shù)據(jù),使模型的分割結(jié)果更精細(xì)、準(zhǔn)確。CSAM-UCTransNet網(wǎng)絡(luò)結(jié)構(gòu)如圖6所示。
CSAM的輸入為CCT輸出的特征圖Oi 以及譯碼器上采樣得到的特征圖Di。CSAM網(wǎng)絡(luò)的主干是由三個池化層交叉連接形成的。以Di 為例,設(shè)其輸入大小為H 、W 、C,分別代表高度、寬度、通道數(shù)。首先Di 經(jīng)過通道層方向上的自適應(yīng)平均池化壓縮為1×1×C 的特征圖C1,將C1 經(jīng)過1×1的卷積與Sigmoid函數(shù)得到C2。將輸入Di 與C2 相乘得到通道加權(quán)后的特征圖Dc。同時,將輸入Di 在空間上先進(jìn)行全局平均池化得到H ×W ×1的平均空間特征圖Svi,S1 經(jīng)過1×1的卷積與Sigmoid函數(shù)得到Sv2。同時,為了更加全面地突出數(shù)據(jù)的空間特性,再將Di 進(jìn)行最大池化得到H ×W ×1的最大空間特征圖Sm1。Sm1 經(jīng)過1×1卷積與Sigmoid函數(shù)得到Sm2。將Sv2 與Sm2 經(jīng)cat操作,融合為一個H ×W ×2的特征圖Svm 。將輸入Di 與Svm 相乘得到空間加權(quán)后的特征圖Dvm 。將Dvm與Dc 相加得到融合了空間與通道多重特征的Dcs。將Dcs 與CCT輸出的經(jīng)過相同操作的Ocs 特征圖結(jié)合,最后經(jīng)過1×1卷積與Sigmoid函數(shù)得到最終輸出的結(jié)果O * 。
2 實驗(Experiment)
2.1 數(shù)據(jù)集與預(yù)處理
本文選用LeicaDM4B低倍顯微鏡,收集小球藻培養(yǎng)液作為實驗樣本。共計拍攝全尺寸小球藻顯微圖片48張,同時使用labeme軟件對圖片進(jìn)行標(biāo)記,藻類全尺寸圖片及標(biāo)簽如圖7所示,其中圖7(a)是采集的小球藻原圖,圖7(b)是小球藻標(biāo)簽。
圖7中有21張圖片的尺寸為1 920×1 200,27張圖片的尺寸為2 448×1 920。圖像尺寸過大,直接進(jìn)行訓(xùn)練會導(dǎo)致設(shè)備內(nèi)存溢出。但是,若將圖片的尺寸直接縮放為256×256,則圖片中的大量信息會損失。所以,將圖像縮放并切割為792張尺寸為480×480的圖片,切割完畢之后將label中像素值小于800的圖片刪除,減少圖片中藻類過少帶來的誤差,剩余401張圖片。每張480×480的圖片含有10~100個小球藻,切分之后的圖片及標(biāo)簽如圖8所示。
本文對數(shù)據(jù)集做了以下預(yù)處理。
(1)將數(shù)據(jù)集以3∶1∶1的比例隨機(jī)劃分為訓(xùn)練集、驗證集和測試集。
(2)在訓(xùn)練之前,對圖像進(jìn)行數(shù)據(jù)增強(qiáng),即對圖片進(jìn)行旋轉(zhuǎn)、翻轉(zhuǎn)、裁剪等操作。
2.2 評價指標(biāo)
本文使用圖像分割領(lǐng)域常用的評價指標(biāo)Dice相似系數(shù)(Dice Similarity Coefficient, DSC)和交并比(Intersection overUnion, IoU)。
Dice相似系數(shù)是用于度量兩個集合相似度的函數(shù),可以用于計算預(yù)測圖像和標(biāo)簽圖像之間的相似度,其取值范圍為[0,1],DSC的值越接近于1,表示分割效果越好,計算公式如下:
2.3 實驗設(shè)置
使用編程語言Python和Pytorch框架進(jìn)行基本的程序運(yùn)行和模型搭建。使用Todesk遠(yuǎn)程連接NVIDIA GeForce RTX3060,其內(nèi)存為12 GB。使用Adam作為優(yōu)化器進(jìn)行迭代。由于顯卡內(nèi)存有限,因此訓(xùn)練時設(shè)置batch-size 為2,并且每次訓(xùn)練結(jié)束后將模型在驗證集上進(jìn)行驗證,記錄下驗證后的Dice相似系數(shù)。若這一次的Dice相似系數(shù)高于以前歷次驗證結(jié)果,則保存當(dāng)前模型。迭代200次,保存200次中的最佳模型。
2.4 實驗結(jié)果分析
本文使用UNet、UNet++、Attention-UNet、UCTransNet、CSAM-UCTransNet模型在同一服務(wù)器上設(shè)置相同的學(xué)習(xí)率和batch-size 進(jìn)行分割訓(xùn)練,UNet是醫(yī)學(xué)圖像分割中最常用的網(wǎng)絡(luò)模型之一,UNet++在UNet的基礎(chǔ)上使用了密集跳躍連接;Attention-UNet在UNet的譯碼模塊中加入了注意力門(Attention Gate)機(jī)制。UCTransNet使用多頭注意力模塊代替UNet跳躍連接。各個模型的分割評價指標(biāo)結(jié)果如表1所示;各個模型的分割效果對比圖如圖9所示。
目前,CSAM-UCTransNet網(wǎng)絡(luò)的Dice相似系數(shù)最高,超過了70%。這是因為CSAM-UCTransNet在保留了UCTransNet完整的CCT模塊的基礎(chǔ)上,在譯碼階段增加了CSAM 機(jī)制對數(shù)據(jù)的空間和通道進(jìn)行多尺度融合,彌補(bǔ)了編碼譯碼階段的語義鴻溝,同時提高了模型對圖片細(xì)節(jié)的感應(yīng)能力。在5組對比模型中,UNet的效果最差,Dice相似系數(shù)僅為56.96%,這是因為UNet僅用最基礎(chǔ)的跳躍連接實現(xiàn)編碼譯碼之間的聯(lián)系,這使得模型對于數(shù)據(jù)細(xì)節(jié)把控較差。UNet++是在UNet的基礎(chǔ)上使用密集的跳躍連接,具有更強(qiáng)的特征提取能力,能對不同尺度的特征進(jìn)行融合提取。雖然UNet++融合了各級感受野的特征,但是由于各級感受野只跟相鄰層級進(jìn)行融合,導(dǎo)致在模型譯碼階段容易忽略信息的細(xì)節(jié)及邊緣。Attention-UNet在UNet的譯碼模塊加入注意力門機(jī)制,其計算效率比UNet++更高,并且能更好地把握全局信息。但是,Attention-UNet是直接將圖片中的細(xì)節(jié)信息進(jìn)行加強(qiáng),導(dǎo)致它極易受噪聲干擾,將雜質(zhì)一起分割出來。UCTransNet使用當(dāng)前較為流行的Transformer模塊代替UNet的跳躍連接,并且在譯碼階段加入通道注意力機(jī)制,但是UCTransNet過于重視圖像的通道信息而忽視了圖像的空間信息,相較于CSAM-UCTransNet網(wǎng)絡(luò),UCTransNet對于小目標(biāo)樣本分割效果相對差一點。
圖9為各個模型對驗證集上圖片的分割效果。對比圖為全尺寸顯微照片切割后尺寸為480×480的圖片。圖9(c)呈現(xiàn)了UCTransNet的分割效果,UCTransNet對小目標(biāo)樣本敏感度不足,對小球藻的邊緣分割效果相對不足。圖9(d)呈現(xiàn)了UNet++的分割效果,其對小球藻的邊緣分割效果較好,但是相鄰細(xì)胞易出現(xiàn)粘連情況,影響整體準(zhǔn)確性。圖9(e)呈現(xiàn)了UNet的分割效果,與UNet++的分割效果一樣,UNet模型對單個細(xì)胞的分割區(qū)域過大,導(dǎo)致相鄰細(xì)胞互相粘連。圖9(f)呈現(xiàn)了Attention-UNet的分割效果,雖然Attention-UNet對細(xì)胞邊緣分割更為細(xì)致,但是出現(xiàn)較多的雜質(zhì)。圖9(g)呈現(xiàn)了CSAM-UCTransNet的分割效果,其對小球藻實現(xiàn)了較為準(zhǔn)確的分割,Dice相似系數(shù)達(dá)到了73.69%,同時IoU 達(dá)到了59.16%,結(jié)果最好。
3 結(jié)論(Conclusion
本文在UCTransNet的基礎(chǔ)上,改進(jìn)譯碼器的注意力機(jī)制,提出了一種結(jié)合了通道空間融合注意力機(jī)制的網(wǎng)絡(luò)模型CSAM-UCTransNet。該模型在譯碼階段不僅考慮了數(shù)據(jù)的通道信息,還將CCT模塊輸出的數(shù)據(jù)與譯碼器的空間信息和通道信息進(jìn)行特征融合,進(jìn)一步加強(qiáng)了模型在編碼和譯碼之間的聯(lián)系,并且強(qiáng)化了模型對樣本在空間上的特征提取能力。相較于其他U形圖像分割模型,CSAM-UCTransNet的分割精度更高,得到的效果圖更好。后續(xù)將進(jìn)一步對模型進(jìn)行改進(jìn),同時引入LoRA[15]等參數(shù)高效性微調(diào)方法,進(jìn)一步提高網(wǎng)絡(luò)模型的精度和收斂速度。