基于U-net模型在煙葉主脈分割中的應用

摘要：煙葉主脈的形狀特征是煙葉分級的參考標準之一，該文在U-net模型的基礎上進行了改進，并將該模型應用于煙葉表面主脈的提取。該模型在所有上采樣層中設置了側輸出，并對這些側輸出進行了約束來提升模型的識別效果。在保證模型性能的同時對U形分割網絡進行了剪枝，減少了模型冗余，提升了模型推理速度。對實驗結果進行對比分析，驗證了該模型在煙葉表面主脈分割任務中的有效性。

關鍵詞：煙葉主脈；語義分割；U-net

doi：10.3969/J.ISSN.1672-7274.2024.05.044

中圖分類號：TP 391.41" " " " " 文獻標志碼：B" " " " " " 文章編碼：1672-7274（2024）05-0-03

Application of U-net Model in Main Vein Segmentation of Tobacco Leaves

XU Kaicheng， ZHANG Wei

（Shanghai Weston Information Technology Co.， Ltd.， Shanghai 200052， China）

Abstract： The shape characteristics of the main veins of tobacco leaves are one of the reference standards for tobacco grading. The article improves on the U-net model and applies it to the extraction of main veins on the surface of tobacco leaves. The model sets side outputs in all upsampling layers and constrains these side outputs to improve the recognition performance of the model. The U-shaped segmentation network was pruned while ensuring model performance， reducing model redundancy and improving model inference speed. The comparative analysis of the experimental results verified the effectiveness of the model in the task of segmenting the main veins on the surface of tobacco leaves.

Keywords： main vein of tobacco leaves; semantic segmentation; U-net

0" "引言

近年來，人工智能技術高速發展，已在工業、農業、醫療等多個領域中取得了不錯的成績。其中，深度學習已成為人工智能技術發展的重要方向。傳統的煙葉分級技術依賴人工經驗，效率相對較低且主觀性強，而深度學習技術能夠為專家提供準確的參考信息，輔助專家更高效地完成煙葉分級工作。

對于不同部位的煙葉，其主脈存在著不同的形狀特征，這些特征包括主脈的粗細、彎曲程度與脈槽開閉程度。本文基于U-net的分割模型對煙葉主脈進行分割，采用深度監督方法對模型上采樣層中不同分辨率特征加以約束，提升了模型整體的分割精度，并通過實驗證明該方法在煙葉主脈分割任務中的表現比傳統U-net更加出色。另外，由于神經網絡對于不同任務所需的模型深度不同，為了減少模型的冗余，提升模型的推理速度，本文測試了模型在不同深度下的性能變化，并根據實驗結果對U型分割網絡進行了剪枝。

1" "相關工作

圖像分割是計算機視覺近幾年來最基本的問題之一。傳統的分割方法可以分為基于閾值分割算法[1]、基于區域的分割算法[2]以及基于邊緣的分割算法[3]三種。隨著深度學習的不斷發展，越來越多用于圖像分割的神經網絡模型被相繼提出。全卷積神經網絡[4]采用端到端的方式實現了圖像分割。在此基礎上U-net[5]對網絡結構進一步完善，成為語義分割模型中最具代表性的模型之一。V-net[6]將U-net模型應用于序列數據中，實現了三維數據的分割。除此之外，一些非基于U-net的卷積神經網絡也被應用于圖像分割任務中，包括生成式對抗網絡[7]、循環神經網絡[8]、transformer[9]模型等。

2" "方法

2.1 預處理

對于完整的煙葉圖像，主脈所占的像素個數與完整圖像的像素個數差距較大，若直接使用原圖尺寸進行訓練，少數量樣本（主脈）很難起到優化模型的作用。為了解決類別不平衡對模型訓練產生的負面影響，在模型訓練前需對訓練數據進行預處理。可根據標注圖像得到主脈在圖像中的位置，隨機選擇主脈上的像素作為小圖的中心，裁剪分辨率為128×128的小圖作為訓練數據中的正樣本，如圖1中實線邊框的小圖所示；同時裁剪等量不包含主脈的小圖作為訓練數據中的負樣本，如圖1中虛線邊框的小圖所示，將裁剪后的圖像作為訓練數據輸入模型中進行訓練。

2.2 模型結構

本文采用的分割模型如圖2所示。該模型主要由下采樣、上采樣以及側輸出3部分組成。

模型前半部分對圖像下采樣，卷積模塊主要由兩個3×3卷積核所組成，每個卷積核后跟批歸一化和激活函數，卷積模塊間通過2×2的最大池化層降低圖像分辨率，輸入圖像經5層卷積模塊后得到小分辨率的圖像特征。模型后半部分對圖像特征進行上采樣，反卷積模塊由一個2×2反卷積核和一個3×3卷積核、歸一化層和激活函數所組成，模型根據提取的圖像特征逐層提升分辨率并最終得到分割結果。分割模型的下采樣與上采樣之間，同尺度的特征以跳躍連接方式連接，下采樣特征圖與上采樣特征圖拼接后再由一個1×1的卷積核完成特征融合。

為了提升模型在反卷積模塊中的表現，本文分別為每個反卷積模塊設置了側輸出，反卷積模塊輸出的特征圖經3×3的卷積后再由插值還原至原圖尺寸，每個側輸出的學習目標與最后的輸出層一致，均為圖像中主脈的二值標注圖，最后由輸出層生成分割結果。深度監督能夠起到輔助模型訓練的作用。由于煙葉表面存在大量褶皺，模型淺層通過跳躍連接傳遞的特征中包含與分割目標無關的干擾特征，增加了學習難度，對反卷積層的特征圖進行約束能夠更好地引導隱藏層篩選特征，進而提升模型最終輸出結果。

2.3 模型剪枝

不同任務所需的模型深度是不同的，裁剪模型冗余參數能夠起到模型加速的作用。為了尋找合適的模型深度，本文針對應用于煙葉主脈分割的U型網絡，進行了不同模型深度的實驗。

主要驗證了以下三種深度的模型。如圖3所示，在保持原有的U型結構基礎上，依次減少卷積模塊和相對應的反卷積模塊，以同樣的訓練數據和模型參數進行訓練，最終根據實驗結果選擇合適的模型深度。

2.4 損失函數

本文選用Focal損失函數和dice損失函數[6]作為訓練模型的約束條件，Focal損失函數的定義如下：

（1）

式中，為類別數，在主脈分割任務中該值為2；為不同類別所對應的權重；為模型對每個像素的預測結果；為調變因子，，在本文中，各類別權重均為1，γ取值為2。

dice損失函數定義如下：

（2）

式中，為主脈標注；為模型的預測結果，分子為標注和預測結果間的交集所對應的像素個數，分母為標注的像素個數和預測的像素個數之和。

上采樣的側輸出中所用的損失函數為dice損失函數：

（3）

模型輸出所用的損失函數為dice損失函數和focal損失函數，定義如下：

（4）

總損失函數定義為：

（5）

式中，n為側輸出個數，該值與模型剪枝后剩余反卷積模塊個數相關。

3" "實驗

3.1 實驗設置和數據

本實驗在Centos 7.7上完成，搭建模型的框架采用Pytorch 1.6.0，cuda版本為10.2，顯卡為RTX TITAN V。

實驗中使用的參數主要如下，初始學習率設置為1×10-4，訓練總迭代次數為100，訓練圖像批尺寸為128，選用的優化器算法為RMSProp，訓練中使用的學習率衰減策略為ReduceLROnPlateau，該策略會在損失函數不發生變化時改變學習率的值。

本文中使用的煙葉數據原圖尺寸為1 152×512，裁剪后的小圖分辨率為128×128，訓練集個數為

1 800，驗證集個數為200。

3.2 實驗結果分析

本節對煙葉主脈分割模型的實驗結果進行展示和分析。選用dice指標作為驗證模型效果的評價指標，其定義與式（2）中dice損失函數一致。

將改進后的U型分割網絡與U-net進行實驗對比，具體的量化結果如表1所示，U-net在煙葉主脈分割任務中，其dice值為78.9%，而改進后的U型網絡，其dice值為81.06%，相比U-net有了2.16%的提升。由此可以證明，相較于U-net，改進后的模型分割的煙葉主脈與人工標注有更高的重合度。模型分割的可視化結果如圖4所示，第一列為原圖的標注圖像，第二列為U-net模型的分割結果，第三列為本文改進后模型的分割結果，觀察后可以發現，改進后的分割模型在主脈提取上連續性更強，且在主脈的邊緣上更為貼合，綜上可以證明，本文改進后的分割模型相較于U-net在煙葉主脈分割任務中表現更加出色。

模型剪枝實驗如下。測試并對比不同深度下的模型的dice值以及單張尺寸為1 152×512圖像的推理時間，如表2所示。未對模型進行剪枝時，模型的Dice值為81.06%，單張平均推理時間為72.56 ms，對模型進行1層的裁剪時，模型的dice值為82.9%，單張平均推理時間為59.96 ms，裁剪2層時，模型的dice值為79.7%，單張平均推理時間為47.11 ms。綜上可以得出，當下采樣和上采樣各裁剪1層時，能夠在不影響性能的同時對模型進行加速。

4" "結束語

本文在U-net模型結構上進行了改進，該模型被應用于煙葉表面的主脈分割任務，模型上采樣中對每個反卷積模塊的特征圖進行約束，并通過實驗證明改進后的U型網絡在主脈分割上有更好的效果。對模型深度進行裁剪，分別測試了不同深度下模型在主脈分割任務中的表現以及單張圖像的推理時間，根據對比實驗結果選擇適用于主脈分割的模型深度。

參考文獻

[1] Otsu N. A threshold selection method from gray-level histograms[J]. Automatica， 1975， 11（285-296）： 23-27.

[2] Adams R， Bischof L. Seeded region growing[J]. IEEE Transactions on pattern analysis and machine intelligence， 1994， 16（6）： 641-647.

[3] Canny J. A computational approach to edge detection[J]. IEEE Transactions on pattern analysis and machine intelligence， 1986 （6）： 679-698.

[4] Long J， Shelhamer E， Darrell T. Fully convolutional networks for semantic segmentation[C]//Proceedings of the IEEE conference on computer vision and pattern recognition. 2015： 3431-3440.

[5] Ronneberger O， Fischer P， Brox T. U-net： Convolutional networks for biomedical image segmentation[C]//Medical image computing and computer-assisted intervention–MICCAI 2015： 18th international conference， Munich， Germany， October 5-9， 2015， proceedings， part III 18. Springer International Publishing， 2015： 234-241.

[6] Milletari F， Navab N， Ahmadi S A. V-net： Fully convolutional neural networks for volumetric medical image segmentation[C]//2016 fourth international conference on 3D vision （3DV）. Ieee， 2016： 565-571.

[7] Luc P， Couprie C， Chintala S， et al. Semantic segmentation using adversarial networks[J]. arXiv preprint arXiv：1611.08408， 2016.

[8] Visin F， Ciccone M， Romero A， et al. Reseg： A recurrent neural network-based model for semantic segmentation[C]//Proceedings of the IEEE conference on computer vision and pattern recognition workshops. 2016： 41-48.

[9] Ranftl R， Bochkovskiy A， Koltun V. Vision transformers for dense prediction[C]//Proceedings of the IEEE/CVF international conference on computer vision. 2021： 12179-12188.