基于改進UCTransNet的海洋微藻圖像分割模型

關鍵詞：海洋微藻;圖像分割;UNet網絡;UCTransNet網絡;注意力機制

0 引言（Introduction）

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

近年來，多頭注意力機制興起，CHEN等[6]提出了TransUNet網絡模型，使用Transformer作為譯碼器。VALANARASU 等[7]提出了Medical Transformer，在Transformer的基礎上加入了可學習的門控機制。JI等[8]將多尺度的語義特征整合到一個框架中，提出了MCTrans模型。GAO 等[9]把Transformer模塊同時加入UNet的編碼與譯碼模塊，提出了UTNet網絡。JI等[10]使用漸進歸一化自注意力網絡PNS-Net來獲取更大的感受野。ZHANG等[11]將多個網絡融合，提出了MBT-Net網絡模型。WANG等[12]提出了使用Transformer代替跳躍連接的模型———UCTransNet，同時在譯碼階段加入了通道注意力模塊[13]。但是，UCTransNet過于重視特征圖通道之間的聯系而忽略了其空間的特征。本文在UCTransNet的基礎上引入空間注意力機制[14]，得到Channel-Space Attention Module（CSAM）模塊，進一步提高了分割效果。

1 算法改進（Algorithm improvement）

1.1 UCTransNet網絡模型

UCTransNet網絡模型整體的編碼譯碼器依然采用UNet模型架構。取消了原始的跳躍連接，將各級編碼完成的特征圖Ei 輸入帶有Transformer的通道交通融合模塊（CCT）中，然后輸出到譯碼器中。在譯碼階段采用通道交叉注意力模塊Channel wise Cross-Attention（CCA）將CCT的輸出和上采樣的輸出進行融合。

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

CCT模塊架構如圖2所示，與原始多頭注意力機制不同的是，MCA是沿著通道軸進行注意力操作，可以加強通道聯系，節省計算量，MCA架構如圖3所示。經過MCA后，各級特征圖進行歸一化和全連接層MLP（Multi-Layer Perceptron），然后輸出到譯碼器上。通過CCT模塊的各級數據加強了通道上的語義聯系，為后續譯碼奠定了基礎。在譯碼階段，CCT輸出的Oi 與譯碼器上采樣的特征圖Di 一起輸入通道交叉注意力模塊（CCA），經過一層全局池化和全連接層與Oi 相乘，形成融合編碼和譯碼數據的特征圖Oi*，然后經上采樣輸出。

CCA模型架構如圖4所示，CCA加強了Di 與Oi 在通道上的語義聯系，并且提高了較為重要的通道的權重。譯碼器將經過4次CCA模塊與上采樣的特征圖，輸入1×1的卷積層與S型函數（Sigmoid）中，最終輸出與原始圖片尺寸一致的分割結果。

1.2CSAM 網絡架構

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

不同于CCA模塊比較重視通道之間的特征融合，CSAM在保留CCA通道能力的同時，加強了特征圖空間上的權重，同時在空間和通道多維度上融合編碼與譯碼的數據，使模型的分割結果更精細、準確。CSAM-UCTransNet網絡結構如圖6所示。

CSAM的輸入為CCT輸出的特征圖Oi 以及譯碼器上采樣得到的特征圖Di。CSAM網絡的主干是由三個池化層交叉連接形成的。以Di 為例，設其輸入大小為H 、W 、C，分別代表高度、寬度、通道數。首先Di 經過通道層方向上的自適應平均池化壓縮為1×1×C 的特征圖C1，將C1 經過1×1的卷積與Sigmoid函數得到C2。將輸入Di 與C2 相乘得到通道加權后的特征圖Dc。同時，將輸入Di 在空間上先進行全局平均池化得到H ×W ×1的平均空間特征圖Svi，S1 經過1×1的卷積與Sigmoid函數得到Sv2。同時，為了更加全面地突出數據的空間特性，再將Di 進行最大池化得到H ×W ×1的最大空間特征圖Sm1。Sm1 經過1×1卷積與Sigmoid函數得到Sm2。將Sv2 與Sm2 經cat操作，融合為一個H ×W ×2的特征圖Svm 。將輸入Di 與Svm 相乘得到空間加權后的特征圖Dvm 。將Dvm與Dc 相加得到融合了空間與通道多重特征的Dcs。將Dcs 與CCT輸出的經過相同操作的Ocs 特征圖結合，最后經過1×1卷積與Sigmoid函數得到最終輸出的結果O * 。

2 實驗（Experiment）

2.1 數據集與預處理

本文選用LeicaDM4B低倍顯微鏡，收集小球藻培養液作為實驗樣本。共計拍攝全尺寸小球藻顯微圖片48張，同時使用labeme軟件對圖片進行標記，藻類全尺寸圖片及標簽如圖7所示，其中圖7（a）是采集的小球藻原圖，圖7（b）是小球藻標簽。

圖7中有21張圖片的尺寸為1 920×1 200，27張圖片的尺寸為2 448×1 920。圖像尺寸過大，直接進行訓練會導致設備內存溢出。但是，若將圖片的尺寸直接縮放為256×256，則圖片中的大量信息會損失。所以，將圖像縮放并切割為792張尺寸為480×480的圖片，切割完畢之后將label中像素值小于800的圖片刪除，減少圖片中藻類過少帶來的誤差，剩余401張圖片。每張480×480的圖片含有10～100個小球藻，切分之后的圖片及標簽如圖8所示。

本文對數據集做了以下預處理。

（1）將數據集以3∶1∶1的比例隨機劃分為訓練集、驗證集和測試集。

（2）在訓練之前，對圖像進行數據增強，即對圖片進行旋轉、翻轉、裁剪等操作。

2.2 評價指標

本文使用圖像分割領域常用的評價指標Dice相似系數（Dice Similarity Coefficient， DSC）和交并比（Intersection overUnion， IoU）。

Dice相似系數是用于度量兩個集合相似度的函數，可以用于計算預測圖像和標簽圖像之間的相似度，其取值范圍為[0，1]，DSC的值越接近于1，表示分割效果越好，計算公式如下：

2.3 實驗設置

使用編程語言Python和Pytorch框架進行基本的程序運行和模型搭建。使用Todesk遠程連接NVIDIA GeForce RTX3060，其內存為12 GB。使用Adam作為優化器進行迭代。由于顯卡內存有限，因此訓練時設置batch-size 為2，并且每次訓練結束后將模型在驗證集上進行驗證，記錄下驗證后的Dice相似系數。若這一次的Dice相似系數高于以前歷次驗證結果，則保存當前模型。迭代200次，保存200次中的最佳模型。

2.4 實驗結果分析

本文使用UNet、UNet++、Attention-UNet、UCTransNet、CSAM-UCTransNet模型在同一服務器上設置相同的學習率和batch-size 進行分割訓練，UNet是醫學圖像分割中最常用的網絡模型之一，UNet++在UNet的基礎上使用了密集跳躍連接;Attention-UNet在UNet的譯碼模塊中加入了注意力門（Attention Gate）機制。UCTransNet使用多頭注意力模塊代替UNet跳躍連接。各個模型的分割評價指標結果如表1所示;各個模型的分割效果對比圖如圖9所示。

目前，CSAM-UCTransNet網絡的Dice相似系數最高，超過了70%。這是因為CSAM-UCTransNet在保留了UCTransNet完整的CCT模塊的基礎上，在譯碼階段增加了CSAM 機制對數據的空間和通道進行多尺度融合，彌補了編碼譯碼階段的語義鴻溝，同時提高了模型對圖片細節的感應能力。在5組對比模型中，UNet的效果最差，Dice相似系數僅為56.96%，這是因為UNet僅用最基礎的跳躍連接實現編碼譯碼之間的聯系，這使得模型對于數據細節把控較差。UNet++是在UNet的基礎上使用密集的跳躍連接，具有更強的特征提取能力，能對不同尺度的特征進行融合提取。雖然UNet++融合了各級感受野的特征，但是由于各級感受野只跟相鄰層級進行融合，導致在模型譯碼階段容易忽略信息的細節及邊緣。Attention-UNet在UNet的譯碼模塊加入注意力門機制，其計算效率比UNet++更高，并且能更好地把握全局信息。但是，Attention-UNet是直接將圖片中的細節信息進行加強，導致它極易受噪聲干擾，將雜質一起分割出來。UCTransNet使用當前較為流行的Transformer模塊代替UNet的跳躍連接，并且在譯碼階段加入通道注意力機制，但是UCTransNet過于重視圖像的通道信息而忽視了圖像的空間信息，相較于CSAM-UCTransNet網絡，UCTransNet對于小目標樣本分割效果相對差一點。

圖9為各個模型對驗證集上圖片的分割效果。對比圖為全尺寸顯微照片切割后尺寸為480×480的圖片。圖9（c）呈現了UCTransNet的分割效果，UCTransNet對小目標樣本敏感度不足，對小球藻的邊緣分割效果相對不足。圖9（d）呈現了UNet++的分割效果，其對小球藻的邊緣分割效果較好，但是相鄰細胞易出現粘連情況，影響整體準確性。圖9（e）呈現了UNet的分割效果，與UNet++的分割效果一樣，UNet模型對單個細胞的分割區域過大，導致相鄰細胞互相粘連。圖9（f）呈現了Attention-UNet的分割效果，雖然Attention-UNet對細胞邊緣分割更為細致，但是出現較多的雜質。圖9（g）呈現了CSAM-UCTransNet的分割效果，其對小球藻實現了較為準確的分割，Dice相似系數達到了73.69%，同時IoU 達到了59.16%，結果最好。

3 結論（Conclusion

本文在UCTransNet的基礎上，改進譯碼器的注意力機制，提出了一種結合了通道空間融合注意力機制的網絡模型CSAM-UCTransNet。該模型在譯碼階段不僅考慮了數據的通道信息，還將CCT模塊輸出的數據與譯碼器的空間信息和通道信息進行特征融合，進一步加強了模型在編碼和譯碼之間的聯系，并且強化了模型對樣本在空間上的特征提取能力。相較于其他U形圖像分割模型，CSAM-UCTransNet的分割精度更高，得到的效果圖更好。后續將進一步對模型進行改進，同時引入LoRA[15]等參數高效性微調方法，進一步提高網絡模型的精度和收斂速度。