基于YOLO v5與層次化分類算法的生活垃圾識別研究

吳 超 蔣鵬飛 呂 剛

（1.合肥學(xué)院生物食品與環(huán)境學(xué)院；2.合肥學(xué)院先進制造工程學(xué)院 安徽合肥 230601）

隨著生活垃圾分類立法、執(zhí)法等工作的順利開展以及人工智能的飛速發(fā)展，視覺處理中的垃圾目標檢測對垃圾處理工廠提升分揀效率、降低工人勞動強度等方面有著重要的研究意義。近年來基于深度學(xué)習(xí)中目標檢測準確率高、模型學(xué)習(xí)速度快等特點，?其在日常生活中的應(yīng)用場景越來越廣泛，眾多學(xué)者也轉(zhuǎn)向了目標檢測的研究領(lǐng)域。深度學(xué)習(xí)目標檢測根據(jù)設(shè)計理念的不同可分為有無anchor[1]兩大類：在anchor出現(xiàn)前，最初R-CNN[2]算法使用的Selective Search[3]方法以及滑動窗口等方法，其中Selective Search方法是能夠分割出多達上千個候選區(qū)域，依據(jù)圖像顏色、紋理、合并后的最小總面積等原則生成兩個大概率存在目標區(qū)域，最終輸出候選區(qū)域。Anchor的首次出現(xiàn)是在Faster R-CNN[1]中，一組3×3的Anchor框是人為地按照3種長寬比(0.5,1,2)，3種縮放比(8,16,32)設(shè)置成，故并不適用于所有數(shù)據(jù)集。直到Y(jié)OLO（You Only Look Once）的出現(xiàn)，在YOLO v2[4]中提出使用K-means聚類來自動生成一組適用于特定數(shù)據(jù)集的特定尺寸的anchor，簡化工作量的同時，可以使網(wǎng)絡(luò)的檢測效果更好。

深度學(xué)習(xí)目標檢測識別算法飛速發(fā)展的成果也呈現(xiàn)在生活垃圾檢測識別領(lǐng)域：Ying Wang等人[5]采用RESNET作為特征提取網(wǎng)絡(luò)，通過在城市自然環(huán)境中融合垃圾圖像的方式擴充得到更多的數(shù)據(jù)集樣本，從而提高了模型的檢測精度和準確率。彭昕昀[6]采用利用數(shù)據(jù)增強對SSD算法進行垃圾識別檢測，但SSD的檢測框準確率較低。Abeywickrama等人[7]將垃圾分類視為圖像分類，結(jié)合支持向量機與卷積神經(jīng)網(wǎng)絡(luò)對垃圾進行了識別分類，得到較高的識別率。寧凱等人[8]對YOLO v3[9]網(wǎng)絡(luò)進行輕量化改進后在自制垃圾數(shù)據(jù)集上檢測準確率較高，但是基于Darknet框架的YOLOv3在環(huán)境配置上較為復(fù)雜。

為了簡化生活垃圾目標識別算法的網(wǎng)絡(luò)結(jié)構(gòu)，選取了輕量化的YOLO v5[10]算法，并且為提高其在生活垃圾上的檢測精度，在識別外觀相似物體時擁有更高的的泛化性，本文提出了基于YOLO v5與層次化分類結(jié)合的生活垃圾識別算法。首先將部分相似類別的垃圾數(shù)據(jù)合并，并對類別較少的數(shù)據(jù)集進行數(shù)據(jù)擴充，在一定程度上緩解了數(shù)據(jù)分布不均衡的問題。對這些預(yù)處理數(shù)據(jù)使用VGG16[11]網(wǎng)絡(luò)依次訓(xùn)練得到八個分類器，將其與YOLO v5網(wǎng)絡(luò)模型的輸出融合，當檢測網(wǎng)絡(luò)輸出類別與分類器對應(yīng)的類別相同時，根據(jù)檢測模型輸出的位置信息將圖片進行裁剪，輸入到對應(yīng)的分類器中進行具體的分類，用于相似類別垃圾的具體分類。再將預(yù)測的類別提供給檢測模型的損失函數(shù)訓(xùn)練，以達到提升準確率的目的，且此改進算法結(jié)構(gòu)簡單清晰，訓(xùn)練方便。

一、算法設(shè)計分析

（一）YOLO v5算法介紹。YOLO系列是非常經(jīng)典的one-stage[12]目標檢測算法，YOLO v1[13]解決了通過直接擬合坐標位置、寬高存在很明顯的缺陷；YOLOv2改為偏移量預(yù)測：采用了anchor機制局部預(yù)測，但YOLO v2還是存在小目標預(yù)測不準確的問題；YOLO v3隨著resnet[14]的出現(xiàn)網(wǎng)絡(luò)更深，backbone改為Darknet-53[15],檢測頭改為多尺度；YOLO v4[16]在bakbone更改為采用CSPDarknet53，數(shù)據(jù)增強采用了cutmix and mosaic,class label smoothing，網(wǎng)絡(luò)結(jié)構(gòu)采用了SPP、PAN、SAM網(wǎng)絡(luò)。

YOLO v5在YOLO v4的基礎(chǔ)上由Ultralytics公司提出，輸入端依然采用Mosaic的數(shù)據(jù)增強方式，即將四張圖片進行隨機裁剪，再拼接到一張圖上作為訓(xùn)練數(shù)據(jù)，提出了新的自適應(yīng)聚類設(shè)定anchor大小，將預(yù)測框和真實框的groundtruth進行比對，計算兩者差距后反向更新迭代網(wǎng)絡(luò)參數(shù)，同時能夠自適應(yīng)檢測圖片的大小縮放。如圖1所示，在Backbone上的主要是采用了Focus結(jié)構(gòu)，主干網(wǎng)絡(luò)采用CSP1_X結(jié)構(gòu)，在Neck上采用了FPN結(jié)構(gòu)和PAN結(jié)構(gòu)。由于YOLO v5模型的目標檢測功能還能夠定位出垃圾所在位置，為我們生活垃圾識別的層次化分類算法研究提供了直接支持。

（二）基于YOLO v5與層次化分類算法的生活垃圾識別算法。YOLO v5極大地提升了算法的靈活性和速度的同時，依然具有較高的識別準確率。但是在實際實驗過程中發(fā)現(xiàn)，該算法在檢測外觀相似物體的過程中，類別的判定效果不佳，這極大的限制了垃圾目標識別準確率的提升。為了緩解這一不足，我們提出了基于YOLO v5與層次化分類算法的生活垃圾識別算法，該算法任務(wù)的框架如圖1所示。

圖1 改進算法框架圖

算法核心是將原YOLO v5網(wǎng)絡(luò)輸出的三個不同尺寸的特征提取面，對特征提取面上的每個點構(gòu)建anchor，Anchor的構(gòu)建是通過K-mean聚類以及對應(yīng)本層特征提取面的尺寸得到的，再通過IOU和非極大值抑制篩選anchor，訓(xùn)練得到預(yù)測框和類別，根據(jù)預(yù)測框的位置對歸并類所指的位置區(qū)域進行裁剪，將裁剪的圖片送入已訓(xùn)練好的VGG16網(wǎng)絡(luò)進行精確分類，將得到的類別重新賦予YOLOv5的類別輸出，完成改進網(wǎng)絡(luò)預(yù)測類別的替換。經(jīng)過類別損失函數(shù)優(yōu)化，進行迭代訓(xùn)練，計算最終準確率。

（三）損失函數(shù)。首先公式化定義基于YOLO v5與層次化分類算法的生活垃圾識別任務(wù)。由于希望控制在檢測過程中的時間成本，需要對分類模型傳入數(shù)據(jù)進行訓(xùn)練，以得到預(yù)訓(xùn)練模型參數(shù)。假設(shè)N表示樣本數(shù)，概率分布p為期望輸出，概率分布q為實際輸出，則分類任務(wù)可以公式化為：

其中L(p,q)表示交叉熵損失函數(shù)，M表示類別數(shù)，p(xij)表示變量如相同即為1，否則為0；q(xij)表示樣本i屬于類別j的預(yù)測概率。

而檢測識別部分class probability score采用了BCELoss（交叉熵損失），objectness score采用了BECWithLogitsLoss（BCEloss和sigmoid融合），bounding box采用了GIOU Loss：

其中A表示預(yù)測框，B表示標簽框，C表示能夠把A與B圍起來的最小框，IoU表示預(yù)測框與真實框的交并比，C(A?B)表示C的面積減去(A?B)的面積。

改進后的算法損失函數(shù)如式（6）所示：

其中Lcls(ci,)是基于原算法改進的交叉熵損失函數(shù)，表示檢測類別損失，其中表示真實類別，ci表示由（1）得到的預(yù)測類別。

改進算法的損失函數(shù)主要在Lcls部分聯(lián)合分類網(wǎng)絡(luò)的輸出做了相應(yīng)的改進，簡化成更簡單的類別細分任務(wù)，同時提高類別分類準確率。

二、試驗及結(jié)果分析

（一）實驗數(shù)據(jù)與實驗環(huán)境。本文試驗數(shù)據(jù)集在VOC2007基礎(chǔ)上擴充到44個類別包含19933張生活垃圾圖片，其中訓(xùn)練集13953張、測試集5980張分別用于原模型的訓(xùn)練和測試。后對其中的18個類別根據(jù)圖片的相似度合并為8個類別，制成總類別為34的數(shù)據(jù)集。通過對原數(shù)據(jù)標簽格式轉(zhuǎn)換、劃分數(shù)據(jù)類別、數(shù)據(jù)增強等方法完成數(shù)據(jù)預(yù)處理，緩解因數(shù)據(jù)不平衡對模型訓(xùn)練的影響。

本次實驗操作系統(tǒng)為Ubuntu18，深度學(xué)習(xí)框架為Pytorch1.4，顯卡為 NVIDIA Quadro P5000，加速環(huán)境為CUDA11。

（二）改進模型訓(xùn)練。本次實驗分類模型的訓(xùn)練采用VGG16網(wǎng)絡(luò)，根據(jù)數(shù)據(jù)標簽的位置信息，將指定的目標圖片對應(yīng)位置裁剪得到分類訓(xùn)練數(shù)據(jù)，對類別較少的數(shù)據(jù)采用數(shù)據(jù)增強進行圖片擴充，圖片預(yù)處理后尺寸統(tǒng)一設(shè)置為90×90后輸入網(wǎng)絡(luò)得到8個預(yù)訓(xùn)練模型。檢測模型部分采用最小的YOLO v5s預(yù)訓(xùn)練模型，原始圖片尺寸統(tǒng)一設(shè)置為416×416，先后訓(xùn)練出原始44個類別的數(shù)據(jù)的檢測模型以及相似類別合并后改進的34個類別的數(shù)據(jù)的檢測模型。圖2為改進模型進行3500次迭代的類別損失函數(shù)圖像。

圖2 改進的YOLO v5損失函數(shù)圖像

訓(xùn)練參數(shù)betch size為32，共迭代次數(shù)為3500次，由圖中可以看出，改進模型的損失函數(shù)在迭代到3000次左右時收斂到0.01附近，在迭代次數(shù)達到3500次后完成訓(xùn)練，得到最終檢測識別模型。

（三）實驗測試及結(jié)果分析。訓(xùn)練得到的改進模型通過在測試集上測試，得到的檢測效果如圖3所示。

圖3 多場景下檢測識別結(jié)果

圖4中顯示為多個類別在各種場景下的檢測識別結(jié)果，效果不錯，訓(xùn)練后的模型基本能夠準確檢測出目標位置，識別出其類別信息，且置信度較高。圖4中在光線不平衡的環(huán)境下依然能夠獲得不錯的檢測識別結(jié)果，對于角落小物體的識別依然較為出色。說明改進模型在應(yīng)對復(fù)雜環(huán)境下的生活垃圾檢測識別具有較高的普適性。

圖4 光線復(fù)雜環(huán)境下的檢測識別結(jié)果

為評估YOLO v5與層次化分類算法對生活垃圾識別的有效性，本文采用平均準確率均值mAP（mean Average Precision）來作為衡量改進的垃圾目標識別模型性能的指標。mAP是指所有類別目標的準確率AP（Average Precision）的均值。對改進模型與原模型多次測試結(jié)果比較，繪制出表1。

表1 網(wǎng)絡(luò)模型改進前后測試數(shù)據(jù)結(jié)果比較

由表1可以得出，原模型的AP88為33.1%，在垃圾目標識別情景下取得了不錯的識別準確率。改進后的模型的AP88比原模型高出了0.5%，比YOLO v3高出了26.2%。此外改進模型的AP89、AP90較各模型均有較大程度的提高，且占據(jù)的優(yōu)勢越來越大，說明改進的模型在預(yù)測框與真實框具有更大的重疊區(qū)域時，表現(xiàn)得比原模型更為出色，準確率也更高，其出現(xiàn)誤識別的幾率更小。

三、結(jié)語

針對復(fù)雜環(huán)境下較多類別的垃圾檢測識別，本文采用了基于YOLO v5與層次化分類算法，包括對數(shù)據(jù)預(yù)處理、算法原理分析、檢測和分類模型的訓(xùn)練及后期優(yōu)化。通過大量的實驗對比發(fā)現(xiàn)，改進后的YOLO v5算法能夠有效提升生活垃圾識別的準確率，在復(fù)雜環(huán)境下依然能夠較原模型發(fā)揮出較大的優(yōu)勢。由于條件限制，本文所使用的數(shù)據(jù)集數(shù)量相對欠缺，對最終模型的表現(xiàn)關(guān)聯(lián)較大，導(dǎo)致小物體的特征學(xué)習(xí)較為困難，在測試數(shù)據(jù)集上的泛化性有所欠缺。接下來的工作中，還需要對數(shù)據(jù)樣本進行擴充優(yōu)化分類模型，在保障現(xiàn)有性能的基礎(chǔ)上解決更低IoU的識別準確率問題，進一步提高本算法的魯棒性。