基于改進(jìn)UCTransNet的海洋微藻圖像分割模型

關(guān)鍵詞：海洋微藻;圖像分割;UNet網(wǎng)絡(luò);UCTransNet網(wǎng)絡(luò);注意力機(jī)制

0 引言（Introduction）

在赤潮發(fā)生的早期，快速鑒定赤潮藻類對(duì)防治藻華有著十分重要的意義。圖像分割算法是將圖像在像素級(jí)別進(jìn)行分類的算法。早期的數(shù)字圖像處理分割主要采用閾值分割、區(qū)域分割、聚類等方法[1]。近年來，深度學(xué)習(xí)技術(shù)開始應(yīng)用到圖像分割領(lǐng)域。SHELHAMER等[2]提出了FCN（Fully Convolutional Networks）網(wǎng)絡(luò)模型，實(shí)現(xiàn)了圖片端到端的分割。RONNEBERGER等[3]提出了UNet網(wǎng)絡(luò)模型，U形分割模型開始興起。ZHOU等[4]提出了UNet++，使用密集型跳躍連接加強(qiáng)編碼譯碼器之間的聯(lián)系。OKTAY等[5]提出了Attention-UNet，將注意力機(jī)制引入U(xiǎn)Net模型之中。

近年來，多頭注意力機(jī)制興起，CHEN等[6]提出了TransUNet網(wǎng)絡(luò)模型，使用Transformer作為譯碼器。VALANARASU 等[7]提出了Medical Transformer，在Transformer的基礎(chǔ)上加入了可學(xué)習(xí)的門控機(jī)制。JI等[8]將多尺度的語(yǔ)義特征整合到一個(gè)框架中，提出了MCTrans模型。GAO 等[9]把Transformer模塊同時(shí)加入U(xiǎn)Net的編碼與譯碼模塊，提出了UTNet網(wǎng)絡(luò)。JI等[10]使用漸進(jìn)歸一化自注意力網(wǎng)絡(luò)PNS-Net來獲取更大的感受野。ZHANG等[11]將多個(gè)網(wǎng)絡(luò)融合，提出了MBT-Net網(wǎng)絡(luò)模型。WANG等[12]提出了使用Transformer代替跳躍連接的模型———UCTransNet，同時(shí)在譯碼階段加入了通道注意力模塊[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)。取消了原始的跳躍連接，將各級(jí)編碼完成的特征圖Ei 輸入帶有Transformer的通道交通融合模塊（CCT）中，然后輸出到譯碼器中。在譯碼階段采用通道交叉注意力模塊Channel wise Cross-Attention（CCA）將CCT的輸出和上采樣的輸出進(jìn)行融合。

原始UCTransNet模型架構(gòu)如圖1所示，在CCT模塊中各級(jí)特征圖首先經(jīng)過層歸一化（LN），使輸入數(shù)據(jù)的分布更規(guī)范、差異性更小。在歸一化結(jié)束后，將各級(jí)數(shù)據(jù)Concat作為多頭交叉注意力機(jī)制（MCA）的V 和K 輸入MCA中，同時(shí)將層歸一化之后的各級(jí)Ei 作為MCA的Q 輸入。

CCT模塊架構(gòu)如圖2所示，與原始多頭注意力機(jī)制不同的是，MCA是沿著通道軸進(jìn)行注意力操作，可以加強(qiáng)通道聯(lián)系，節(jié)省計(jì)算量，MCA架構(gòu)如圖3所示。經(jīng)過MCA后，各級(jí)特征圖進(jìn)行歸一化和全連接層MLP（Multi-Layer Perceptron），然后輸出到譯碼器上。通過CCT模塊的各級(jí)數(shù)據(jù)加強(qiáng)了通道上的語(yǔ)義聯(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 在通道上的語(yǔ)義聯(lián)系，并且提高了較為重要的通道的權(quán)重。譯碼器將經(jīng)過4次CCA模塊與上采樣的特征圖，輸入1×1的卷積層與S型函數(shù)（Sigmoid）中，最終輸出與原始圖片尺寸一致的分割結(jié)果。

1.2CSAM 網(wǎng)絡(luò)架構(gòu)

由于顯微照片下的藻類細(xì)胞過小，為了進(jìn)一步突出特征圖的空間與通道特征，同時(shí)消除編碼與譯碼器之間的語(yǔ)義鴻溝。本研究在UCTransNet的基礎(chǔ)上進(jìn)一步改進(jìn)了CCA模塊，提出了通道空間融合注意力機(jī)制（CSAM），CSAM 模型架構(gòu)如圖5所示。

不同于CCA模塊比較重視通道之間的特征融合，CSAM在保留CCA通道能力的同時(shí)，加強(qiáng)了特征圖空間上的權(quán)重，同時(shí)在空間和通道多維度上融合編碼與譯碼的數(shù)據(jù)，使模型的分割結(jié)果更精細(xì)、準(zhǔn)確。CSAM-UCTransNet網(wǎng)絡(luò)結(jié)構(gòu)如圖6所示。

CSAM的輸入為CCT輸出的特征圖Oi 以及譯碼器上采樣得到的特征圖Di。CSAM網(wǎng)絡(luò)的主干是由三個(gè)池化層交叉連接形成的。以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。同時(shí)，將輸入Di 在空間上先進(jìn)行全局平均池化得到H ×W ×1的平均空間特征圖Svi，S1 經(jīng)過1×1的卷積與Sigmoid函數(shù)得到Sv2。同時(shí)，為了更加全面地突出數(shù)據(jù)的空間特性，再將Di 進(jìn)行最大池化得到H ×W ×1的最大空間特征圖Sm1。Sm1 經(jīng)過1×1卷積與Sigmoid函數(shù)得到Sm2。將Sv2 與Sm2 經(jīng)cat操作，融合為一個(gè)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 實(shí)驗(yàn)（Experiment）

2.1 數(shù)據(jù)集與預(yù)處理

本文選用LeicaDM4B低倍顯微鏡，收集小球藻培養(yǎng)液作為實(shí)驗(yàn)樣本。共計(jì)拍攝全尺寸小球藻顯微圖片48張，同時(shí)使用labeme軟件對(duì)圖片進(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)練會(huì)導(dǎo)致設(shè)備內(nèi)存溢出。但是，若將圖片的尺寸直接縮放為256×256，則圖片中的大量信息會(huì)損失。所以，將圖像縮放并切割為792張尺寸為480×480的圖片，切割完畢之后將label中像素值小于800的圖片刪除，減少圖片中藻類過少帶來的誤差，剩余401張圖片。每張480×480的圖片含有10～100個(gè)小球藻，切分之后的圖片及標(biāo)簽如圖8所示。

本文對(duì)數(shù)據(jù)集做了以下預(yù)處理。

（1）將數(shù)據(jù)集以3∶1∶1的比例隨機(jī)劃分為訓(xùn)練集、驗(yàn)證集和測(cè)試集。

（2）在訓(xùn)練之前，對(duì)圖像進(jìn)行數(shù)據(jù)增強(qiáng)，即對(duì)圖片進(jìn)行旋轉(zhuǎn)、翻轉(zhuǎn)、裁剪等操作。

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

本文使用圖像分割領(lǐng)域常用的評(píng)價(jià)指標(biāo)Dice相似系數(shù)（Dice Similarity Coefficient， DSC）和交并比（Intersection overUnion， IoU）。

Dice相似系數(shù)是用于度量?jī)蓚€(gè)集合相似度的函數(shù)，可以用于計(jì)算預(yù)測(cè)圖像和標(biāo)簽圖像之間的相似度，其取值范圍為[0，1]，DSC的值越接近于1，表示分割效果越好，計(jì)算公式如下：

2.3 實(shí)驗(yàn)設(shè)置

使用編程語(yǔ)言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í)設(shè)置batch-size 為2，并且每次訓(xùn)練結(jié)束后將模型在驗(yàn)證集上進(jìn)行驗(yàn)證，記錄下驗(yàn)證后的Dice相似系數(shù)。若這一次的Dice相似系數(shù)高于以前歷次驗(yàn)證結(jié)果，則保存當(dāng)前模型。迭代200次，保存200次中的最佳模型。

2.4 實(shí)驗(yàn)結(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跳躍連接。各個(gè)模型的分割評(píng)價(jià)指標(biāo)結(jié)果如表1所示;各個(gè)模型的分割效果對(duì)比圖如圖9所示。

目前，CSAM-UCTransNet網(wǎng)絡(luò)的Dice相似系數(shù)最高，超過了70%。這是因?yàn)镃SAM-UCTransNet在保留了UCTransNet完整的CCT模塊的基礎(chǔ)上，在譯碼階段增加了CSAM 機(jī)制對(duì)數(shù)據(jù)的空間和通道進(jìn)行多尺度融合，彌補(bǔ)了編碼譯碼階段的語(yǔ)義鴻溝，同時(shí)提高了模型對(duì)圖片細(xì)節(jié)的感應(yīng)能力。在5組對(duì)比模型中，UNet的效果最差，Dice相似系數(shù)僅為56.96%，這是因?yàn)閁Net僅用最基礎(chǔ)的跳躍連接實(shí)現(xiàn)編碼譯碼之間的聯(lián)系，這使得模型對(duì)于數(shù)據(jù)細(xì)節(jié)把控較差。UNet++是在UNet的基礎(chǔ)上使用密集的跳躍連接，具有更強(qiáng)的特征提取能力，能對(duì)不同尺度的特征進(jìn)行融合提取。雖然UNet++融合了各級(jí)感受野的特征，但是由于各級(jí)感受野只跟相鄰層級(jí)進(jìn)行融合，導(dǎo)致在模型譯碼階段容易忽略信息的細(xì)節(jié)及邊緣。Attention-UNet在UNet的譯碼模塊加入注意力門機(jī)制，其計(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對(duì)于小目標(biāo)樣本分割效果相對(duì)差一點(diǎn)。

圖9為各個(gè)模型對(duì)驗(yàn)證集上圖片的分割效果。對(duì)比圖為全尺寸顯微照片切割后尺寸為480×480的圖片。圖9（c）呈現(xiàn)了UCTransNet的分割效果，UCTransNet對(duì)小目標(biāo)樣本敏感度不足，對(duì)小球藻的邊緣分割效果相對(duì)不足。圖9（d）呈現(xiàn)了UNet++的分割效果，其對(duì)小球藻的邊緣分割效果較好，但是相鄰細(xì)胞易出現(xiàn)粘連情況，影響整體準(zhǔn)確性。圖9（e）呈現(xiàn)了UNet的分割效果，與UNet++的分割效果一樣，UNet模型對(duì)單個(gè)細(xì)胞的分割區(qū)域過大，導(dǎo)致相鄰細(xì)胞互相粘連。圖9（f）呈現(xiàn)了Attention-UNet的分割效果，雖然Attention-UNet對(duì)細(xì)胞邊緣分割更為細(xì)致，但是出現(xiàn)較多的雜質(zhì)。圖9（g）呈現(xiàn)了CSAM-UCTransNet的分割效果，其對(duì)小球藻實(shí)現(xiàn)了較為準(zhǔn)確的分割，Dice相似系數(shù)達(dá)到了73.69%，同時(shí)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)化了模型對(duì)樣本在空間上的特征提取能力。相較于其他U形圖像分割模型，CSAM-UCTransNet的分割精度更高，得到的效果圖更好。后續(xù)將進(jìn)一步對(duì)模型進(jìn)行改進(jìn)，同時(shí)引入LoRA[15]等參數(shù)高效性微調(diào)方法，進(jìn)一步提高網(wǎng)絡(luò)模型的精度和收斂速度。