基于多頭自注意力機制的茶葉采摘點語義分割算法

宋 彥 楊 帥 鄭子秋 寧井銘

(1.安徽農業大學工學院,合肥 230036; 2.安徽省智能農機裝備工程實驗室,合肥 230036;3.茶樹生物學與資源利用國家重點實驗室,合肥 230036)

0 引言

目前我國名優茶采摘仍以人工為主,在春茶加工季節,需要大量的人力資源,成本高。而現有的采摘機械[1-3],嫩芽和老葉不加區分的“一刀切”,原料質量得不到保證[4]。但隨著機器視覺技術的興起,為解決茶葉選擇性采摘問題提供了一種新的途徑,茶芽的識別和可采摘區域定位則是實現選擇性自動化采摘的前提。

早期的研究多基于圖像的閾值分割方法進行茶芽的識別,如文獻[5-7]中采用單一色彩閾值或聯合色彩閾值的方法從圖像中識別茶芽目標。然而,茶園環境具有非結構特性,包括不受控制的光照、茶葉嫩芽和老葉之間的高度相似性等,在這些條件下,閾值分割法可能面臨魯棒性不足的問題。

近年來,卷積神經網絡得益于較為優秀的特征提取能力,對復雜背景的目標定位表現出較好的魯棒性,受光照、背景影響較小。王琨等[8]、呂軍等[9]選取茶葉的部分顏色和形狀特征作為特征集,輸入卷積神經網絡模型,通過訓練最終使得模型能夠識別茶葉嫩芽。孫肖肖等[10]使用改進的YOLO v3識別茶葉,經測試平均準確率達到了84.2%,召回率達到了82%。呂軍等[11]使用結合圖像預處理的改進YOLO v5模型,解決由于茶園的光照變化茶芽難以準確檢測的問題。由前述可知,基于卷積神經網絡的目標檢測技術為茶芽的識別提供了可行的技術途徑。

一旦圖像中的茶芽目標被定位,準確采摘的下一步就是分割可采摘點或可采摘區域。然而,由于采摘點區域不盡相同,形態大多不規則,而茶葉選擇性采摘又需要準確的像素坐標,此時再使用目標檢測不再合適。語義分割是機器視覺中的典型任務,可用于像素級的對象檢測,能夠準確定位茶葉采摘點的區域。該方法已被應用于需精準定位目標像素坐標的相關農業領域,如作物采摘。LI等[12]使用DeepLabv3+模型[13]分割龍眼果串的主果枝,分割結果的像素準確率達到了94.52%。ZHU等[14]為自動準確地識別甜椒采摘區域,使用自建的全分辨率殘差網絡分割彩椒圖像,在測試集的像素準確率達到了97.94%。YU等[15]使用Mask R-CNN模型分割草莓圖像中的可采摘區域,100幅測試圖像的檢測結果顯示,平均檢測準確率為95.78%,召回率為95.41%,分割的平均區域重合度為89.85%。由于語義分割結果能精準獲取采摘區域,極少包含無關信息,在采摘點定位領域,可以給予后續處理器有效的目標信息,簡化下一步需要數據過濾工作。但上述研究的試驗環境,部分目標與背景相似程度較小,如草莓果實與背景的顏色特征差異明顯;部分研究工作在室內環境或實驗室環境開展,可能難以直接應用于復雜的茶園環境。

本文針對茶芽采摘點分割時面臨的目標尺度小、背景復雜等問題,提出RMHSA-NeXt語義分割模型,其特點在于使用殘差多頭自注意力模塊和結合條形池化的ASPP[13]模塊,能夠提高對關鍵特征的關注程度,并降低對背景與目標的干擾,達到準確分割采摘點的目的,以期為自然環境下茶葉機械選擇性采摘提供可靠識別依據。

1 材料和方法

1.1 茶葉采摘點數據來源與數據集構建

茶葉圖像數據采集于六安市金寨縣青山鎮抱兒村,采集時間為2022年4月上旬到10月下旬。圖像采集設備為STEREOLABS公司生產的ZED雙目相機,為了使圖像數據更加貼合采茶機械的工作狀況,采用視頻錄制的模式,視頻分辨率為1 920像素×1 080像素。所有數據采集完成后通過STEREOLABS公司提供的視頻轉換軟件,將視頻數據逐幀轉換為圖像數據,篩選后,得到原始數據集圖像數量為3 192幅。

為了開展茶葉采摘點語義分割模型研究,首先采用目標檢測模型對原始數據集做茶芽目標檢測,經過對目標檢測結果的篩選,最終采摘點數據集圖像數量為26 793幅,本文使用Labelme軟件標注出茶葉可采摘區域,按照9∶1的比例將數據劃分為訓練集和測試集。

分析采摘點數據集后,發現本文的分割目標的特點在于:茶葉采摘點目標較小(普遍尺度約100像素)、光照強度變化范圍大、采摘點的背景多為茶葉,紋理和色彩差距不明顯、采摘點形狀普遍為長條形等,圖1 為茶葉目標檢測與分割的場景分析圖。

圖1 茶葉目標檢測與分割的場景分析

1.2 RMHSA-NeXt語義分割模型

通用的語義分割模型由編碼器和解碼器構成,編碼器進行特征提取,解碼器利用反卷積或者上采樣運算將編碼器輸出的低分辨率特征映射到高分辨率像素空間,得到密集的像素預測分類[16]。為解決茶葉采摘點較小、光照變化范圍大、采摘點與老葉較為相似等問題,本文使用編碼-解碼架構設計提出了一種新的語義分割模型——RMHSA-NeXt,其網絡結構如圖2所示。為了準確分割茶芽采摘點,該模型提出了以下方法:①在模型的編碼階段,采用ConvNeXt作為特征提取單元,其可以有效地篩除無用特征,增強有效特征的表達,減少茶園環境中的光照變化和茶芽采摘點與背景相似度較大的影響。②提出殘差多頭自注意力模塊(Residual multi-head self-attention,RMHSA),其根據茶芽特征相關性為各特征分配權重,將模型注意力集中于茶葉采摘點目標,減少不相關的老葉背景的干擾,進一步強化模型的特征提取能力。③設計結合條形池化的ASPP結構,由于茶芽采摘點多為長條形,故將常用的正方形池化改變為條形池化,便于抑制背景干擾,且由于其內部的多尺度融合機制,有望解決茶芽采摘點小、目標特征信息不足的問題。

圖2 RMHSA-NeXt語義分割模型

1.2.1特征提取網絡

近年來,Swin Transformer[17]越來越多地被應用于計算機視覺領域,并且展現了優異的性能。ConvNeXt吸收Swin Transformer的優點,通過結合Swin Transformer的層結構、倒置瓶頸和深度可分離卷積(Depthwise Convolution)等技巧,進一步提高了模型的特征提取能力,因此,本文的語義分割模型使用ConvNeXt作為初始特征的提取單元。

圖3是ConvNeXt網絡結構圖,其借鑒了ResNet的設計思想,在ResNet傳統結構上,增加了下采樣方法,該網絡主要由ConvNeXt block組成。圖4是ConvNeXt block的結構圖,首先將初始特征與提取后的特征相融合;其次在特征提取過程中使用深度可分離卷積,只在每個通道上進行空間信息的交互,減少了參數,同時激活函數使用GeLU,相較于ReLU更加平滑和連續,也有更高的收斂幾率;而后使用層歸一化(Layer normalization),將這一中間層的神經元參數進行歸一化,對神經網絡中隱藏層的輸入進行歸一化,從而使得網絡更容易訓練[18];最后使用Drop path隨機地將深度學習中的多分支結構刪除,即讓某些神經元失效,添加正則化能力,防止模型過擬合。

圖3 ConvNeXt網絡結構圖

圖4 ConvNeXt block 網絡結構圖

1.2.2殘差多頭自注意力模塊

注意力機制主要根據上下文內容或像素間的相關性快速提取數據或者圖像中的重要特征。多頭自注意力模塊[19](Multi-head self-attention,MHSA)是多個自注意力模塊[20]平行堆疊而來,其可以根據采摘點目標相關性為各特征動態地分配權重,將網絡注意力集中于茶葉采摘點目標,減少茶園中不相關背景的干擾,提高網絡特征提取性能。

本文基于殘差結構和多頭自注意力機制,構建了如圖5所示的 RMHSA 結構,其中多頭自注意力模塊代替了原本殘差結構中的卷積模塊,同時使用兩次的跳連接,讓未被處理過的原始特征與被多頭自注意力模塊處理過的特征充分融合。在圖5中多頭自注意力機制部分是由多個自注意力機制平行計算,最后拼接而來,自注意力機制作用在于減少了對外部信息的依賴,更有利于捕捉茶芽采摘點數據或特征的內部相關性,圖6展示了自注意力機制的結構。

圖5 殘差多頭自注意力模塊結構示意圖

圖6 自注意力機制結構示意圖

模塊的輸出Z為基于查詢向量WQ、鍵向量WK、值向量WV的加權和,計算式為

Z=Softmax(WKWQ)WV+X

(1)

式中X——輸入特征

自注意力模塊的輸入通過線性變換得到WQ、WK、WV。其中WQ的目的是計算當前位置與其他位置之間的相似度;WK用于計算當前位置和與其他位置之間的相似度,在運算過程中,WK通常會與WQ進行矩陣乘法,得到每個位置向量的權值。而WV與WQ、WK計算出的權值矩陣進行矩陣乘法,得到特征的加權平均結果。

1.2.3結合條形池化的ASPP

多尺度結構可以通過不同尺度的特征提取,幫助網絡獲得不同感受野下的圖像信息,以增強模型對小目標的感知能力。因此,多尺度信息捕獲能力對于解決茶葉采摘點圖像有效信息過少具有重要意義,本文選擇了ASPP結構作為模型中的多尺度結構。ASPP通常由1個1×1卷積、3個不同采樣率的3×3擴張卷積和1個空間池化組成。在這種結構中,通過設置不同的采樣率可以得到不同比例的特征圖。

由于初始ASPP中的空間池化的采樣窗口為正方形,而本文中的茶葉采摘區域大部分為長條形,此時正方形窗口不可避免地會包含其他不相關區域的干擾信息。因此,本文采用了條形池化結構,由于長條形采樣窗口采樣時會減少不相關信息的獲取,從而降低無關信息的干擾[21]。改進后的ASPP結構如圖7所示,而條形池化計算過程如圖8所示。

圖7 改進后的ASPP結構示意圖

圖8 條形池化示意圖

在條形池化時,首先將輸入特征圖進行水平和豎直條形池化后變為H×1和1×W兩個特征圖,隨后經過卷積核為3的一維卷積對2個特征圖分別沿著左右和上下進行擴容,擴容后兩個特征圖尺寸相同,擴容后的兩特征圖對應相同位置求和得到H×W的特征圖,之后通過1×1的卷積與sigmoid處理后與原輸入圖對應像素相乘得到了條狀池化輸出結果,在本文中H和W均為5。

1.2.4解碼部分

在編解碼結構中,編碼器的作用是提取特征,而解碼器的作用是將經過編碼器處理后的特征進行采樣和映射,最終實現逐像素分類。在本文模型中,首先會將從改進ASPP中獲得的多尺度特征和從殘差多頭自注意力機制中獲得的特征進行通道拼接,在此過程中,部分多尺度結構會被上采樣以保證特征圖大小一致;而后會經過一次1×1卷積和3×3處理,降低特征圖的維度,篩選特征;最后使用上采樣將圖像恢復到與輸入圖像一樣的大小,再使用Softmax函數進行結果預測,確定每一個特征圖的類別。

2 實驗結果與分析

2.1 實驗平臺

計算機配置CPU型號為Intel Core i7-9700 CPU,GPU型號為NVIDIA GTX2080Ti,內存為32 GB,1 TB固態硬盤。計算機操作系統為Ubuntu 20.04,配置PyTorch深度學習框架用于所有模型的訓練和測試,模型評價等程序均在Python語言環境下編寫。

2.2 模型評價指標

本文語義分割模型采用準確率、檢測速度、參數量等指標進行性能評價。準確率是根據網絡模型預測圖像與人工標注圖像之間的像素誤差計算得到,設語義類別總數為k+1(k個目標類與1個背景類),Pii表示屬于第i類且被預測為第i類的像素數,Pij表示屬于第i類卻被預測為第j類的像素數,在本文中k+1為2,即采摘區域和非采摘區域。

像素準確率(Pixel accuracy,PA)為正確預測像素數量與圖像像素總量的比值,計算公式為

(2)

平均區域重合度(Mean intersection over union,MIoU)為每類預測像素數量與真實像素數量交集與并集比值,然后取所有類別的平均值。平均區域重合度反映了預測結果與圖像真實結果的重合程度,是語義分割模型常采用的準確率度量標準,計算公式為

(3)

檢測速度(每秒檢測幀數,FPS)用來表明每秒能檢測的數據幀數,其數值越高表明訓練好的模型檢測速度越快。

2.3 結果分析

2.3.1消融實驗

為測試本文所提出的語義分割算法的有效性,設計消融實驗分析各功能模塊對模型性能的影響。構建基礎模型,由ConvNeXt網絡和解碼部分組成。在基礎模型上逐步加入殘差多頭自注意力機制、改進的ASPP等結構構成基礎模型+RMHSA、基礎模型+ASPP、基礎模型+改進ASPP以及基礎模型+RMHSA+改進ASPP模型等。通過像素準確率、平均區域重合度、參數量和檢測速度對模型性能進行分析,表1為測試集在上述5種模型的運行結果。

表1 不同模型的分割結果

由表1可知,改進ASPP通過聚合不同尺度的信息增強了模型的判別能力,模型的PA與MIoU相較于基礎模型也增加15.46%和16.54%,且由于采摘點目標形狀均為長條形,使用條形池化減少了不相關信息的獲取,提高了分割指標;對比基礎模型和增加RMHSA模型的結果發現,增加RMHSA模型的PA與MIoU提升到57.14%、51.42%,表明增加RMHSA模型后,由于高效注意力機制的引入,模型更加傾向于尋找更有顯著性的特征,實現強化目標、弱化背景的目的;對比基礎模型+RMHSA+改進ASPP與基礎模型的結果發現,相較于基礎模型,本文模型的PA和 MIoU增加35.74%和37.90%,相較于基礎模型+改進ASPP與基礎模型+RMHSA,本文模型在準確率上有了較大的提升,且在高效注意力機制以及針對性多尺度結構的同時作用下,模型最大程度地保留了有效特征,保證了其向后傳播,減少無用信息的干擾,促使網絡更加關注目標物體不同部位的細節特征,也提升了目標區域的定位精度,有效組合不同尺度下的特征信息,提高分割精度。

在實時性方面,隨著功能模塊的加入,分割模型包含的參數不斷增加,檢測速度逐漸降低。其中,基礎模型的參數數量最小,檢測速度最快,RMHSA模塊的加入使基礎模型參數量增加25.35%,檢測速度降低12.99%;改進ASPP模塊的增加使參數量增加34.9%,檢測速度降低22.80%。基于上述數據定性分析,可以得知改進ASPP模塊對模型的運行效率和計算開銷影響最大,RHMSA模塊次之。

圖9為表1中5種類別模型識別結果圖,針對于分割目標的3個難點:目標較小、光照強度變化范圍大、背景復雜,本文選取7幅圖像展示。由圖9可知,基礎模型由于缺少多尺度結構以及有效的注意力機制,7幅圖像均出現了較大面積的錯檢和漏檢。基礎模型+RMHSA雖然增加了高效的注意力機制,相較于基礎模型有較大的進步,如圖9中的第1行和第5行,但由于網絡結構中的多次下采樣,模型損失較多的有效特征,所以造成結果仍有較多的漏檢和錯檢。而基礎模型+ASPP和基礎模型+改進ASPP結果中,由于增加了多尺度結構,減少了下采樣的特征丟失,但是由于缺少高效的注意力機制的引導,仍然效果不佳,如圖9中的第2、3、6、7行。而本文所提出的模型結合殘差多頭自注意力機制以及改進的ASPP的優點,其識別結果相較于標注圖像沒有明顯差距,沒有大面積的漏檢和錯檢,證明使用這兩種結構組合,可有效加強模型的茶芽采摘點的分割能力。

圖9 不同模型的語義分割結果

2.3.2不同語義分割模型性能對比

為驗證本文所提出的模型具體性能情況,選擇HRNet V2[22]、EfficientUNet++[23]、DeeplabV3+[24]、BiSeNet V2[25]等模型與本文模型進行對比測試,通過像素準確率、區域重合度、檢測速度和參數數量等指標對模型性能做出評價。表2為不同語義分割模型性能參數對比。

表2 不同網絡模型性能對比

由表2可以看出,在像素準確率方面,本文模型的平均像素準確率與平均區域重合度分別為75.20%、70.78%,比DeeplabV3+分別高7.28、5.36個百分點;比EfficientUNet++分別高0.54、0.92個百分點;比BiSeNet V2分別高9.85、6.87個百分點,主要原因是本文模型引入的殘差多頭自注意力機制和條形池化的ASPP模塊能夠強化模型各階段判別特征的能力,聚合不同尺度池化區域獲取有效的全局上下文信息,增加模型對茶芽采摘點對象的分割精度;盡管DeeplabV3+也采用空洞卷積金字塔模塊聚合多尺度特征以提高分割準確性,但由于數據集的特殊性,目標形態均為長條形,造成DeeplabV3+存在像素預測不一致性,因此在分割精度方面略低于本文模型。在實時性方面,本文采用ConvNeXt作為前置基礎網絡降低模型計算量,檢測速度達到8.97 f/s,為DeeplabV3+、EfficientUNet++、BiSeNet V2、HRNet V2模型的2.81、2.17、1.37、2.10倍。在參數量方面,本文模型為2.06×106,低于其他網絡模型。

在所有對比模型中,BiSeNet V2的特點是雙路同時計算,模型的設計過程中缺乏多個尺度的圖像特征交換,所以BiSeNet V2并不能很好地解決本文問題;HRNet V2以及通過多次高分辨率特征淺層網絡與快速下采樣深度網絡融合,對于小目標的識別有很好的效果,但是由于自身設計的原因,參數量約為本文模型的7倍,雖在識別精度上超過了本文模型,但是在識別速度上確實有所差距;而EfficientUNet++使用EfficientNet[26]作為每層的特征提取網絡,同時使用U-Net[27]的結構來構建網絡,U-Net結構的特點是同時存在多層特征提取網絡,各層之間通過下采樣和上采樣相互聯通,從各項指標來看,雖然EfficientUNet++的精度指標與RMHSA-NeXt相近,但是由于EfficientNet的參數量較高,造成整體的檢測速度不佳。綜合評價上述精度和速度指標后,可得MHSA-NeXt模型可以實現精度與速度的均衡,具有良好的分割性能。

針對茶葉采摘點存在的目標較小、光照強度變化范圍大和背景復雜問題,選出了7幅圖像,圖10為5種模型對測試集的分割結果。

圖10 不同網絡模型語義分割結果對比

從圖10中可知,DeeplabV3+雖然在背景較為簡單的場景如第2、7行分割效果與標注結果覆蓋區域相似,但是在較為復雜的場景如第1、3、4、5行中,均有較大的誤差,尤其在第3行出現了較大面積的錯檢,基于結果,本文認為DeeplabV3+分割效果不佳的原因在于缺少高效的注意力機制和針對于茶葉采摘點特性的針對性設計。BiSeNet V2識別效果不佳,雖然在第6行沒有出現大面積漏檢,但是在第1、3行等都有較為明顯的誤差,說明BiSeNet V2的雙路并行計算的方式并不能很好地解決目標小、背景復雜、光照強度變化范圍大的問題。HRNet V2的識別效果與本文模型識別效果也如表2所示,檢測結果與標注差別較小,但其在PA和MIoU指標上略高于RMHSA-NeXt,表明多次的高低層級特征融合能有效利用小目標的圖像信息,但是由于層次堆疊過多模型的整體數據量較高,檢測速度較慢。從本文模型的識別結果來看,圖10中除了在第3行中出現部分漏檢,在其他目標中識別結果與標注圖像基本相差較小,其根據目標的實際特點,設計了殘差多頭自注意力模塊和結合條形池化的ASPP模塊,以較小的參數量,獲得了較好分割性能。

3 結論

(1)本文模型在實際茶園場景下,檢測結果的像素準確率達到75.20%,MIoU為70.78%,運行速度達到8.97 f/s,解決了該場景下茶葉采摘點目標較小、背景復雜、光照強度變化范圍大等困難,較好地完成了茶葉采摘點語義分割的任務。

(2)選擇HRNet V2、EfficientUNet++、DeeplabV3+、BiSeNet V2模型與本文模型進行對比測試,通過比較像素準確率、區域重合度、檢測速度、參數數量等指標,可以發現多次不同階段的高低層特征融合可以有效提取小目標的顯著特征,但是檢測速度就會下降,而本文中提出的兩種特殊結構,以相對較少的參數量達到了較高的準確率,平衡了模型的檢測速度與準確率。