基于卷積神經網絡的垃圾圖像分類算法①

董子源,韓衛光

1(中國科學院大學,北京 100049)

2(中國科學院 沈陽計算技術研究所,沈陽 110168)

1 引言

垃圾回收利用作為發展循環經濟的必經之路,是根治污染、提高環境治理效果的關鍵所在.隨著我國生產力水平的發展,生活垃圾、工業垃圾數量不斷增加,困擾著很多城市.據統計,僅2018年,中國垃圾的清運量已經達到了2.28 億噸[1].在人們將垃圾投放進垃圾箱之后,垃圾被運送到垃圾處理廠統一處理.當前國內的垃圾處理廠,更多依靠人工在流水線上作業去分揀垃圾,對作業者健康不利且分揀效率較低,已不能滿足大量垃圾處理需求.此外,人工分揀的垃圾種類極為有限,大部分垃圾無法重新回收利用,造成很大浪費.隨著深度學習技術的發展,卷積神經網絡使圖像分類算法在精度和速度上得到了巨大的提升,讓我們看到了借助視覺技術自動分揀垃圾的可能性.通過攝像頭拍攝垃圾圖片,利用卷積神經網絡檢測出垃圾的類別,之后就可以借助機械手或推板自動完成分揀任務,可以降低人工成本,提高分揀效率.因此,開展垃圾圖像分類算法的研究,具有重要的應用價值.

2 相關工作

早期,學者們只能借助經典的圖像分類算法[2–5]完成垃圾圖像分類任務,這要通過手動提取的圖像特征并結合相應的分類器完成.吳健等[6]利用顏色和紋理特征,初步完成了廢物垃圾識別.由于不同數據集的圖像背景、尺寸、質量不盡相同,傳統算法需要根據相應數據人工提取不同的特征,算法的魯棒性較差,并且處理方式復雜,所需時間較長,無法達到實時的效果.隨著卷積神經網絡(Convolution Neural Network,CNN)的飛速發展,深度學習廣泛應用于圖像識別領域.作為數據驅動的算法,CNN 具有強大的特征擬合能力,可以有效、自動地提取圖像特征,并具有較快的運行速度.2012年,AlexNet[7]取得了ImageNet 圖像分類競賽的冠軍,標志著深度學習的崛起.隨后幾年,GoogleNet[8]、VGGNet[9]、ResNet[10]等算法提升了圖像分類的精度,并成功應用于人臉識別、車輛檢測等多個領域.垃圾圖像分類,在深度學習算法的幫助下同樣取得了較大的突破.斯坦福大學的Yang 等建立了TrashNet Dataset 公開數據集,包含6 個類別,共計2527 張圖片.Ozkaya 等[11]通過對比不同CNN 網絡的分類能力,搭建神經網絡(本文稱之為TrashNet)并進行參數微調,在數據集TrashNet Dataset 上取得了97.86%的準確率,是目前這一數據集上最佳分類網絡.在非公開數據集方面,Mittal 等[12]自制了2561 張的垃圾圖片數據集GINI,使用GarbNet 模型,得到了87.69%的準確率.國內方面,鄭海龍等[13]用SVM 方法進行了建筑垃圾分類方面的研究.向偉等[14]使用分類網絡CaffeNet,調整卷積核尺寸和網絡深度,使其適用于水面垃圾分類,在其自制的1500 張圖片數據集上取得了95.75%的識別率.2019年,華為舉辦垃圾圖像分類競賽,構建了樣本容量為一萬余張的數據集,進一步促進了該領域的發展.

我國各地區生活垃圾分類標準有所不同,大致可分為可回收垃圾、有害垃圾、廚余垃圾和其他垃圾這4 大類,且每個類別下又包含若干子類別,種類繁多且十分復雜.按照這樣的分類標準做的垃圾圖像識別研究,國內目前還處于起步階段.現有的圖形分類算法在垃圾處理領域的應用較少,且存在準確率不足、泛化性能差、處理效率低的缺點.針對現有方法的不足,本文提出一種基于卷積神經網絡的垃圾圖像分類算法(Garbage Classification Net,GCNet),在網絡結構中融合了注意力機制模塊與特征融合模塊,提高了模型在垃圾分類任務上的準確性與魯棒性.

3 算法設計

3.1 模型結構

本文構建的GCNet 模型包括特征提取器、分類器兩部分,整體結構如圖1所示.圖中特征提取器由Resnet101 作為主干部分,共包括5 個bottleneck,并在不同的bottleneck 后加入注意力機制模塊,同時對不同模塊提取到的特征進行特征融合(如圖1中虛線所示)以從輸入x中提取圖像的特征信息F1:

其中,Me表示特征提取器.

分類器由兩層全連接層和一個Softmax 分類器組成,對提取到的特征信息F1進行分類,以得到圖像在每個類別下的最終得分yi:

其中,Mc表示分類器.

3.2 注意力機制

注意力機制源于對人類視覺的研究,人類會根據實際需要選擇視網膜內特定區域進行集中關注,可以將有限的處理資源分配至重要的部分.由于相同類別垃圾的特征表征差異性可能較大,不利于圖片的正確分類,這就要求準確地關注圖像中的顯著區域.受這一思想的啟發,通過構建注意力機制模塊,使網絡模型重點關注有利于分類的特征區域,以實現更好的特征提取功能,其具體結構如圖2所示.

圖1 GCNet 網絡結構圖

圖2 注意力機制示意圖

將每個bottleneck 提取到的特征以Fi表示,其大小為N×C×W×H.利用格拉姆矩陣[15]構建局部注意力機制模塊,將Fi與其轉置相乘,獲得大小為N×C×C的局部特征:

該操作可以獲得特征圖中每個元素自身間的關聯,使原本數值較大的元素越大、數據較小的元素越小,即凸顯每個局部的重要程度,使有助于判斷類別的特征更加明顯,同時抑制影響判斷的特征.

隨后,對Fi進行全局平均池化操作(Global Average Pooling,GAP)[16],可以保留特征Fi提取到的空間信息和語義信息,并獲得大小為N×C×1×1的全局特征:

代 表了特征Fi中最顯著的特征,這也符合所定義的全局注意力機制.

最后,將局部特征與全局特征相乘,從而獲得包含局部和全局信息的總體特征:

3.3 特征融合機制

垃圾分類任務中,從屬于同一類別的垃圾往往差異性較大,如可回收物中包括的紙板箱和玻璃杯,二者存在較大的外觀差異但卻屬于同一類別,這會給分類器的判別增加一定的難度.此外,隨著網絡層數的加深,單一的圖像特征會丟失某些區域的信息,從而導致模型的分類性能變差.針對這一問題,常用的方法是通過不同卷積核的池化操作構建多尺度的特征融合模塊,但這種操作會增大模型的運算量,同時也會使模型訓練更加困難.

由于構建的特征提取器中包含了不同的池化操作,因此只需要提取不同bottleneck 下的特征便能獲得不同尺寸的特征.在此基礎上,本文提出了改進的特征融合機制,如圖1中的虛線所示,將每個bottleneck 經過注意力機制后所提取到的總體特征進行融合:

其中,wc為1×1的卷積核,bc為偏置,Concat表示融合操作.該操作旨在利用不同尺度的特征信息,避免信息的丟失,進一步提升模型的魯棒性.

4 實驗及結果分析

4.1 實驗平臺

本文實驗在Ubuntu16.04 系統下,利用Python 語言,Pytorch 深度學習框架完成.硬件環境為CPU Intel I7-9700K,內存32 GB,顯卡Nvidia GeForce RTX 2080Ti.

4.2 實驗數據

本文采用華為垃圾分類挑戰杯數據集,該數據集已全部標注類別,包括廚余垃圾、可回收物、其他垃圾和有害垃圾4 個大類,每個大類下又包含若干子類別,合計40 個小類別,共計14 683 張圖片.

在實驗中,采用4:1的比例劃分數據集,80%作為訓練數據,20%作為測試數據.此外,為了增強模型的泛化能力和魯棒性,還對訓練樣本進行了數據增強操作,包括隨機旋轉、隨機翻折、隨機裁剪等.

4.3 結果分析

實驗選用ADAM[17]優化算法訓練模型,動量系數為0.9,共設置50 個迭代周期,初始學習率設置為0.01,每隔10 個迭代周期學習率衰減0.1 倍,一階矩估計的指數衰減率為0.99,二階矩估計的指數衰減率為0.999.此外,使用交叉熵損失函數來訓練優化模型.

在訓練過程中,對GCNet 模型通過消融實驗以分別驗證注意力機制和特征融合機制的功能.實驗1(Experiment_1)為含有注意力機制和特征融合機制的模型,實驗2(Experiment_2)為只含有注意力機制的模型,實驗3(Experiment_3)為只含有特征融合機制的模型,實驗4(Experiment_4)為不含有注意力機制和特征融合機制的模型.各個實驗的訓練過程迭代曲線如圖3所示.

圖3 GCNet 訓練過程曲線

從圖3中可以看到,實驗1的準確率提升與損失值下降都快于其他3 個實驗,說明訓練過程收斂更快.實驗2和實驗3的速度幾乎相近,而實驗4的速度最慢.為了進一步驗證上述訓練的準確性,在對應的測試集上進行測試,實驗結果如表1所示.包含注意機制模塊和特征融合模塊的實驗1 取得了96.73%的最優準確率,同時在各個子類別上均取得了最高的準確率,說明本文構建的模型泛化能力較好.在類內差異性較大的“其他垃圾”這一類別中,注意力機制和特征融合機制均能夠顯著提高模型的準確率.

表1 GCNet 模型消融實驗準確率對比(單位:%)

本文模型與其他模型的準確率對比結果如表2所示.可以看出,本文構建的GCNet 平均準確率最高,在各個類別中也取得了最高的準確率,說明本文構建的注意力機制和特征融合機制充分提取了有利于圖像分類的特征,使得分類結果更加準確.TrashNet的平均準確率略差于GCNet,而CaffeNet 結果最差.

表2 不同模型實驗準確率對比(單位:%)

圖4為3 種算法GCNet,TrashNet和CaffeNet 對同樣4 張圖片的測試結果,顯示了預測類別以及屬于這一類別的概率.GCNet 在各個類別的樣例中,也取得了最好的結果.特別是在第4 幅(其他垃圾)的測試中,由于該類別的圖像類內差異過大,CaffeNet 將其識別為其他垃圾的概率僅為62.75%,將其錯誤地識別成有害垃圾的概率卻高達35.61%,在某些干擾情況下,這可能會導致判定錯誤,對分類結果造成較大影響.本文算法成功的改進了這一問題,確保了在較難分辨的類別上的準確率.

圖4 各模型測試結果對比圖

5 結論

本文針對垃圾圖像分類問題,構建了一種基于卷積神經網絡的算法GCNet,該網絡通過構建注意力機制和特征融合機制,能夠有效地提取圖像特征、降低類別差異性帶來的影響,并在相關數據集上取得了96.73%的平均準確率,相較于現有的分類算法提升了約4%的準確率,滿足了實際的應用需求,具有良好的應用前景.