基于YOLOv5的鋼材數(shù)量識別方法

劉洪寧，李容露

（大連民族大學(xué)機(jī)電工程學(xué)院，遼寧 大連 116600）

在工業(yè)建筑過程中，不可避免會(huì)消耗大量的工業(yè)建材。傳統(tǒng)工業(yè)建材計(jì)數(shù)工作一直由人工完成，手工計(jì)數(shù)方法易受到統(tǒng)計(jì)人員的注意力、反應(yīng)速度、疲勞等因素的影響，無法滿足現(xiàn)代化鋼鐵企業(yè)快速生產(chǎn)需求。近年來，計(jì)算機(jī)視覺技術(shù)蓬勃發(fā)展，在圖像處理，目標(biāo)檢測領(lǐng)域取得了豐碩成果。卷積神經(jīng)網(wǎng)絡(luò)CNN 可以提取到圖片數(shù)據(jù)集更好層次的特征信息，使得視覺處理模型獲得更高的準(zhǔn)確率，現(xiàn)已被廣泛應(yīng)用于各類工業(yè)任務(wù)中。在鋼材計(jì)數(shù)這一圖像處理任務(wù)中，傳統(tǒng)圖像處理方法首先手動(dòng)提取鋼材特征，再進(jìn)行比對計(jì)數(shù)，這就需要根據(jù)經(jīng)驗(yàn)手動(dòng)設(shè)置參數(shù)。鋼材圖像數(shù)據(jù)具有數(shù)據(jù)密集、成像背景噪聲復(fù)雜、目標(biāo)尺寸較小等多種特性。傳統(tǒng)方法不具備較好的泛化能力，難以獲得穩(wěn)定準(zhǔn)確的檢測效果?；诰矸e網(wǎng)絡(luò)的深度學(xué)習(xí)模型具有更好的準(zhǔn)確性與魯棒性。采用深度學(xué)習(xí)的方法訓(xùn)練目標(biāo)檢測模型應(yīng)用于鋼材計(jì)數(shù)任務(wù)中，不但快速準(zhǔn)確地完成對工業(yè)建材的清點(diǎn)工作，更能將建筑工人從枯燥繁重工作中解脫出來。本文擬采用深度學(xué)習(xí)的方法，以YOLOv5 為基礎(chǔ)模型，構(gòu)建卷積神經(jīng)網(wǎng)絡(luò)對鋼材圖像進(jìn)行目標(biāo)檢測，完成鋼筋計(jì)數(shù)任務(wù)。

1 基于YOLOv5 的鋼材數(shù)量識別算法

1.1 目標(biāo)檢測模型搭建

YOLO（You Only Look Once）是由Joseph Redmon 編寫的one-stage 目標(biāo)檢測算法，現(xiàn)已更新到第五代版本即YOLOv5。在YOLO 模型出來之前，先前的目標(biāo)檢測系統(tǒng)大多使用分類器對測試圖像的不同切片進(jìn)行評估，即滑動(dòng)窗口這一十分經(jīng)典的模型。DPM（Deformable Part Model）等一些早期目標(biāo)檢測算法模型都采用了滑動(dòng)窗口的方式，并在這些窗口區(qū)域上運(yùn)行分類器。R-CNN（Region-based Convolutional Neural Networks）等采用候選區(qū)方法配合CNN 模型進(jìn)行預(yù)測，最后得到1 個(gè)特征向量進(jìn)行分類。無論是滑動(dòng)窗口還是候選區(qū)模型，其本質(zhì)上是一種圖像分割算法。YOLO之前的系統(tǒng)一直被速度慢、計(jì)算量大，優(yōu)化困難等問題困擾。YOLO 算法與傳統(tǒng)兩階段的目標(biāo)檢測算法不同，其單階段的實(shí)現(xiàn)端到端的檢測任務(wù)。YOLO 將對象檢測重新定義為一個(gè)回歸問題，及端到端的直接預(yù)測，而不再采取傳統(tǒng)的圖像分割累加預(yù)測的形式。YOLO 的推理速度是傳統(tǒng)R-CNN 模型的推理速度的百倍以上。YOLOv5 是目前YOLO 系列最新的實(shí)時(shí)目標(biāo)檢測算法，YOLOv5 檢測精度的提升使得網(wǎng)絡(luò)在目標(biāo)檢測方面更加具有實(shí)用意義。YOLO 算法將單個(gè)卷積神經(jīng)網(wǎng)絡(luò)CNN 應(yīng)用于整個(gè)圖像，將圖像分成多個(gè)網(wǎng)格，選取每個(gè)網(wǎng)格作為預(yù)測中心，計(jì)算落在該網(wǎng)格內(nèi)的目標(biāo)的信息，每個(gè)網(wǎng)格經(jīng)過運(yùn)算直接輸出物體的位置信息和置信度。并且以經(jīng)驗(yàn)預(yù)設(shè)的瞄框?yàn)榛绢A(yù)測邊框，對每個(gè)網(wǎng)格中心計(jì)算可能有物體的圖片區(qū)域進(jìn)行識別。本文算法以YOLOv5 為基礎(chǔ)搭建卷積神經(jīng)網(wǎng)絡(luò)，網(wǎng)絡(luò)結(jié)構(gòu)如圖1所示。

圖1 YOLO 網(wǎng)絡(luò)結(jié)構(gòu)模型

本文針對數(shù)據(jù)集特性搭建卷積神經(jīng)網(wǎng)絡(luò)，分別構(gòu)建輸入端、主干網(wǎng)絡(luò)Backbone、特征融合網(wǎng)絡(luò)Neck 和預(yù)測端Detection。在輸入端對數(shù)據(jù)進(jìn)行多項(xiàng)預(yù)處理工作。在主干網(wǎng)絡(luò)部分，層疊搭建瓶頸層C3（CSP Bottleneck with 3 convolutions）和卷積層Conv，C3 是簡化版本的Bottleneck 瓶頸層模塊，在主干網(wǎng)絡(luò)的末尾使用SPP 特征池化模塊對特征進(jìn)行解碼處理，送入neck網(wǎng)絡(luò)進(jìn)行進(jìn)一步的特征融合。SPP 模塊融合局部特征和全局特征融合，豐富特征圖的表達(dá)，有利于檢測精度的提升。在Neck 特征融合網(wǎng)絡(luò)部分，構(gòu)建Upsample 特征上采樣模塊與Concat 特征圖融合模塊實(shí)現(xiàn)特征的多尺度融合，保證鋼材數(shù)據(jù)的特征被網(wǎng)絡(luò)有效接收。每個(gè)模塊內(nèi)涉及到的基本卷積操作都使用了Conv+BatchNormalization+激活函數(shù)的標(biāo)準(zhǔn)形式。

1.2 數(shù)據(jù)集預(yù)處理

本文訓(xùn)練數(shù)據(jù)集為公開發(fā)布的鋼筋計(jì)數(shù)圖像。仔細(xì)分析鋼筋數(shù)據(jù)圖像特點(diǎn)，可以看出鋼筋圖像目標(biāo)十分密集，單張圖片往往存在30～50 左右的鋼筋目標(biāo)。其次，圖像中目標(biāo)數(shù)據(jù)和背景相似度較高，存在大量的噪聲干擾情況。此外，成像背景為工地建筑場地，背景較為復(fù)雜。

數(shù)據(jù)集預(yù)處理。為了使得鋼材數(shù)據(jù)集目標(biāo)特征有效表達(dá)，使得模型獲得較好的檢測精度。首先對鋼材數(shù)據(jù)圖像進(jìn)行mosaic 增強(qiáng)處理，mosaic 數(shù)據(jù)增強(qiáng)方法通過模糊濾波使得檢測物體暴露，模型更容易區(qū)分目標(biāo)和背景的關(guān)系，減少噪聲干擾。為了避免數(shù)據(jù)特征不足導(dǎo)致模型精度低的情況發(fā)生，可使用坐標(biāo)旋轉(zhuǎn)，圖像拼接等數(shù)據(jù)擴(kuò)充方法對數(shù)據(jù)集進(jìn)行增強(qiáng)，提高鋼材數(shù)據(jù)集的分布狀態(tài)質(zhì)量及目標(biāo)檢測模型的泛化性能。鋼材數(shù)據(jù)集標(biāo)注使用csv 格式，將標(biāo)簽文件標(biāo)準(zhǔn)化為yolo 格式，將Annotations 文件夾內(nèi)的xml 文件標(biāo)注鋼筋。

經(jīng)過數(shù)據(jù)預(yù)處理后，劃分訓(xùn)練集與測試機(jī)比例8∶2，采用矩形訓(xùn)練的方式送入模型。

K-means 聚類先驗(yàn)框。由于鋼材計(jì)數(shù)數(shù)據(jù)集的特殊性，為了獲得更好的模型精度，使用K-means 聚類得到最適合數(shù)據(jù)集的錨框。K-means 本質(zhì)是基于距離的數(shù)據(jù)劃分算法，其認(rèn)為2 個(gè)目標(biāo)的距離越近，相似度越高，均值和方差大的維度將對數(shù)據(jù)的聚類產(chǎn)生決定性影響。K-means 算法首先選k 個(gè)樣本作為初始聚類中心即a=a1，a2...ak，再針對數(shù)據(jù)集中的每個(gè)樣本xi，計(jì)算其到聚類中心距離，將其分到距離最小的中心簇中并重新計(jì)算該簇中的所有樣本的聚類中心重復(fù)運(yùn)算直至終止條件。在YOLO 模型中K-means 使用的聚類距離公式如下

由于鋼材數(shù)據(jù)集的特殊性，對數(shù)據(jù)集進(jìn)行3 種尺度的劃分，即期望模型使用3 種小尺度的錨框擬合目標(biāo)，每類物體3 個(gè)預(yù)測框，聚類結(jié)果為（18，18），（21，21），（25，24），（28，29），（33，28），（33，34），（38，38），（45，44），（54，53）。

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

2.1 實(shí)驗(yàn)環(huán)境

實(shí)驗(yàn)環(huán)境為服務(wù)器Ubuntu 20.04，6 張NVIDIA GeForce GTX 1080 Ti（11G）深度學(xué)習(xí)環(huán)境為Pytorch 1.9.1+CUDA 11.1+cudnn 8.0.5。訓(xùn)練batch_size=12，學(xué)習(xí)率設(shè)置為0.02，共訓(xùn)練200 epoch。

2.2 實(shí)驗(yàn)方法

本文以交并比（Intersection over Union）、準(zhǔn)確率（Precision）、召回率（Recall）、平均精度均值mAP（mean Average Precision）、完全參數(shù)量、GFlops（Giga Floating-point Operations Per Second）計(jì)算量和模型體積作為評估模型性能的指標(biāo)，以此客觀評價(jià)本文算法的性能。mAP 以準(zhǔn)確率和召回率為基礎(chǔ)計(jì)算多個(gè)類別AP 的平均值，mAP 值越大則認(rèn)為算法精度越好。完全參數(shù)量為模型訓(xùn)練時(shí)需要計(jì)算的全部參數(shù)，GFlops 代表每秒10 億次的浮點(diǎn)運(yùn)算數(shù)。GFlops 計(jì)算量是衡量算法計(jì)算復(fù)雜度的重要指標(biāo)。

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

實(shí)驗(yàn)結(jié)果由表1可知，訓(xùn)練100 個(gè)epoch 作為基礎(chǔ)，并且以相同環(huán)境下測試數(shù)據(jù)集作為算法返回結(jié)果。

表1 本文模型測試性能表

通過實(shí)驗(yàn)結(jié)果可知，本文算法平均檢測精度可達(dá)98.7%以上，模型較好的完成了鋼材計(jì)數(shù)任務(wù)。但是由于數(shù)據(jù)圖像的尺寸，數(shù)量及規(guī)模較小，圖像數(shù)據(jù)非均勻分布等原因。模型訓(xùn)練曲線出現(xiàn)了一些震蕩收斂情況。后續(xù)可改進(jìn)數(shù)據(jù)集質(zhì)量避免這一情況發(fā)生。本文算法的Loss 曲線，Precision-Recall 曲線及mAP 曲線如圖2所示。

圖2 模型Loss，Precision-Recall 以及mAP 曲線

本文算法模型訓(xùn)練完成，得到包含最優(yōu)網(wǎng)絡(luò)結(jié)構(gòu)參數(shù)的模型權(quán)重文件大小僅為14.4 M。終端設(shè)備可通過遷移學(xué)習(xí)方法部署模型，將會(huì)大大的減少內(nèi)存資源的占用。對訓(xùn)練完成的算法模型進(jìn)行測試，測試結(jié)果如圖3所示。

圖3 模型檢測結(jié)果

3 結(jié)束語

本文提出了一種基于YOLOv5 的鋼材圖像計(jì)數(shù)算法。實(shí)驗(yàn)結(jié)果表明，本文算法平均檢測精度可達(dá)98.7%以上，模型準(zhǔn)確率可98.9%，該模型較好完成了鋼材計(jì)數(shù)檢測任務(wù)，且具有一定魯棒性。后續(xù)模型研究可以考慮擴(kuò)充數(shù)據(jù)集，提高數(shù)據(jù)集質(zhì)量獲得更好的訓(xùn)練效果?；?qū)W(wǎng)絡(luò)結(jié)構(gòu)優(yōu)化，引入注意力機(jī)制或者多尺度融合機(jī)制提升模型擬合能力。為了更好部署于邊緣終端設(shè)備，可考慮搭建輕量化卷積網(wǎng)絡(luò)訓(xùn)練模型，方便部署，加快推理速度。