基于改進YOLO v5的礦山石塊實例分割算法

曹士杰,張竹林

(山東交通學院 汽車工程學院,山東 濟南 250023)

礦山工程建設面臨復雜的地質與環境條件,存在較強的安全風險。以云計算、人工智能算法等共性支撐技術為基礎的數字孿生技術的快速發展及應用,使數據的實時獲取、瞬時傳輸、快速計算、智能決策等成為了現實,為實現安全、綠色、高效、智能發展與賦能零碳礦山的建設提供了有力的技術支撐[1]。作為液壓挖掘機的一個重要作業工具,對于從事開礦、二次破碎等工作發揮著重要作用,而完成對破碎目標的準確識別是實現液壓破碎錘挖掘機無人駕駛的關鍵一步。

由于礦山石塊大多為開山作業形成,其結構特征通常較為復雜。而且礦山環境存在目標堆積、遮擋等問題,目標的漏檢、誤檢很容易發生。研究人員提出了一系列目標檢測改進的方法,楊麗榮等[2]結合混合注意力模塊改進Mask RCNN的分割網絡結構,減少卷積運算造成的礦石細節信息丟失問題,提高了分割精度。司晨冉等[3]提出了一種基于Mask R-CNN和分水嶺算法的巖石顆粒圖像分割方法,首先利用改進的CNN方法防止大塊巖石出現過分割。湯文聰等[4]提出了一種基于FCM-WA聯合算法的礦石圖像分割方法,將模糊C均值聚類(FCM)算法與分水嶺(WA)算法相結合,對FCM算法輸出的礦石圖像邊緣粘連部分進行分割,獲取最佳的分割圖像。

以上方法雖然保障了檢測精度,但檢測速度有所犧牲,且計算量較大,操作較為復雜。本文通過采集到礦山石塊相關圖片信息,考慮圖像處理方法的通用性,研究利用機器視覺技術識別石塊的方法。為改善目標漏檢、誤檢問題,并保證模型的檢測精度與輕量化,提出了基于YOLO v5的實例分割模型。

1 改進的YOLO v5算法

YOLO v5的7.0版本提供了5種實例分割模型,綜合考慮各模型的檢測速度、模型大小及檢測精度,本文選擇YOLO v5l-seg作為進一步優化的模型。

1.1 SA注意力機制

注意力模塊是網絡設計的重要組成部分,為關注有用的特征信息,提升模型性能,可以引入ECA[5]、CBAM[6]、CA[7]等注意力模塊。ECA注意力機制通過不降維的局部跨信道交互策略,使用了1-D卷積來生成通道間注意力機制,降低模型的復雜度,但其準確度無法保證;CBAM注意力機制主要聚焦于空間信息的聚合,但效果未達到最優。本文通過引入Shuffle-Attention(SA)[8-9]注意力機制提升模型效果。SA注意力模塊結合了組卷積,空間注意力機制和通道注意力機制。不僅能取不同通道間的信息,還能降低計算量,有助于實現目標的精確定位和識別。SA注意力模塊將輸入序列打亂再重新排序,并計算每個位置的重要性,得到注意力權重,與傳統的注意力機制相比,在計算效率提高的同時,提高模型的泛化能力,使模型在訓練數據和測試數據上的表現更加一致。SA注意力模塊如圖1所示。

圖1 SA注意力模塊的整體結構

從圖1給出的SA注意力模塊整體結構示意可以看到,它首先采用Feature Grouping對輸入尺寸為X∈RC×H×W的特征圖X,沿著通道維度劃分為g個組:X=[X1,…,XG],RC/G×H×W,特征X會被拆分成2個分支:Xk1,Xk2∈RC/2G×H×W,分別用于學習通道注意力特征和空間注意力特征。對每組特征通過SA Unit模塊Concat的方式將組內信息進行融合,生成不同的重要性系數。SA內部綠色部分為通道注意力機制,采用GAP+Scale+Sigmoid的組合,該過程可以描述如下。

(1)

(2)

SA內部藍色部分為空間注意力機制,采用Group Norm(GN)對Xk2進行處理得到空域層面的統計信息,然后采用Fc(·)進行增強,該過程可以描述如下。

(3)

1.2 ASPP模塊

圖片中的相鄰像素點一般都存在信息的冗余問題,傳統的下采樣雖可增加感受野,但會降低空間分辨率。為了獲得更大的感受野,減少信息的冗余問題,實現礦山石塊的精確分割,選擇將主干網絡中的SPPF模塊替換為Atrous Spatial Pyramid Pooling(ASPP)[10]模塊。

ASPP模塊使用并行空洞卷積層進行采樣,獲得不同采樣率的提取特征,處理后的分支特征融合后,最終結果在保證分辨率的同時可以獲得更大的感受野和多尺度的物體信息。通過賦值不同的膨脹因子實現不同的感受野,進而獲取多尺度信息,具體實現如圖2所示。

圖2 ASPP模塊結構

ASPP模塊由一個1×1的卷積(最左側綠色)、池化金字塔和ASPP Pooling 3個部分組成。池化金字塔各層分別為卷積核大小為3×3的3個卷積層,每個卷積層的padding和dilation分別為6、12、18?？斩淳矸e[11]層通過控制padding以及dilation 2個參數控制膨脹率,獲取不同尺度的感受野來實現多尺度的特征信息提取。ASPP Polling 部分首先是一個指定最后輸出尺寸為1×1的自適應均值池化(AdaptiveAvgPool2d)層,通過將各通道的特征圖分別壓縮至1×1,從而提取各通道的特征,獲取全局的特征。然后是一個1×1的卷積層,將獲取的特征進一步提取并降維。通過Upsample利用傳統插值方法進行上采樣還原成原始輸入大小,出層后對ASPP五層輸出進行Concat融合,對疊加后的輸出用1×1卷積層卷積操作降維至給定通道數,得到最終輸出結果。Concat融合過程表述如式(4)所示。

n=k+(k-1)×(d-1),

(4)

式中:n為現在的卷積核大小;k為原卷積核大小;d為超參數;d-1的值為加入的空格數。令o為空洞卷積后的特征圖大小,i為輸入空洞卷積的大小,s為步長,則有

(5)

全局平均池化后上采樣的輸出等于每個空洞卷積的輸出,可以實現在通道維度的Concat融合。

1.3 加權雙向特征融合模塊設計

YOLO v5的實例分割算法主干特征提取網絡為CNN網絡,頸部層采用PANet結構,由于存在冗余信息的融合,計算量過大,通過用改進后的加權雙向特征金字塔[12]網絡(BiFPN)替換原PAN結構,并將multi-head self-attention(BoT)[13]處理框架融合到網絡頸部中,形成BIFPN+Transformer(BBIFPN)架構,充分發揮兩者優點,提高對密集目標的檢測效果。

BoT使用sgd+momentum構成的優化器進行訓練,通過Batch Normalization模塊完成歸一化操作,并在FPN block模塊中作用了3個非線性激活。改善基線的同時減少了參數量,從而實現了延遲最小化。圖3(a)是傳統的Transformer[14]模型,圖3(b)是BoTNet模型,圖3(c)是BoTblock模型。

圖3 BoT模塊

本文特征融合結構改進點如下:

(1) 融合Bottleneck Transformers處理框架。在實例分割任務中納入自注意力,用全局自注意力替換空間卷積,減少參數,從而使延遲最小化,更有利于實現密集預測任務。

(2) 刪除冗余節點。特征融合后的節點,對于特征網絡的貢獻相對更大,因此保留網絡中進行特征融合的節點,刪除單獨輸出的節點,有效減少模型計算量。對相同尺度的輸入節點與輸出節點之間添加跳躍連接,如圖6中間虛線所示,在不增加太多計算成本的同時,融合了更多的特征。

(3) 加權特征融合。引入可學習的權重,使網絡更關注關鍵特征學習。通過Fast normalized fusion集成雙向跨尺度連接和快速歸一化融合。具體計算過程可以表示為

(6)

式中:Ii表示輸入特征圖;wi為權重系數,使用Relu激活函數保證該值始終大于0;ε為學習率,值為10-4。

以level4的融合特征為例,計算過程如式(7)～(8)。

(7)

(8)

圖4 特征融合模型

1.4 網絡輕量化

在引入BBIFPN架構后,檢測精度有了很好的改善,為了進一步減少計算量,提高推理檢測速度,降低模型體積,采用輕量化模塊優化網絡。YOLO v5特征提取網絡原始C3結構,參數量與模型過于龐大,容易產生檢測速度、響應速度慢,內存不足的情況。本文通過引入Ghost[15]網絡,減少特征圖中的冗余來實現高效的神經網絡。

Ghost模塊由3步操作完成,假設原始卷積操作中特征圖輸入數據為h×w×c,卷積核維度為n,尺寸為k×k,卷積后輸出的本征圖(intrinsic)數據為h′×w′×m。

(1)對于輸出通道較少的特征圖使用m組k×k的Kernel與input進行原始卷積。

(2)其他的ghost特征圖通過廉價操作生成,即將本征圖的每個通道單獨進行線性變換產生。

(3)將前兩步得到的特征intrinsic和ghost進行identity連接得到最終結果。

假設在Ghost模塊的線性操作中有m個恒等映射和n-m個線性運算,且每個線性運算的平均核大小等于d×d,則用Ghost模塊替換普通卷積的理論加速比(speed up ratio)為

(9)

理論壓縮比(compression ratio)為

(10)

由此可知,通過使用Ghost模塊,理論上可以獲得s倍的加速比和壓縮比,有效減少模型的體積,提高模型的運算速度。本文構建了Ghostconv、GhostBottleNeck和GhostC3模塊,如圖5所示。

圖5 Ghost系列模型

1.5 Rock-YOLO v5模型整體結構

Rock-YOLO v5模型從以下4個方面對模型進行了優化:

(1)融合Shuffle-Attention(SA)注意力機制。

(2)優化主干網絡,將SPPF模塊替換為ASPP模塊。

(3)設計BBIFPN框架優化路徑聚合網絡(Path Aggregation Network,PAN)。

(4)網絡輕量化。

最終建立的模型結構如圖6所示。

圖6 Rock_YOLO v5模型結構

2 試驗結果分析

2.1 數據集及試驗環境介紹

本文所有試驗均在Windows10系統上進行,CPU為Intel(R) Core(TM) i7-10700,GPU為NVIDIA GeForce RTX 3070,顯存為8 GB,采用的深度學習框架為Pytorch1.7.1,參數環境CUDA版本為11.0,編譯語言為Python-3.9.13,訓練批尺寸(batch_size)設置為4;迭代次數(epoch)為100;輸入圖片大小為640×640。試驗所使用的數據集為實地采集的石塊數據集,共有4 556張高清圖片。石塊數據集按照7∶1∶2的比例隨機劃分訓練集、測試集、驗證集3個部分,分別為訓練集3 188張,測試集456張,驗證集912張。

2.2 結果分析

為驗證引入注意力機制、ASPP模塊、Ghost模塊和BBIFPN架構后對實例分割任務產生的影響,本文通過消融試驗進行測試,評估指標包括參數量(params)、計算量(FLOPs)、mAP50等。mAP50為IOU>5時的mAP。

對于三類損失曲線:定位損失(box_loss)、分割損失(seg_loss)、置信度損失(obj_loss),改進模型收斂速度更快,訓練損失圖像如圖7所示。為驗證不同注意力機制對模型性能提升的影響,進行對照試驗,對比加入不同注意力機制后對網絡性能的提升效果,結果如表1所示。

表1 注意力機制性能對照

圖7 訓練損失圖像

由試驗結果可知,各注意力機制均對分割精度有提升效果,添加SA和CA注意力機制后,準確率提升最高,達到86.4%,但SA注意力機制的計算量與參數更少,因此選擇SA注意力機制對模型進行優化。

進行消融試驗驗證改進前后相同條件下對算法性能的影響,結果如表2所示。

表2 消融試驗

改進后的模型獲得了mAP的提升,雖然FPS值有所下降,但還是可以滿足無人駕駛破碎錘視覺實時性的要求。引入SA模塊后,mAP提升了1.4%;引入ASPP模塊時,mAP提升了2.8%;引入BBIFPN架構時,mAP提升了2.2%,證明本文提出的特征融合結構對于實例分割任務有正面作用,通過引入Ghost模塊極大的減少了網絡復雜度。使用GhostConv系列模塊替換掉原始模塊后,參數量減少了42.4%,保障了模型的輕量化。

設置對照試驗,驗證對于密集目標實例分割任務的效果。圖8為YOLO v5原始模型與改進后模型對于密集目標任務的檢測對比圖,由圖8可知,引入空洞卷積、SA注意力機制和加權跨層特征融合模塊后,對目標關注度提升更大,對位置信息的捕捉更為精確,從而提高模型精度。當石塊分布為密集時,原算法出現了漏檢的情況,一些石塊沒有被分割出來,而Rock_YOLO v5模型檢測結果有不錯的提高,漏檢的情況有所改善,驗證了改進算法具有更高的檢測精度。

圖8 密集目標檢測對比效果

3 結語

本文建立了Rock_YOLO v5實例分割模型,結合通道注意力機制與空間注意力機制,對原始算法進行改進,并在礦山石塊數據集上進行了試驗驗證。針對密集目標實例分割任務的漏檢、誤檢問題,通過引入SA注意力機制、ASPP模塊和改進的BBIFPN模塊對實例分割網絡進行優化。SA注意力機制和ASPP模塊使網絡更有利于關注位置信息和獲取感受野,更有利于網絡對密集目標物的定位和識別,改進后的BBIFPN模塊通過多尺度特征信息和融合,使模型提取到更豐富的特征,有效提高了模型的準確率。引入Ghost模塊后,減輕了模型的計算量和參數,降低了顯存和要求。試驗結果表明,相比YOLO v5l-seg,本文提出的改進算法具有更好的準確性和實時性,可以完成無人駕駛破碎錘識別目標石塊的任務。本文改進后的模型在解決行業難題,特別是處理遮擋和密集問題方面具有參考價值。但是也存在訓練和推理時間增加的問題,可以從邊緣平滑和分割體驗優化方面,進一步提升Rock_YOLO v5模型的性能,更有效地分配和利用資源。