基于YOLOv5的鋼材數量識別方法

劉洪寧，李容露

（大連民族大學機電工程學院，遼寧 大連 116600）

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

1 基于YOLOv5 的鋼材數量識別算法

1.1 目標檢測模型搭建

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

圖1 YOLO 網絡結構模型

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

1.2 數據集預處理

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

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

經過數據預處理后，劃分訓練集與測試機比例8∶2，采用矩形訓練的方式送入模型。

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

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

2 實驗結果分析

2.1 實驗環境

實驗環境為服務器Ubuntu 20.04，6 張NVIDIA GeForce GTX 1080 Ti（11G）深度學習環境為Pytorch 1.9.1+CUDA 11.1+cudnn 8.0.5。訓練batch_size=12，學習率設置為0.02，共訓練200 epoch。

2.2 實驗方法

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

2.3 實驗結果分析

實驗結果由表1可知，訓練100 個epoch 作為基礎，并且以相同環境下測試數據集作為算法返回結果。

表1 本文模型測試性能表

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

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

本文算法模型訓練完成，得到包含最優網絡結構參數的模型權重文件大小僅為14.4 M。終端設備可通過遷移學習方法部署模型，將會大大的減少內存資源的占用。對訓練完成的算法模型進行測試，測試結果如圖3所示。

圖3 模型檢測結果

3 結束語

本文提出了一種基于YOLOv5 的鋼材圖像計數算法。實驗結果表明，本文算法平均檢測精度可達98.7%以上，模型準確率可98.9%，該模型較好完成了鋼材計數檢測任務，且具有一定魯棒性。后續模型研究可以考慮擴充數據集，提高數據集質量獲得更好的訓練效果。或對網絡結構優化，引入注意力機制或者多尺度融合機制提升模型擬合能力。為了更好部署于邊緣終端設備，可考慮搭建輕量化卷積網絡訓練模型，方便部署，加快推理速度。