基于改進YOLOv5的奶牛乳頭目標檢測與數(shù)據(jù)集構(gòu)建

夏事成, 王 磊,2*, 王成軍, 席橫流

(1.安徽科技學院 機械工程學院,安徽 鳳陽 233100;2.巢湖學院 電子工程學院,安徽 巢湖 238024;3.安徽理工大學 人工智能學院,安徽 淮南 232001)

目標檢測作為計算機視覺領(lǐng)域的核心問題之一[1-4],是找出圖像中所有感興趣的目標,確定它們的類別和位置,其檢測效果的準確性和實時性是衡量機器視覺的重要指標。目標檢測已在眾多領(lǐng)域得到廣泛的運用[5-10]。但由于在目標檢測過程中存在外界因素的干擾,目標檢測依然是計算機視覺領(lǐng)域最具有挑戰(zhàn)性的問題。目前主流的檢測方法是基于深度學習的方法。分為兩種[11]:一種是存在候選區(qū)域的二階目標檢測法(Tow Stage),包括經(jīng)典的R-CNN、SPP-Net、faster R-CNN等算法。另一種是不存在候選區(qū)域的一階目標檢測法(One Stage),主要包含YOLO、SSD、RetinaNet等算法[12-15]。其中,YOLO算法因其迭代迅速、檢測速度快,加快了目標檢測領(lǐng)域中實時性問題的解決,因而廣受歡迎[16-17]。

隨著機器視覺技術(shù)的不斷發(fā)展,科研人員一直致力于YOLO算法的維護與改進。Lalak等[18]使用YOLO算法從無人機圖像中自動檢測非典型航空障礙物;Wang等[19]將YOLOv5與注意機制結(jié)合通過計算重疊率來構(gòu)建數(shù)據(jù)集,提出YOLO_CBAM網(wǎng)絡(luò)模型的精度與召回率均略高于傳統(tǒng)的YOLOv5算法;Yue等[20]對YOLO算法增加了一個新的預測層YOLO_head,實現(xiàn)相比于YOLOv4算法對小目標檢測的平均精度(AP)提高約12%,總平均精度提高約3%;Yuan等[21]基于YOLOv5增加了注意力機制突出特征圖中的重要信息,弱化次要信息,增強了網(wǎng)絡(luò)的識別能力。

YOLOv5算法的發(fā)展雖然迅速,但在環(huán)境比較復雜以及目標比較密集的場景下,因其計算復雜度較大,導致檢測速度降低,檢測的實時性變差[22]。針對此問題,基于YOLOv5對中大目標的檢測提出CN-YOLOv5算法,并基于YOLOv5s創(chuàng)建奶牛乳頭圖像數(shù)據(jù)集CN-YOLOv5s model,對改進后的算法進行效果驗證。

1 YOLOv5算法

YOLOv5算法的識別過程主要包括輸入端、主干(Backbone)、脖子(Neck)以及輸出端(Prediction)等4個部分,其網(wǎng)絡(luò)模型如圖1所示。

圖1 YOLOv5網(wǎng)絡(luò)模型圖Fig.1 YOLOv5 network model

輸入端包括Mosaic數(shù)據(jù)增強、自適應(yīng)錨框計算、自適應(yīng)圖片縮放等3個部分,達到預處理圖片、提高檢測速度的目的。在日常生活中小目標的數(shù)量一般比中目標和大目標低很多。而Coco數(shù)據(jù)集包含大量的小目標,但是小目標的分布并不均勻,因此使用Mosaic數(shù)據(jù)增強處理,在提高網(wǎng)絡(luò)魯棒性、小目標檢測準確性、降低設(shè)備檢測要求、加速檢測速度等方面起到了良好效果。

網(wǎng)絡(luò)主干部分主要是Focus結(jié)構(gòu)。Focus結(jié)構(gòu)是對輸入端輸入的圖片進行一個類似切面的操作(圖2)。這種結(jié)構(gòu)可以有效保留圖片中更多的特征。再利用3×3,步長為2的卷積核對圖片進行卷積以減少信息梯度重復,起到減小計算量并增強圖片梯度表現(xiàn)的作用。

圖2 切片結(jié)構(gòu)示意圖Fig.2 Schematic diagram of the slice structure

YOLOv5網(wǎng)絡(luò)結(jié)構(gòu)脖子部分的重點是FPN(Feature Pyramid Network)與PAN(Path Aggregation Network)結(jié)構(gòu)。FPN是通過對圖片特征的上采樣方式將高層信息進行傳遞融合,得到進行預測的特征圖。PAN是把特征圖進行步幅為2的3×3卷積,并與上一層下采樣后的層進行橫向連接,二者進行特征融合,以此類推。YOLOv5網(wǎng)絡(luò)結(jié)構(gòu)的輸出端主要包含CIOU_Loss損失函數(shù)與NMS(Non Maximum Suppression)非極大值抑制兩個部分。CIOU損失函數(shù)如圖3所示。

圖3 CIOU示意圖Fig.3 Schematic diagram of CIOU

(1)

(2)

圖3中c是預測框和真實框的最小外接矩陣對角線距離;d是預測框與真實框的中心點距離;v是用來衡量長寬比一致性的參數(shù)。式(2)中w為寬,h為高,gt代表真實框。YOLOv5中非極大值抑制是通過將圖片分割成多個網(wǎng)格。每個網(wǎng)格的大小相同,并且讓多個網(wǎng)格每個都可以預測出多個預測框。預測出來的每個邊界框都包含物體的中心位置(x,y),物體的高h,物體的寬w,預測的置信度和預測類別等5個信息。

2 YOLOv5網(wǎng)絡(luò)結(jié)構(gòu)的改進

基于YOLOv5算法,對原始模型與網(wǎng)絡(luò)結(jié)構(gòu)進行優(yōu)化,以更優(yōu)秀地適應(yīng)奶牛乳頭目標檢測,具體改進方法:(1)在主干部分后增加SE注意力模塊結(jié)構(gòu),該結(jié)構(gòu)通過全連接網(wǎng)絡(luò),根據(jù)損失將主干部分獲取的特征權(quán)重再進行自主學習,對特征重新校準,增大有效特征通道的權(quán)重,提高識別效率。但也增加了一些參數(shù)和計算量。(2)將網(wǎng)絡(luò)結(jié)構(gòu)頸部中的SPP結(jié)構(gòu)改進為SPPF結(jié)構(gòu),SPPF利用多個小尺寸池化核連接的方法代替了SPP中等差增大的單池化,大幅度減少計算量,達到在不影響識別準確度的同時,加快奶牛乳頭識別速度的結(jié)果。(3)擠奶機的相機采集到的奶牛乳頭照片中幾乎沒有小目標,因此在原PAN結(jié)構(gòu)中,輕量化對小目標檢測的特征融合過程,起到提升檢測速度的作用。

2.1 YOLOv5網(wǎng)絡(luò)結(jié)構(gòu)中SE注意力模塊的引入

如圖4所示,在主干部分后增加SE注意力模塊,YOLOv5原模型使用上采樣獲得大尺寸的特征圖,并和主干網(wǎng)中相同尺寸的特征圖以拼接的方式實現(xiàn)目標信息的聚合。SE注意力模塊可以建立特征通道之間的相互依賴關(guān)系并執(zhí)行特征重新校準,通過自適應(yīng)學習使用全局信息選擇性地強調(diào)可靠的信息性特征,并抑制對當前任務(wù)不太可靠的特征。達到對目標更加準確的識別。

圖4 CN-YOLOv5網(wǎng)絡(luò)結(jié)構(gòu)模型圖Fig.4 CN-YOLOv5 network structure model

SE通過壓縮(Squeeze)模塊和激勵(Exciation)模塊實現(xiàn)所述功能。如圖5所示,SE注意力模塊主要組成分為Squeeze和Excitation兩部分。輸入特征圖大小為W×H×C,C表示通道數(shù),W、H表示特征圖寬、高。

在輸入端輸入數(shù)據(jù)后通過Squeeze操作,對每個特征圖做全局平均池化,平均成一個實數(shù)值,如圖6所示。

圖6 SE模塊中的squeeze示意圖Fig.6 Schematic diagram of squeeze in SE model

Excitaton操作如圖7所示,經(jīng)過Squeeze操作后,網(wǎng)絡(luò)輸出1×1×C大小的特征圖,利用權(quán)重W來學習C個通道的相關(guān)性。

得到1×1×C向量后將SE模塊計算出來的各通道權(quán)重值分別和原特征圖對應(yīng)通道的二維矩陣相乘,得出的結(jié)果輸出。原有特征向量為W×H×C,SE模塊的引入一方面降低了網(wǎng)絡(luò)計算量,一方面增加了網(wǎng)絡(luò)的非線性能力。實現(xiàn)提升重要特征、抑制不重要特征的功能。

2.2 SPP結(jié)構(gòu)的改進

YOLO采用如圖8所示的統(tǒng)一步長但不同尺寸的卷積核進行卷積實現(xiàn)SPP,統(tǒng)一步長代表輸出的特征圖尺寸一樣,只是對區(qū)域的敏感性不一樣,再通過Concat函數(shù)按通道拼接后用1×1卷積,實現(xiàn)特征融合。

如圖9所示的SPPF結(jié)構(gòu)用3個池化核大小分別為5×5、9×9、13×13,步長為1的最大池化串行的方法,大幅降低計算量,使SPPF結(jié)構(gòu)相較于SPP結(jié)構(gòu),在不影響檢測準確度的情況下,實現(xiàn)更快的目標檢測。

2.3 CN-PAN結(jié)構(gòu)的引入

在擠奶機器人工作時,奶牛處于擠奶工位上,擠奶機器人采集到的奶牛乳頭照片中,幾乎不會出現(xiàn)小目標的情況。根據(jù)這一特點,對YOLOv5的網(wǎng)絡(luò)結(jié)構(gòu)模型進行優(yōu)化,簡化了對小目標識別的過程。具體結(jié)構(gòu)如圖10所示。

圖10 CN-PAN結(jié)構(gòu)示意圖Fig.10 Schematic diagram of CN-PAN structure

3 實驗與分析

在Windows 10 操作系統(tǒng)中使用YOLOv5框架進行網(wǎng)絡(luò)模型部署,所用CPU為Intel Core i7 9750H處理器,顯卡為NVIDIA GeForce GTX1650;Torch版本為1.12.1;opencv-python版本為4.6.0.66。

3.1 數(shù)據(jù)集創(chuàng)建

數(shù)據(jù)集訓練圖片共208張。為保證實驗數(shù)據(jù)集的可靠性,對不同姿勢與位置、密集奶牛乳頭及不同遮擋程度的奶牛乳頭進行采集。使用Labeling軟件進行手動標注,標注為Pascalvoc格式。將數(shù)據(jù)集的80%作為訓練集,20%作為測試集。模型訓練輸入圖像尺寸設(shè)置為640×640,epochs為300,batch-size為8。

3.2 評估指標

使用目標檢測通過評估指標定量分析改進YOLOv5算法對奶牛乳頭檢測的性能,分別為精確率(Precision)、召回率(Recall)、平均精度率(MAP)、驗證損失Validation-Loss、檢測速度FPS。設(shè)定置信度閾值為0.45。

(3)

(4)

(5)

其中,P表示所有檢測出的目標檢測正確率;R表示所有正樣本中檢測正確率;AP表示不同召回率下精確率的均值;TP為正樣本預測為正類個數(shù);FP為負樣本預測為正類個數(shù);FN為正樣本預測為負類個數(shù)。

利用YOLOv5算法建立對奶牛乳頭檢測的模型庫CN-YOLOv5 model。模型指標如表1所示。

表1 模型評價指標

檢測結(jié)果如圖11所示,第一行3張為原圖,第二行3張為檢測后的圖片。

圖11 奶牛乳頭檢測效果Fig.11 Detection effects of cows' nipples

3.3 模型改進前后檢測速度指標

選取204張圖片作為樣本。利用CN-YOLOv5 model數(shù)據(jù)集對比CN-YOLOv5s與YOLOv5s對這204張圖片的檢測速度。測試結(jié)果如表2所示,改進后的算法相較于原算法,在相同的檢測環(huán)境中檢測速度提升了約13%。

表2 2種算法在奶牛乳頭模型的檢測結(jié)果

4 結(jié)論

本研究為實現(xiàn)對奶牛乳頭的快速識別與檢測,針對擠奶機器人拍攝到的奶牛乳頭大小及形狀特性展開研究,提出CN-YOLOv5算法。該算法在不影響識別準確度的同時,加快了奶牛乳頭識別速度。同時利用YOLOv5算法,建立牛乳頭檢測模型庫CN-YOLOv5 model。模型指標中Precision為0.84,Recall為0.45,MAP為0.49,Box_loss為0.03,并用此模型庫進行對比檢測。實驗結(jié)果表明,改進后的YOLOv5算法相對于傳統(tǒng)YOLOv5算法在幾乎不影響準確率的情況下,檢測速度提高約13%。本研究通過對YOLOv5的改進,實現(xiàn)了在復雜環(huán)境和密集目標場景中對擠奶機器人拍攝到的奶牛乳頭進行快速識別,并創(chuàng)建奶牛乳頭數(shù)據(jù)集。為擠奶機器人的智能化、自動化提供有益的技術(shù)支撐。