基于YOLO神經(jīng)網(wǎng)絡(luò)的垃圾檢測(cè)與分類

張 偉，劉 娜，江 洋，李清都

(1.上海理工大學(xué) 光電信息與計(jì)算機(jī)工程學(xué)院，上海 200093； 2.上海理工大學(xué) 醫(yī)療器械與食品學(xué)院，上海 200093； 3.重慶市互聯(lián)網(wǎng)信息辦公室，重慶 401120)

隨著經(jīng)濟(jì)的迅速發(fā)展，我國城鎮(zhèn)人口數(shù)量急劇增加，城市生活垃圾也日益增多，而垃圾分類則是處理垃圾公害的最佳解決方案[1]。當(dāng)下的關(guān)注熱點(diǎn)在于將垃圾進(jìn)行無害化、資源化處理，因此合理有效地對(duì)垃圾進(jìn)行回收利用具有一定的意義。對(duì)垃圾進(jìn)行有效地分類是處理垃圾回收問題的前提[2]。目前，針對(duì)人工分揀垃圾環(huán)境差、分揀效率低等問題，研究人員提出建立垃圾識(shí)別與分揀智能化和自動(dòng)化系統(tǒng)，通過利用深度學(xué)習(xí)技術(shù)來降低人工分類成本，節(jié)約人力資源[3]。

垃圾檢測(cè)屬于計(jì)算機(jī)視覺領(lǐng)域的目標(biāo)檢測(cè)任務(wù)。傳統(tǒng)的檢測(cè)方法以HOG(Histogram of Gridie) 檢測(cè)器[4]、基于滑動(dòng)窗口的VJ(Viola-Jones)檢測(cè)器和DPM(Deformable Part Model)檢測(cè)器[5]為代表。卷積神經(jīng)網(wǎng)絡(luò)(Convolutional Neural Network，CNN)架構(gòu)于2012年被首次提出，并逐漸成為了熱門研究方向。經(jīng)過多年發(fā)展，深度學(xué)習(xí)領(lǐng)域已涌現(xiàn)出VGGNet、GoogLeNet、ResNet等多種神經(jīng)網(wǎng)絡(luò)算法。最先被用于目標(biāo)檢測(cè)的卷積神經(jīng)網(wǎng)絡(luò)是R-CNN[6](Region-CNN)算法。該算法使用卷積神經(jīng)網(wǎng)絡(luò)進(jìn)行特征提取，使用Bounding Box Regression修正目標(biāo)邊界框。不過R-CNN算法也有其弊端，串行式CNN前向傳播耗時(shí)長，且CNN特征提取、SVM(Support Vector Machines)分類和邊框修正3個(gè)模塊是分開訓(xùn)練的，對(duì)于存儲(chǔ)空間的消耗較大。為了克服R-CNN的耗時(shí)問題，F(xiàn)ast R-CNN算法直接采用一個(gè)CNN提取全圖特征，減少了對(duì)存儲(chǔ)空間的占用。R-CNN和Fast R-CNN算法存在相同的缺陷，即存在耗時(shí)的Selectice Search，因此這兩種算法無法實(shí)現(xiàn)端到端的物體檢測(cè)。為了解決這個(gè)問題，文獻(xiàn)[7]提出了一種區(qū)域提議網(wǎng)絡(luò)(Region Proposal Network， RPN)，該網(wǎng)絡(luò)可以直接連接到Fast R-CNN 網(wǎng)絡(luò)的后一個(gè)卷積層，以得到對(duì)應(yīng)的檢測(cè)對(duì)象框和分?jǐn)?shù)信息，該網(wǎng)絡(luò)也被稱作為 Faster R-CNN。上述目標(biāo)檢測(cè)算法都是二階段算法，針對(duì)二階段目標(biāo)檢測(cè)算法普遍存在運(yùn)算速度慢的缺點(diǎn)，研究人員提出了YOLO(You Only Look Once)[6]算法，該算法可在一個(gè)步驟中完成對(duì)物體的分類和定位。YOLO和SSD[8](Single Shot MultiBox Detector)這類一階段算法僅使用一個(gè)CNN網(wǎng)絡(luò)就能直接預(yù)測(cè)不同目標(biāo)的類別與位置。SSD的網(wǎng)絡(luò)結(jié)構(gòu)基于VGG16構(gòu)造，其檢測(cè)速度較快。YOLO的網(wǎng)絡(luò)結(jié)構(gòu)類似于GoogLeNet，其先將圖片尺寸放縮到固定尺寸，然后將圖片傳入CNN網(wǎng)絡(luò)進(jìn)行處理，用非極大值抑制進(jìn)行邊界框[9]的冗余裁剪，從而處理掉大批冗余，得到最后的預(yù)測(cè)結(jié)果。YOLO將物體檢測(cè)作為一個(gè)回歸問題進(jìn)行求解，僅一次求解便可以得到圖像中所有物體的位置、類別以及相應(yīng)的置信概率。

本文探討了CNN的精細(xì)特征對(duì)垃圾分類任務(wù)的適用性，針對(duì)不同實(shí)驗(yàn)場(chǎng)景，制作了特定的垃圾數(shù)據(jù)集。本文提出采用基于YOLO的神經(jīng)網(wǎng)絡(luò)模型來進(jìn)行垃圾檢測(cè)與分類，可有效解決垃圾分類的人工預(yù)處理問題。

1 垃圾分類檢測(cè)流程

本文提出的垃圾檢測(cè)與分類方法設(shè)計(jì)包括數(shù)據(jù)集制作以及模型調(diào)整兩部分。垃圾檢測(cè)與分類的方法設(shè)計(jì)整體流程如圖1所示。

圖1 設(shè)計(jì)流程圖

本文實(shí)驗(yàn)使用的數(shù)據(jù)來源于相機(jī)拍攝，數(shù)據(jù)集共計(jì)20 000張圖片。垃圾樣本共計(jì)9類，分別是易拉罐、煙盒、煙頭、果皮、包裝袋、紙、塑料瓶、瓜子殼、樹葉。YOLO模型相較于R-CNN等兩階段算法擁有一定的速度優(yōu)勢(shì)。相較于YOLO模型早期版本，YOLOv4模型的識(shí)別精度更高，且有利于邊緣計(jì)算，對(duì)于項(xiàng)目落地有較強(qiáng)的可行性，因此本文選用YOLOv4模型進(jìn)行研究。

2 垃圾分類檢測(cè)方法設(shè)計(jì)

2.1 數(shù)據(jù)集的獲取

由于項(xiàng)目機(jī)器人作業(yè)地點(diǎn)為鄭州東站，故在鄭州東站進(jìn)行數(shù)據(jù)采集，采集時(shí)間為夏季，天氣以晴天為主，少量雨天。因鄭州東站環(huán)境較復(fù)雜，所以數(shù)據(jù)采集的背景選擇非常重要。本實(shí)驗(yàn)數(shù)據(jù)采集的背景選擇了30余處，基本覆蓋所有場(chǎng)景，確保了實(shí)驗(yàn)數(shù)據(jù)的完整性和準(zhǔn)確性。通過D435雙目相機(jī)獲取視頻，視頻分辨率為1 920×1 080，每秒傳遞幀數(shù)為30幀。為了獲得最佳的訓(xùn)練數(shù)據(jù)，本文充分考慮了垃圾的各種形態(tài)，選擇多角度環(huán)繞拍攝，避免垃圾之間產(chǎn)生重疊。對(duì)拍攝完的視頻進(jìn)行解析，每秒提取2張圖片，格式為JPEG。為了保證圖像的質(zhì)量和清晰度，對(duì)解析后的圖片進(jìn)行人工篩選，然后對(duì)篩選后的圖片使用labellmg工具進(jìn)行標(biāo)注。標(biāo)注時(shí)，為了保證數(shù)據(jù)的準(zhǔn)確性，需確保目標(biāo)框?yàn)槔鴺颖镜淖畲笸饨泳匦?。?duì)標(biāo)注后的數(shù)據(jù)，進(jìn)行人工驗(yàn)收，防止出現(xiàn)漏標(biāo)誤標(biāo)。后期則根據(jù)模型測(cè)試的效果，對(duì)易誤檢的目標(biāo)或背景再進(jìn)行樣本補(bǔ)充，提高識(shí)別的準(zhǔn)確率。采集的垃圾樣本如圖2所示。

圖2 垃圾樣本

2.2 數(shù)據(jù)集的劃分

對(duì)所有圖片進(jìn)行命名排序，根據(jù)標(biāo)注產(chǎn)生的txt格式文件編寫相應(yīng)的代碼使其生成訓(xùn)練集、驗(yàn)證集和測(cè)試集。

訓(xùn)練集是模型訓(xùn)練過程中使用的圖像數(shù)據(jù)[10]。模型效果與訓(xùn)練集的選擇有直接關(guān)系。訓(xùn)練集對(duì)初始模型擬合與參數(shù)優(yōu)化有重要作用。本文的訓(xùn)練集分為單類垃圾數(shù)據(jù)和多類垃圾數(shù)據(jù)。

驗(yàn)證集被用于在每個(gè)Epoch完成之后進(jìn)行模型狀態(tài)和收斂情況的檢驗(yàn)，它不參與梯度下降的過程，僅負(fù)責(zé)調(diào)整迭代次數(shù)和學(xué)習(xí)率等超參數(shù)。此過程不參與學(xué)習(xí)參數(shù)的確定。

測(cè)試集用于對(duì)訓(xùn)練完成的最終模型進(jìn)行測(cè)試，通過測(cè)試來確定該模型檢測(cè)與分類的準(zhǔn)確率，確保該模型在未來實(shí)際應(yīng)用中的正確性和有效性。

2.3 垃圾識(shí)別分類方法

本文探討的垃圾檢測(cè)與分類方法使用基于YOLOv4的網(wǎng)絡(luò)作為實(shí)驗(yàn)訓(xùn)練模型，并將其用于全局和局部的特征提取。首先，YOLOv4網(wǎng)絡(luò)主體結(jié)構(gòu)如圖3所示，主要分為3部分：主干網(wǎng)絡(luò)、頸部網(wǎng)絡(luò)、頭部網(wǎng)絡(luò)。主干網(wǎng)絡(luò)采用CSPDarknet53，頸部網(wǎng)絡(luò)是SPP(Spatial Pyramid Pooling)模塊+PANet(Path Aggregation Network)，頭部網(wǎng)絡(luò)則采用YOLOv3。

圖3 YOLOv4網(wǎng)絡(luò)主體結(jié)構(gòu)

圖3中，CBM和CBL模塊是融合了批量歸一化并使用Mish和Leaky ReLU激活函數(shù)的卷積層；CSP模塊由多個(gè)殘差塊疊加而成；SPP模塊中的Concat操作則是將大小相同的特征圖進(jìn)行通道拼接。

為了在不增加訓(xùn)練成本的基礎(chǔ)上增加模型的準(zhǔn)確率，本文采用了數(shù)據(jù)增強(qiáng)和邊界框回歸損失函數(shù)。數(shù)據(jù)增強(qiáng)能增加輸入圖片的多樣性，從而增強(qiáng)復(fù)雜環(huán)境下模型的魯棒性[11]。傳統(tǒng)的目標(biāo)檢測(cè)方法使用的是均方誤差(Mean Squared Error，MSE)，與檢測(cè)框評(píng)測(cè)指標(biāo)IOU(Intersection over Union)存在誤差，故引入了 LossIou，其定義如式(1)所示。

LossIou=1-IoU

(1)

注意力機(jī)制模塊可以根據(jù)全局圖像，從大量信息中篩選有用信息，抑制其它無用信息，增強(qiáng)特征圖中重要的空間和通道特征，使網(wǎng)絡(luò)在訓(xùn)練過程中可抓住目標(biāo)特征重點(diǎn)學(xué)習(xí)。本文選用卷積注意力機(jī)制模塊(Convolutional Block Attention Module，CBAM)，并將其嵌入到Y(jié)OLOv4模型中。如圖4所示，本文選擇在每個(gè)檢測(cè)頭前嵌入一個(gè)CBAM模塊。

圖4 CBAM嵌入示意圖

在CBAM模塊中[6]，對(duì)于一個(gè)中間層的特征圖F∈RC×H×W，C表示特征圖的通道，H表示特征圖的高度，W表示特征圖寬度。一維的通道注意力特征圖MC∈RC×1×1，二維的空間注意力特征圖MS∈R1×H×W，整個(gè)過程為

F1=MC(F)?F，F(xiàn)2=MS(F1)?F1

(2)

式中，?表示逐元素相乘。先將通道注意力特征圖與輸入的特征圖相乘得到F1；然后計(jì)算F1的空間注意力特征圖；再將兩者相乘得到輸出F2。

通道注意力模塊以及空間注意力模塊的計(jì)算過程如下文所述。CBAM通道注意力模塊使用最大池化和平均池化對(duì)特征圖進(jìn)行空間維度上的壓縮，得到兩個(gè)不同的通道背景描述，即Fmax和Favg。再通過由多層感知器(Multilayer Percep-tron，MLP)組成的共享網(wǎng)絡(luò)對(duì)Fmax和Favg進(jìn)行計(jì)算[12]，得到通道注意力特征圖 ，如式(3)所示。

MC(F)=σ(MLP(AvgPool(F))+MLP(MaxPool(F)))=

σ(W1(W0(Favg))+W1(W0(Favg)))∈RC×1×1

(3)

式中，σ表示Sigmoid激活函數(shù)；MLP表示共享全連接層[12]；W0表示共享全連接層的第1層；W1表示共享全連接層的第2層；輸出向量長度為C。

通道注意力模塊(Channel Attention Module)如圖5所示。

圖5 CBAM通道注意力模塊

CBAM空間注意力模塊先在通道維度上進(jìn)行最大池化和平均池化，得到兩個(gè)不同的空間背景描述，再使用卷積操作[13]生成空間注意力特征圖

MS(F)=σ(f7×7([AvgPool(F)；MaxPool(F)])) =

σ(f7×7([Favg；Fmax]))∈RH×W

(4)

式中，f7×7表示7×7的卷積層。

圖6所示為空間注意力模塊示意圖(Spatial Attention Module)。

圖6 CBAM空間注意力模塊

相比于FPN(Feature Pyramid Networks)，PANet的效果更好，但計(jì)算量也更大。本文中YOLOv4特征集成采用PANet，故需先將PANet進(jìn)行簡(jiǎn)化，將PANet的addition改成concatenation。PANet主要包含F(xiàn)PN、Bottom-up path augmentation、Adaptive feature pooling、Fully-connected fusion共4個(gè)部分[14]。FPN通過融合高低層特征來提升目標(biāo)檢測(cè)的效果，其對(duì)小尺寸目標(biāo)檢測(cè)效果的提升更為顯著。淺層特征一般是邊緣形狀等特征，考慮到網(wǎng)絡(luò)淺層特征信息對(duì)于實(shí)例分割較為重要，故引入Bottom-up path augmentation。Adaptive feature pooling用每個(gè)感興趣區(qū)域提取不同層的特征來做融合。Fully-connected fusion則向原有分割支路引入一個(gè)前背景二分類的全連接支路[15]，通過融合這兩條支路的輸出來得到更精確的分割結(jié)果。

本文選用Mish函數(shù)作為激活函數(shù)，如圖7所示，其正值無上界的特點(diǎn)避免了由于封頂而導(dǎo)致的飽和。另外，Mish函數(shù)保證了每一點(diǎn)的平滑性，從而獲得了較好的梯度下降效果[16]。

圖7 Mish函數(shù)示意圖

Mish(x)=x×tanh(In(1+ex))

(5)

2.4 數(shù)據(jù)集訓(xùn)練部分

本次實(shí)驗(yàn)平臺(tái)基于Ubuntu16.04系統(tǒng)，顯卡為TITAN RTX 2080Ti，內(nèi)存為32 GB。訓(xùn)練部分基于YOLOv4-custom.cfg對(duì)數(shù)據(jù)集進(jìn)行模型訓(xùn)練，其中batch值為64，subdivision值為16。因本實(shí)驗(yàn)數(shù)據(jù)集共有9類，所以max_batches值被設(shè)置為18 000次。Steps=xxx，本文將xxx設(shè)置成max_batches的80%和90%。YOLOv4有3個(gè)YOLO層，將每個(gè)層里面的classes值設(shè)為數(shù)據(jù)集類型數(shù)目9。數(shù)據(jù)集按照二八原則進(jìn)行劃分，即20%為驗(yàn)證集，剩余80%作為訓(xùn)練集，并編寫Python腳本文件來分離訓(xùn)練集和測(cè)試集，在Main中生成train.txt和val.txt文件。訓(xùn)練時(shí)打開Trainning的batch和subdivisions，將Testing部分注釋掉。本實(shí)驗(yàn)基于YOLO預(yù)訓(xùn)練模型開始訓(xùn)練，當(dāng)訓(xùn)練的loss趨于平穩(wěn)時(shí)，停止訓(xùn)練。使用訓(xùn)練得到的模型對(duì)垃圾進(jìn)行檢測(cè)，對(duì)易誤檢的垃圾進(jìn)行數(shù)據(jù)補(bǔ)充，從而繼續(xù)提升模型精度。

3 實(shí)驗(yàn)結(jié)果與分析

針對(duì)本文提出的模型調(diào)整與算法改進(jìn)措施，分別進(jìn)行對(duì)比實(shí)驗(yàn)來檢測(cè)性能。目前平均精度(Average Precision，AP)是多類別目標(biāo)檢測(cè)及其分類中最重要的評(píng)價(jià)指標(biāo)，因此本次實(shí)驗(yàn)采用AP作為垃圾檢測(cè)與分類的評(píng)價(jià)指標(biāo)[17]。各類別垃圾在YOLOv4以及改進(jìn)的模型上的精確度表現(xiàn)如表1所示。

表1 不同垃圾類別的AP比較

由表1可知，采用本文提出的算法模型后，有8類垃圾的檢測(cè)精確度高于YOLOv4的精確度，其中對(duì)PlasticBottle的檢測(cè)精確度提升最高，提升了6.36%；識(shí)別效果最好的類別則為Cans，識(shí)別精度可達(dá)94.56%。本實(shí)驗(yàn)還對(duì)比了CSPResNEXt50和CSPDarknet53在垃圾檢測(cè)與分類方面的表現(xiàn)，具體結(jié)果如表2所示。

表2 CSPResNEXt53與CSPDarket53的參數(shù)對(duì)比

由表2可知，CSPResNeXt50僅包含16個(gè)卷積層，其感受野為425×425，包含20.6 MB參數(shù)； CSPDarkNet53包含29個(gè)卷積層，725×725的感受野，共有27.6 MB參數(shù)。因此，雖然CSPResNEXt50在分類方面的表現(xiàn)優(yōu)于CSPDarkNet53，但其在檢測(cè)方面的表現(xiàn)則不如CSPDarkNet53[18]。

此外，本實(shí)驗(yàn)還對(duì)比了使用K-means聚類算法優(yōu)化Anchor box后和未經(jīng)優(yōu)化的YOLOv4以及本文引入注意力機(jī)制和K-means的YOLOv4模型在本文數(shù)據(jù)集上進(jìn)行訓(xùn)練后的性能變化，具體結(jié)果如表3所示。

表3 不同模型的效能評(píng)估

由表3可知，當(dāng)IOU閾值設(shè)置為0.5時(shí)，相較于原始YOLOv4，使用K-means聚類算法的YOLOv4-1的map提升了1.08%，使用K-means聚類算法并引入注意力模塊的模型的map提升最為明顯，提升了2.81%。改進(jìn)YOLO模型對(duì)垃圾的識(shí)別結(jié)果如圖8所示。

圖8 改進(jìn)后的算法識(shí)別效果

4 結(jié)束語

本文針對(duì)現(xiàn)實(shí)生活中人工分揀垃圾效率低、任務(wù)重且環(huán)境差等問題，基于卷積神經(jīng)網(wǎng)絡(luò)的模型進(jìn)行改進(jìn)，提出了進(jìn)行垃圾分類與檢測(cè)的新方法。該方法能夠在一定程度上改善垃圾分類耗時(shí)耗力的問題，并準(zhǔn)確檢測(cè)垃圾，完成分類。本文在YOLO模型中融入了K-means聚類算法，并向該模型中加入CBAM注意力機(jī)制，研究了注意力機(jī)制對(duì)垃圾檢測(cè)算法性能的影響，并選擇經(jīng)過優(yōu)化的網(wǎng)絡(luò)架構(gòu)來完成垃圾檢測(cè)與分類任務(wù)。下一步將對(duì)模型進(jìn)行壓縮，嘗試構(gòu)造更好的輕量級(jí)網(wǎng)絡(luò)結(jié)構(gòu)并將其用于模型訓(xùn)練，從而提升網(wǎng)絡(luò)模型的實(shí)時(shí)性。