基于注意力YOLOv5模型的自動(dòng)水果識(shí)別①

曹秋陽(yáng),邵葉秦,尹 和

1(南通大學(xué) 信息科學(xué)技術(shù)學(xué)院,南通 226019)

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

近年來(lái),隨著科學(xué)技術(shù)的快速發(fā)展,人工智能給人們生活帶來(lái)了便捷和智能化的服務(wù). 水果自動(dòng)識(shí)別在超市、菜市場(chǎng)、果園等很多場(chǎng)景有著重要的應(yīng)用. 超市以及菜市場(chǎng)可以結(jié)合水果稱(chēng)重,自動(dòng)計(jì)算水果的價(jià)格,提高顧客購(gòu)買(mǎi)的效率. 果園可以通過(guò)水果的檢測(cè)與識(shí)別,估計(jì)水果的收成,并利于機(jī)械化自動(dòng)采摘.

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

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

1 YOLOv5 模型

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

2 基于注意力YOLOv5 模型的自動(dòng)水果識(shí)別

本文實(shí)現(xiàn)基于注意力YOLOv5 模型的自動(dòng)識(shí)別水果,流程如圖1 所示. 首先,將數(shù)據(jù)集進(jìn)行預(yù)處理,接著輸入主干網(wǎng)絡(luò)提取特征,并使用SENet 注意力模塊得到一個(gè)與通道對(duì)應(yīng)的一維向量作為評(píng)價(jià)分?jǐn)?shù); 其次,將評(píng)價(jià)分?jǐn)?shù)通過(guò)乘法操作作用到feature map 的對(duì)應(yīng)通道上,得到用于水果識(shí)別的有效特征; 然后,經(jīng)過(guò)feature pyramid networks (FPN)[7]和path aggregation network(PAN)[8]結(jié)構(gòu)將特征融合并獲得語(yǔ)義信息更強(qiáng),定位信息更準(zhǔn)的特征圖; 最后,經(jīng)過(guò)類(lèi)別分類(lèi)與預(yù)測(cè)框回歸計(jì)算得到精準(zhǔn)檢測(cè)結(jié)果.

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

2.1 預(yù)處理

2.1.1 Mosaic 數(shù)據(jù)增強(qiáng)

Mosaic 數(shù)據(jù)增強(qiáng)的方式參考了CutMix[9]數(shù)據(jù)增強(qiáng)思想. CutMix 數(shù)據(jù)增強(qiáng)將兩張圖片進(jìn)行拼接,而Mosaic 采用4 張圖片的拼接,增加數(shù)據(jù)量的同時(shí)可以豐富檢測(cè)物體的背景,如圖2 所示.

圖2 Mosaic 數(shù)據(jù)增強(qiáng)

2.1.2 自適應(yīng)錨框

在YOLO 系列算法中,通常對(duì)不同的數(shù)據(jù)集都會(huì)設(shè)定初始長(zhǎng)寬的錨框. 在YOLOv3、YOLOv4 中,初始錨框都是通過(guò)單獨(dú)算法得到的,常用的是K-means 算法. 本文將這種功能嵌入至代碼中,實(shí)現(xiàn)了每次訓(xùn)練可以自適應(yīng)的計(jì)算不同訓(xùn)練集中的最佳初始錨框. 本文的初始錨框?yàn)閇10,13,16,30,33,23]、[30,61,62,45,59,119]、[116,90,156,198,373,326],經(jīng)過(guò)計(jì)算本文最佳初始錨框?yàn)閇111,114,141,121,127,141]、[150,149,159,169,195,212]、[256,173,173,292,326,298].

2.1.3 自適應(yīng)縮放圖片

數(shù)據(jù)集的大小往往都是大小不一,需要對(duì)其尺寸歸一化. 然而,實(shí)際項(xiàng)目中的很多圖片長(zhǎng)寬比不一致,縮放并填充后,兩端填充部分較多,存在很多冗余信息,影響模型速度及效果. 本文方法對(duì)原始圖像進(jìn)行自適應(yīng)填充最少的灰度值,使得圖像高度或?qū)挾葍啥说幕叶戎底钌?計(jì)算量也會(huì)隨之減少,速度也得到提升. 具體步驟如下.

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

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

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

2.2 主干網(wǎng)絡(luò)

2.2.1 特征提取網(wǎng)絡(luò)

為了在水果圖像上提取豐富的特征,受到Y(jié)OLOv5 的啟發(fā),本文使用CSPDarknet-53 作為主干網(wǎng)絡(luò). CSPDarknet-53 可以增強(qiáng)卷積網(wǎng)絡(luò)的學(xué)習(xí)能力,降低內(nèi)存消耗.

CSPDarknet-53 主干網(wǎng)絡(luò)包括Focus、Mosaic、多次卷積、殘差結(jié)構(gòu)等,其中CSP1_X 用來(lái)調(diào)整殘差組件的數(shù)量,如圖3 所示. Neck 中的CSP2_X 則是用來(lái)對(duì)卷積層數(shù)量的調(diào)整,如圖4 所示. CSPDarknet-53 提取的特征后續(xù)用于得到通道注意力.

圖3 CSP1_X 結(jié)構(gòu)

圖4 CSP2_X 結(jié)構(gòu)

2.2.2 SELayer

為了得到不同特征通道的權(quán)重,強(qiáng)化重要通道,減弱次要通道,本文使用SENet[10]注意力機(jī)制學(xué)習(xí)通道權(quán)重. SENet 可以學(xué)習(xí)通道之間的相關(guān)性,生成通道注意力. 雖然計(jì)算量有所增加,但是提取的特征更加有效.圖5 是SENet 模型示意圖. 首先,使用全局平均池化作為Squeeze 操作; 其次,使用兩個(gè)全連接層得到通道間的相關(guān)性,同時(shí)減少參數(shù)與計(jì)算量; 然后,通過(guò)Sigmoid歸一化權(quán)重; 最后,通過(guò)Scale 操作將歸一化后的權(quán)重作用在原始通道的特征上. 本文是將SELayer 嵌入至SPP[11]模塊,如圖6 所示. SPP 作為一種Inception 結(jié)構(gòu),嵌入了水果多尺度信息,聚合了不同感受野上的特征,因此使用SELayer 能夠?qū)矸e特征通道重新加權(quán),增強(qiáng)重要特征之間的相互依賴(lài),可以學(xué)習(xí)到不同通道特征的重要程度,從而產(chǎn)生更好的效果并提升識(shí)別性能.

圖5 SENet 結(jié)構(gòu)

圖6 改進(jìn)YOLOv5 模型結(jié)構(gòu)

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

圖7 特征圖對(duì)比圖

2.3 FPN+PAN 特征融合

為了獲得更強(qiáng)的語(yǔ)義信息以及更為精準(zhǔn)的位置信息實(shí)現(xiàn)水果準(zhǔn)確識(shí)別,本文采用特征金字塔FPN+PAN提取多層次的特征,頂層特征包含豐富的語(yǔ)義信息,而底層特征具有精準(zhǔn)的位置信息,如圖8 所示,其中,(a)區(qū)域?yàn)镕PN 部分,(b)區(qū)域?yàn)镻AN 部分.

圖8 FPN+PAN 結(jié)構(gòu)圖

FPN 設(shè)計(jì)了自頂向下和橫向連接的結(jié)構(gòu),這樣的好處是既利用了頂層語(yǔ)義特征(利于分類(lèi)),又利用了底層的高分辨率信息(利于定位),如圖9 所示.

圖9 FPN 結(jié)構(gòu)圖

本文在FPN 后增加自底向上的特征金字塔PAN,將底層的特征信息通過(guò)下采樣的方式進(jìn)行融合,將底層定位信息傳送至頂層,這樣的操作是對(duì)FPN 的補(bǔ)充,將底層的強(qiáng)定位特征傳遞上去.

通過(guò)組合FPN+PAN 兩個(gè)模塊,對(duì)不同的檢測(cè)層進(jìn)行參數(shù)的聚合,增強(qiáng)語(yǔ)義信息的同時(shí),提高目標(biāo)的定位精度從而全面的提升模型的魯棒性和準(zhǔn)確率.

2.4 損失函數(shù)

2.4.1GIOU

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

圖10 GIOU 示意圖

相比于IOU,GIOU一方面解決了當(dāng)預(yù)測(cè)框與目標(biāo)框不相交(IOU=0)時(shí)損失函數(shù)不可導(dǎo)的問(wèn)題; 另一方面,當(dāng)兩個(gè)預(yù)測(cè)框大小相同、IOU相同時(shí),IOU損失函數(shù)無(wú)法區(qū)分兩個(gè)預(yù)測(cè)框相交的不同之處,GIOU則緩解這種情況的發(fā)生.

但是,如圖11 所示,當(dāng)預(yù)測(cè)框與目標(biāo)框重疊時(shí),則GIOU的值與IOU值相同,它們的效果一致,因此難以區(qū)分兩者相對(duì)的位置關(guān)系.

圖11 目標(biāo)框與預(yù)測(cè)框重疊,GIOU=IOU=0.85

2.4.2CIOU

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

CIOU是將真實(shí)框與預(yù)測(cè)框之間的距離、重疊率、邊框尺度以及懲罰因子均考慮進(jìn)去,使得目標(biāo)邊框回歸更加穩(wěn)定,有效的解決IOU在訓(xùn)練過(guò)程中發(fā)散的問(wèn)題,如圖12 所示.

圖12 CIOU 示意圖

式(1)為CIOU公式:

其中,ρ2(b,bgt)即圖10 中預(yù)測(cè)框與真實(shí)框中心點(diǎn)之間的歐式距離d,c表示同時(shí)包含真實(shí)框與預(yù)測(cè)框最小閉包矩形框的對(duì)角線(xiàn)距離.

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

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

其中,wgt和hgt分別表示真實(shí)框的寬和高,w和h分別表示預(yù)測(cè)框的寬和高.

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

如圖13 所示,目標(biāo)框與預(yù)測(cè)框重合時(shí),CIOU值也不相同.c值相同時(shí),通過(guò)目標(biāo)框與預(yù)測(cè)框中心點(diǎn)的歐式距離與對(duì)角線(xiàn)的比值d,有效度量?jī)烧呶恢藐P(guān)系,損失函數(shù)能夠有效收斂.

圖13 目標(biāo)框與預(yù)測(cè)框重疊,CIOU 值不同

3 實(shí)驗(yàn)

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

本文的水果數(shù)據(jù)集部分來(lái)自于網(wǎng)上公開(kāi)數(shù)據(jù)集,部分來(lái)自于手機(jī)拍攝的數(shù)據(jù),所用數(shù)據(jù)均為模擬稱(chēng)重時(shí)俯拍的水果圖片. 水果類(lèi)別共有15 種,共計(jì)13676張,訓(xùn)練集、驗(yàn)證集、測(cè)試集的比例為8:1:1 (訓(xùn)練集10940 張,驗(yàn)證集和測(cè)試集均為1368 張),具體類(lèi)別及數(shù)量如表1 所示.

表1 數(shù)據(jù)集表

3.2 實(shí)驗(yàn)配置

本文實(shí)驗(yàn)是在深度學(xué)習(xí)開(kāi)發(fā)框架PyTorch 下進(jìn)行,工作站的配置為Ubuntu 16.04.6、內(nèi)存64 GB、顯存12 GB、GPU 為NVIDIA TITAN Xp、CUDA 10.2 版本以及CUDNN 7.6.4.

3.3 模型訓(xùn)練

模型訓(xùn)練過(guò)程中,epoch 共100 次,學(xué)習(xí)率為0.01,batch_size 為16,權(quán)重衰減數(shù)為0.000 5. 訓(xùn)練過(guò)程中,模型訓(xùn)練集損失函數(shù)損失值(box、objectness、classification)、驗(yàn)證集損失值(val box、val objectness、val classification)、查準(zhǔn)率(precision)、召回率(recall)以及平均精度(mAP@0.5、mAP@0.5:0.95)如圖14. 圖15給出15 類(lèi)水果在驗(yàn)證集上的P-R 曲線(xiàn)圖.

圖14 各項(xiàng)性能指標(biāo)

圖15 P-R 曲線(xiàn)圖

3.4 CIOU 效果驗(yàn)證

為了證明CIOU的有效性,我們進(jìn)行了對(duì)比實(shí)驗(yàn).在YOLOv5 模型的基礎(chǔ)上,將GIOU損失函數(shù)改為對(duì)應(yīng)的CIOU損失函數(shù). 實(shí)驗(yàn)結(jié)果如表2 所示.

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

表2 CIOU 效果驗(yàn)證性能對(duì)比

3.5 SELayer 效果驗(yàn)證

為了證明SELayer 的有效性,我們同樣進(jìn)行了對(duì)比實(shí)驗(yàn). 在YOLOv5+CIOU 模型的基礎(chǔ)上增加注意力模塊SELayer. 實(shí)驗(yàn)結(jié)果如表3 所示.

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

表3 SELayer 效果驗(yàn)證性能對(duì)比

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

圖16 SELayer 效果對(duì)比

3.6 模型魯棒性檢驗(yàn)

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

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

圖17 光照較強(qiáng)下測(cè)試效果對(duì)比

圖18 光照較弱下測(cè)試效果對(duì)比

(2)有遮擋. 如圖19.

圖19 有遮擋測(cè)試效果對(duì)比

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

圖20 同類(lèi)別不同品種測(cè)試效果對(duì)比

通過(guò)對(duì)比發(fā)現(xiàn),本文模型在遮擋、不同光照、多目標(biāo)等情況下水果的識(shí)別效果更好、魯棒性更好,輸出的預(yù)測(cè)框相比更符合目標(biāo)水果.

3.7 與最新方法的對(duì)比

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

由表4 可見(jiàn),Faster-RCNN 的mAP 為95.49%,YOLOv4 的mAP 為95.39%,YOLOv5 的mAP 為96.15%,而本文方法的mAP 為99.10%,識(shí)別速度到82 幀/s,在準(zhǔn)確率及速度上都優(yōu)于其他主流的對(duì)比方法.

表4 模型性能對(duì)比

4 結(jié)語(yǔ)

本文采用基于注意力YOLOv5 算法模型實(shí)現(xiàn)15類(lèi)水果的自動(dòng)識(shí)別. 實(shí)驗(yàn)表明,本文方法是魯棒的,并且在水果識(shí)別準(zhǔn)確率和識(shí)別速度上都優(yōu)于主流Faster-RCNN、YOLOv4 和傳統(tǒng)YOLOv5 算法. 在后續(xù)的研究中,將考慮更多種類(lèi)的水果,在保證水果種類(lèi)的多樣性時(shí)也能夠保證模型的泛化能力以及識(shí)別準(zhǔn)確率與速度.