基于注意力YOLOv5模型的自動水果識別①

曹秋陽,邵葉秦,尹 和

1(南通大學 信息科學技術學院,南通 226019)

2(南通大學 交通與土木工程學院,南通 226019)

近年來,隨著科學技術的快速發展,人工智能給人們生活帶來了便捷和智能化的服務. 水果自動識別在超市、菜市場、果園等很多場景有著重要的應用. 超市以及菜市場可以結合水果稱重,自動計算水果的價格,提高顧客購買的效率. 果園可以通過水果的檢測與識別,估計水果的收成,并利于機械化自動采摘.

目前,越來越多的國內外研究人員聚焦果蔬識別.彭紅星等[1]提出一種改進的single shot multibox detector(SSD)水果檢測模型,將SSD 模型主干網絡VGG16 替換為ResNet-101 網絡,并通過隨機梯度下降算法以及遷移學習思想優化SSD 模型,在4 種水果上的檢測精度達到88.4%. 王輝等[2]在Darknet-53 網絡的基礎上使用組歸一化代替原先的批量歸一化,繼而引入YOLOv3[3]算法構建水果檢測模型,實現水果的準確識別. Bargoti 等[4]設計了基于Faster-RCNN 的目標檢測模型實現自然環境下3 種水果的檢測. Liu 等[5]提出了single shot detector 方法,用于對象的檢測和識別,在保證準確率的同時提高了效率. 這些方法普遍存在如下問題: (1)數據集中水果種類過少; (2)模型傾向于對象的全局信息,容易忽略某些關鍵及重要的水果局部信息; (3)目標框與預測框重合時未考慮它們之間的相互關系,容易出現預測結果不精確問題.

因此,本文采用包括不同光照、不同角度等的15 種水果組成的數據集,并使用基于注意力的YOLOv5模型實現水果的準確分類和識別. 具體來說,該模型在主干網絡后增加注意力機制squeeze-and-excitation networks (SENet),通過神經網絡計算通道注意力權重,以增強水果的重要特征,減弱不重要的特征,使提取的特征更具代表性且保留局部的重要信息,提升水果識別的準確率. 同時,將原先的GIOU損失函數替換為包括邊框長寬比信息和中心點位置關系的CIOU損失函數,使預測框更加接近真實框. 實驗證明,本文基于注意力的YOLOv5 模型在準確率及速度上都優于目前最新的水果識別算法.

1 YOLOv5 模型

YOLOv5 是由Ultralytics LLC 公司提出的深度神經網絡模型. 相比于早期的YOLO 模型[3,6],YOLOv5模型體積小、速度快、精度高,受到工業界的青睞. 具體來說,對比于YOLOv4,YOLOv5 進行了如下改進.首先,對輸入圖片經過Focus 切片操作,保留了更完整的圖片下采樣的信息; 其次,采用CSPDarknet-53 主干網絡進行特征提取,分別在主干網絡以及Neck 部分設計了兩種CSP 結構用來調整殘差組件的數量以及卷積層數量; 最后,在Neck 部分輸出小、中、大3 層特征.雖然YOLOv5 主干網絡后的spatial pyramid pooling(SPP)層解決了輸入圖像特征尺寸不統一的問題,但是沒有對特征圖進行通道間的加權融合. 為此,本文通過軟自注意力的方式融合圖像特征,強調有效特征,提高水果識別的準確率.

2 基于注意力YOLOv5 模型的自動水果識別

本文實現基于注意力YOLOv5 模型的自動識別水果,流程如圖1 所示. 首先,將數據集進行預處理,接著輸入主干網絡提取特征,并使用SENet 注意力模塊得到一個與通道對應的一維向量作為評價分數; 其次,將評價分數通過乘法操作作用到feature map 的對應通道上,得到用于水果識別的有效特征; 然后,經過feature pyramid networks (FPN)[7]和path aggregation network(PAN)[8]結構將特征融合并獲得語義信息更強,定位信息更準的特征圖; 最后,經過類別分類與預測框回歸計算得到精準檢測結果.

圖1 本文方法的處理流程圖

2.1 預處理

2.1.1 Mosaic 數據增強

Mosaic 數據增強的方式參考了CutMix[9]數據增強思想. CutMix 數據增強將兩張圖片進行拼接,而Mosaic 采用4 張圖片的拼接,增加數據量的同時可以豐富檢測物體的背景,如圖2 所示.

圖2 Mosaic 數據增強

2.1.2 自適應錨框

在YOLO 系列算法中,通常對不同的數據集都會設定初始長寬的錨框. 在YOLOv3、YOLOv4 中,初始錨框都是通過單獨算法得到的,常用的是K-means 算法. 本文將這種功能嵌入至代碼中,實現了每次訓練可以自適應的計算不同訓練集中的最佳初始錨框. 本文的初始錨框為[10,13,16,30,33,23]、[30,61,62,45,59,119]、[116,90,156,198,373,326],經過計算本文最佳初始錨框為[111,114,141,121,127,141]、[150,149,159,169,195,212]、[256,173,173,292,326,298].

2.1.3 自適應縮放圖片

數據集的大小往往都是大小不一,需要對其尺寸歸一化. 然而,實際項目中的很多圖片長寬比不一致,縮放并填充后,兩端填充部分較多,存在很多冗余信息,影響模型速度及效果. 本文方法對原始圖像進行自適應填充最少的灰度值,使得圖像高度或寬度兩端的灰度值最少,計算量也會隨之減少,速度也得到提升. 具體步驟如下.

(1)圖像縮放比例. 假設原始圖像為1000×800,縮放至416×416. 將416×416 除以原始圖像相應寬高,得到系數分別為0.416 和0.52,取其較小值0.416.

(2)縮放后的尺寸. 將原始圖片寬高乘以較小的系數0.416,則寬為416,高為332.

(3)灰邊的填充值. 先將416?332=84,并采用取余的方式得需要填充的像素值84%32=20 (32 是由于網絡經過了5 次下采樣,2 的5 次方為32),兩端各10 個像素. 在測試過程中采用灰色填充,訓練過程依舊使用原始的resize 操作以提高物體的檢測、計算速度.

2.2 主干網絡

2.2.1 特征提取網絡

為了在水果圖像上提取豐富的特征,受到YOLOv5 的啟發,本文使用CSPDarknet-53 作為主干網絡. CSPDarknet-53 可以增強卷積網絡的學習能力,降低內存消耗.

CSPDarknet-53 主干網絡包括Focus、Mosaic、多次卷積、殘差結構等,其中CSP1_X 用來調整殘差組件的數量,如圖3 所示. Neck 中的CSP2_X 則是用來對卷積層數量的調整,如圖4 所示. CSPDarknet-53 提取的特征后續用于得到通道注意力.

圖3 CSP1_X 結構

圖4 CSP2_X 結構

2.2.2 SELayer

為了得到不同特征通道的權重,強化重要通道,減弱次要通道,本文使用SENet[10]注意力機制學習通道權重. SENet 可以學習通道之間的相關性,生成通道注意力. 雖然計算量有所增加,但是提取的特征更加有效.圖5 是SENet 模型示意圖. 首先,使用全局平均池化作為Squeeze 操作; 其次,使用兩個全連接層得到通道間的相關性,同時減少參數與計算量; 然后,通過Sigmoid歸一化權重; 最后,通過Scale 操作將歸一化后的權重作用在原始通道的特征上. 本文是將SELayer 嵌入至SPP[11]模塊,如圖6 所示. SPP 作為一種Inception 結構,嵌入了水果多尺度信息,聚合了不同感受野上的特征,因此使用SELayer 能夠對卷積特征通道重新加權,增強重要特征之間的相互依賴,可以學習到不同通道特征的重要程度,從而產生更好的效果并提升識別性能.

圖5 SENet 結構

圖6 改進YOLOv5 模型結構

針對全局特征差別不大(大小、形狀、顏色等),某些局部特征有差異的水果,注意力機制SENet 能夠增強水果的重要特征,減弱不重要的特征,使得提取的水果特征更加具有代表性且保留局部重要信息. 如圖7特征圖所示,本文選取前16 張特征圖,青蘋果與番石榴的大小、形狀、顏色等全局特征相似,而部分區域顏色、表面紋理以及根蒂等有所不同. 如圖7(b)、圖7(e)所示,在沒有進行SENet 操作前,兩者特征信息類似,特征像素未體現出特征的重要程度,經過SENet 操作后,如圖7(c)、圖7(f)所示,根據特征重要程度將特征像素進行重新加權計算,一方面減弱了周邊不重要的信息,另一方面突出了兩種水果局部紋理、形狀等重要特征,有利于準確識別出青蘋果與番石榴.

圖7 特征圖對比圖

2.3 FPN+PAN 特征融合

為了獲得更強的語義信息以及更為精準的位置信息實現水果準確識別,本文采用特征金字塔FPN+PAN提取多層次的特征,頂層特征包含豐富的語義信息,而底層特征具有精準的位置信息,如圖8 所示,其中,(a)區域為FPN 部分,(b)區域為PAN 部分.

圖8 FPN+PAN 結構圖

FPN 設計了自頂向下和橫向連接的結構,這樣的好處是既利用了頂層語義特征(利于分類),又利用了底層的高分辨率信息(利于定位),如圖9 所示.

圖9 FPN 結構圖

本文在FPN 后增加自底向上的特征金字塔PAN,將底層的特征信息通過下采樣的方式進行融合,將底層定位信息傳送至頂層,這樣的操作是對FPN 的補充,將底層的強定位特征傳遞上去.

通過組合FPN+PAN 兩個模塊,對不同的檢測層進行參數的聚合,增強語義信息的同時,提高目標的定位精度從而全面的提升模型的魯棒性和準確率.

2.4 損失函數

2.4.1GIOU

YOLOv5 采用GIOU_Loss[12]作為bounding box的損失函數. 具體來說,對于兩個bounding boxA、B(如圖10),首先,算出A、B的最小外接矩形C; 其次,計算C中沒有覆蓋A和B的面積(即差集)占C總面積的比值; 最后,用A與B的IOU減去這個比值:

圖10 GIOU 示意圖

相比于IOU,GIOU一方面解決了當預測框與目標框不相交(IOU=0)時損失函數不可導的問題; 另一方面,當兩個預測框大小相同、IOU相同時,IOU損失函數無法區分兩個預測框相交的不同之處,GIOU則緩解這種情況的發生.

但是,如圖11 所示,當預測框與目標框重疊時,則GIOU的值與IOU值相同,它們的效果一致,因此難以區分兩者相對的位置關系.

圖11 目標框與預測框重疊,GIOU=IOU=0.85

2.4.2CIOU

針對GIOU_Loss 損失函數所產生的問題,本文采用CIOU_Loss[13]替換了GIOU_Loss. GIOU_Loss 解決了邊框不重合的問題,而CIOU_Loss 在其基礎上不僅考慮了邊框重合問題,而且將邊框高寬比和中心的位置關系等信息也考慮進去,使得預測框的回歸速度與精度更高.

CIOU是將真實框與預測框之間的距離、重疊率、邊框尺度以及懲罰因子均考慮進去,使得目標邊框回歸更加穩定,有效的解決IOU在訓練過程中發散的問題,如圖12 所示.

圖12 CIOU 示意圖

式(1)為CIOU公式:

其中,ρ2(b,bgt)即圖10 中預測框與真實框中心點之間的歐式距離d,c表示同時包含真實框與預測框最小閉包矩形框的對角線距離.

式(2)為懲罰項αv中α 的公式:

式(3)為懲罰項 αv中v的公式:

其中,wgt和hgt分別表示真實框的寬和高,w和h分別表示預測框的寬和高.

式(4)為CIOU在回歸時Loss 的計算公式:

如圖13 所示,目標框與預測框重合時,CIOU值也不相同.c值相同時,通過目標框與預測框中心點的歐式距離與對角線的比值d,有效度量兩者位置關系,損失函數能夠有效收斂.

圖13 目標框與預測框重疊,CIOU 值不同

3 實驗

3.1 數據采集與預處理

本文的水果數據集部分來自于網上公開數據集,部分來自于手機拍攝的數據,所用數據均為模擬稱重時俯拍的水果圖片. 水果類別共有15 種,共計13676張,訓練集、驗證集、測試集的比例為8:1:1 (訓練集10940 張,驗證集和測試集均為1368 張),具體類別及數量如表1 所示.

表1 數據集表

3.2 實驗配置

本文實驗是在深度學習開發框架PyTorch 下進行,工作站的配置為Ubuntu 16.04.6、內存64 GB、顯存12 GB、GPU 為NVIDIA TITAN Xp、CUDA 10.2 版本以及CUDNN 7.6.4.

3.3 模型訓練

模型訓練過程中,epoch 共100 次,學習率為0.01,batch_size 為16,權重衰減數為0.000 5. 訓練過程中,模型訓練集損失函數損失值(box、objectness、classification)、驗證集損失值(val box、val objectness、val classification)、查準率(precision)、召回率(recall)以及平均精度(mAP@0.5、mAP@0.5:0.95)如圖14. 圖15給出15 類水果在驗證集上的P-R 曲線圖.

圖14 各項性能指標

圖15 P-R 曲線圖

3.4 CIOU 效果驗證

為了證明CIOU的有效性,我們進行了對比實驗.在YOLOv5 模型的基礎上,將GIOU損失函數改為對應的CIOU損失函數. 實驗結果如表2 所示.

從表2 中可以看出,利用CIOU作為邊框回歸損失函數,模型mAP 值為97.72%,提升1.57%,證明了CIOU損失函數的有效性.

表2 CIOU 效果驗證性能對比

3.5 SELayer 效果驗證

為了證明SELayer 的有效性,我們同樣進行了對比實驗. 在YOLOv5+CIOU 模型的基礎上增加注意力模塊SELayer. 實驗結果如表3 所示.

從表3 中,可以看出,在YOLOv5+CIOU 的基礎上增加SENet 注意力機制模塊,即本文基于注意力YOLOv5模型,mAP 值為99.10%,提升了1.38%,精度提升的同時,模型的速度并沒有下降,證明了SELayer 的有效性.

表3 SELayer 效果驗證性能對比

如圖16 所示,在形狀、顏色、紋理、大小類似的兩種水果中,圖16(a)為蘋果,圖16(b)為番石榴,模型能夠準確識別.

圖16 SELayer 效果對比

3.6 模型魯棒性檢驗

為了驗證本文方法的魯棒性,本文檢測了15 種水果,并分別考慮了光照、遮擋等因素. 如圖17–圖20.

(1)不同光照. 如圖17、圖18.

圖17 光照較強下測試效果對比

圖18 光照較弱下測試效果對比

(2)有遮擋. 如圖19.

圖19 有遮擋測試效果對比

(3)同類別不同品種. 如圖20.

圖20 同類別不同品種測試效果對比

通過對比發現,本文模型在遮擋、不同光照、多目標等情況下水果的識別效果更好、魯棒性更好,輸出的預測框相比更符合目標水果.

3.7 與最新方法的對比

為了驗證方法的有效性,除了對比YOLOv5 模型,本文對比了最新主流的Faster-RCNN、YOLOv4 模型,如表4 所示.

由表4 可見,Faster-RCNN 的mAP 為95.49%,YOLOv4 的mAP 為95.39%,YOLOv5 的mAP 為96.15%,而本文方法的mAP 為99.10%,識別速度到82 幀/s,在準確率及速度上都優于其他主流的對比方法.

表4 模型性能對比

4 結語

本文采用基于注意力YOLOv5 算法模型實現15類水果的自動識別. 實驗表明,本文方法是魯棒的,并且在水果識別準確率和識別速度上都優于主流Faster-RCNN、YOLOv4 和傳統YOLOv5 算法. 在后續的研究中,將考慮更多種類的水果,在保證水果種類的多樣性時也能夠保證模型的泛化能力以及識別準確率與速度.