基于U-net和YOLOv4-tiny的鎖孔中心定位算法*

謝孟添 刁云峰 程文明 唐 鑫 肖施睿

1西南交通大學機械學院 成都 610031 2軌道交通運維技術與裝備四川省重點實驗室 成都 610031

0 引言

近年來，隨著我國經濟的飛速發展，鐵路集裝箱貨運需求量大大增加。為了應對運量的上升，起重機的自動化是至關重要的一步。目前，在集裝箱裝卸工作中，起重機操作人員需要在地面員工的輔助下進行集裝箱的裝卸，大大降低了集裝箱的裝卸效率。更重要的是，長時間的手工工作可能導致工作事故。調查顯示[1]，2015年～2020年，人為因素造成的事故占到集裝箱碼頭所有事故的一半以上。因此，目前迫切需要一種高效、安全的自動化方法來替代人工操作。

隨著計算機視覺領域的發展，基于圖像視覺的識別和定位技術逐漸運用在了集裝箱領域[2-4]。Yoon H J等[5]利用雙目視覺系統實現了集裝箱高度距離的檢測，但該方法計算復雜度高，對環境光照敏感，且無法定位集裝箱的橫向與縱向位置，故難以應用于實際的集裝箱貨場。李巖等[6]提出了一種基于機器學習的集裝箱鎖孔檢測和定位方法，該方法運用改進的HOG算法提取鎖孔特征，并利用SVM分類器實現了鎖孔的實時檢測和定位，但檢測準確率并未達到實際應用的水平。文獻[7]利用傳統的邊緣檢測算法來提取鎖孔的邊緣，并在理想的實驗條件下通過微縮集裝箱模型進行了測試，得到了較好的對位精度。但實際條件下的集裝箱通常會受到環境腐蝕，鎖孔周圍的銹跡會大大降低邊緣檢測算法的精度。Li Y等[8]利用目標檢測的方法對集裝箱鎖孔進行檢測，通過檢測框的中心位置來進行鎖孔定位。雖然該方法取得了較好的檢測精度，但僅通過檢測框中心難以計算出準確的鎖孔中心。

為了解決上述問題，本文提出了一種基于圖像分割的鎖孔中心定位方法。該方法分為2個階段：第1階段利用基于深度學習的目標檢測方法對鎖孔進行檢測，提取鎖孔所在的矩形區域；第2階段采用圖像分割方法對矩形框進行預測，從而精確計算出鎖孔區域及其中心位置。同時，在鐵路集裝箱貨場進行數據采集，并在該數據集下進行了實驗評估，證實了本文提出方法的有效性。

1 算法

1.1 算法整體框架

本文采用兩段式的方法對鎖孔中心進行定位。因為基于圖像分割的方法在處理大分辨率的圖像時會產生較大計算量，同時集裝箱貨場環境復雜，難以保證分割精度。因此，在算法的第1步中，采用基于深度學習的目標檢測算法YOLOv4-tiny[9]來對原始圖片進行檢測，利用矩形框將鎖孔從原始圖像中提取出來。第2步，利用基于U-net[10]的圖像分割的方法對提取的檢測框進行預測，從而準確地得出鎖孔區域，并根據圖像處理的方法計算出區域中心的準確位置。其算法流程如圖1所示。

圖1 算法流程圖

1.2 基于YOLOv4-tiny的鎖孔檢測

本文選用了YOLOv4-tiny算法作為第1階段的目標檢測算法。因為與YOLOv4相比，YOLOv4-tiny結構更輕量化，運算速度更快。此外，相比于two-stage目標檢測算法和one-stage目標檢測算法，YOLOv4-tiny在保持檢測準確率的基礎上還能達到較快的運算速度，更利于集裝箱貨場這類需求實時檢測的應用環境。其網絡結構如圖2所示。

圖2 YOLOv4-tiny網絡結構

其網絡結構主要由2部分組成，骨干網絡和主干網絡。其中，骨干網絡為CSPDarknet53-tiny。圖片輸入后，骨干網絡會對其進行特征提取和融合，以生成不同深度的特征。隨后，主干網絡會對尺寸為13×13的深層特征282進行上采樣操作，并與尺寸為26×26的淺層特征17進行特征融合，豐富淺層特征的語義信息，提升網絡檢測能力。融合的特征在頭部網絡中進行預測，其中尺寸為26×26的特征用于小目標的預測，13×13的特征用于大目標的預測。根據生成預測目標的置信度和邊界框位置，通過置信度閾值和非極大值抑制來篩選出準確的目標，得到的鎖孔區域檢測結果如圖3所示。

1.3 基于U-net的鎖孔中心計算

在深度學習領域主要有2類圖像分割算法：1）基于卷積神經網絡（CNN）的方法，該方法通過卷積運算進行圖像特征的提取，但在卷積過程中容易損失圖像信息。且由于其類別預測采用的全連接層，故難以做到像素級的精確分割。2）利用反卷積層替代全連接層，采用全卷積神經網絡（FCN）對輸入圖片進行處理。輸出與原圖尺寸相同的分割圖，實現對圖像的精確分割。

本文選用了U-net模型作為第2階段的圖像分割模型，因為U-net模型結構簡單，同時能夠在較小的數據集下取得較好的分割效果，大大降低了實際應用時數據籌備的時間。該網絡是經典的全卷積神經網絡（見圖4），網絡采用3×3的卷積核對輸入圖片進行一系列的卷積和池化，得到4個不同尺寸的特征，然后對小尺寸的深層特征進行上采樣操作，同時采用級聯的方式與同尺寸的淺層特征進行拼接，使得深層特征和淺層特征得到了充分的利用，大大提高了圖像分割效果。

圖4 U-net模型框架

將第1階段提取的鎖孔矩形框輸入U-net圖像分割算法，可得到精確分割后的鎖孔區域灰度圖。由于鎖孔區域是規則的幾何形狀，故通過擬合該區域的形心即可精確計算出鎖孔的實際中心。

不同行業中的企業對人才需求的類別、特質不同，同一行業中的企業對人才需求的時間、能力也有所不同，如何抽取企業需求，分解能力指標，區分能力度，形成可初初始化的需求庫，是構建共享平臺的關鍵問題之一，這需要對行業和企業進行深度調研，有針對性的構建人才需求平臺。

本文采用基于圖像處理的方法對分割后的灰度圖進行二值化處理，精確計算出鎖孔內所有白色像素點的坐標，并通過式（1）計算出鎖孔中心在圖像中的像素坐標。

式中：xa、ya分別為鎖孔中心的橫、縱坐標，n為白色像素點的總數,xi、yi為第i個白色像素點的橫、縱坐標。

圖5 U-net鎖孔中心計算效果圖

2 實驗

2.1 數據集及標注

由于目前在集裝箱鎖孔檢測領域并沒有公共的數據集，為了保證算法檢測的準確度和有效性，本文在鐵路集裝箱貨場進行了數據采集，同時為了增加數據環境的多樣性，提高模型的泛化能力，在晴天、多云以及雨天的環境下進行數據采集，共采集1 409張集裝箱圖片和5 237鎖孔圖片。

對于集裝箱圖片，本文采用LabelImg標注軟件將其轉換成VOC2007[11]數據集格式，并按照7:2:1的比例將其劃分為訓練集、驗證集和測試集，以此構成第1階段的數據集。為了保證算法在鎖孔檢測時矩形框能完整的提取鎖孔，在標注過程中，標注框略應大于鎖孔邊緣。

對于鎖孔圖片，本文采用二階貝塞爾曲線[12]對鎖孔邊緣進行精確的標注，根據獲取的曲線坐標重新繪制鎖孔邊緣，并分別用白色像素和黑色像素填充鎖孔邊緣的內部和外部。最終將其作為第2階段的數據集。完整的數據集處理過程如圖6所示。圖中黃色邊緣線為二階貝塞爾曲線標注的鎖孔邊緣。

圖6 U-net標注的鎖孔數據集

2.2 實驗環境和實驗參數

本實驗基于Pytorch深度學習平臺進行，操作系統為Ubuntu18.04，CPU為AMD Ryzen 5 3500X，GPU為NVIDIA GTX 1060SUPER 6G。第1階段的模型訓練采用ADAM優化器，訓練批尺寸為16，初始學習率設置為0.001，并采用學習率衰減的方式進行訓練，一共迭代100次。第2階段的模型訓練同樣采用Adam優化器，網絡分辨率為256×256，訓練批尺寸為2，學習率初始化為10-5，一共循環100次。

2.3 評價指標

為了評價算法能否準確地從集裝箱圖片中檢測鎖孔，本文采用準確率acc作為鎖孔檢測的評價指標。該指標代表所有預測為鎖孔的對象N中正確預測T的比例，計算公式為

2.3.2 鎖孔中心檢測評價指標

本文選用了實際中心與預測中心像素坐標的歐氏距離dis作為第2階段鎖孔中心檢測的評價指標。當距離越小時，代表預測中心與實際中心越接近，算法精度越高。

式中：xt、yt為像素坐標系下鎖孔實際中心的橫坐標與縱坐標，xp、yp為像素坐標系下鎖孔預測中心的橫坐標與縱坐標。

2.4 實驗結果與分析

2.4.1 鎖孔檢測實驗結果

為了驗證在第一階段的鎖孔檢測操作中，YOLOv4-tiny算法能夠準確快速地提取鎖孔，本文將YOLOv4-tiny與Faster R-CNN[13]和RetinaNet[14]算法進行對比。在相同數據集下的檢測結果如表1所示。

表1 鎖孔檢測算法對比

由表1可看出，RetinaNet的鎖孔檢測平均成功率和YOLOv4-tiny較為接近，但本文選用的YOLOv4-tiny算法在檢測速度上擁有明顯的優勢，完全能滿足第一階段鎖孔檢測的實時性需求。

2.4.2 鎖孔中心定位檢測結果

本文模型在相同的數據集和評價標準下與基于目標檢測的方法和DIAO等[7]提出的方法做了精度對比。其中，基于目標檢測的方法選用了YOLOv4-tiny作為對比，因為該方法是目前實時檢測算法中檢測精度較高的算法，能保證本次對比實驗的公平性。實驗對比的可視化結果如圖7所示。其中紅圈和黃圈中心分別為實際鎖孔中心和預測鎖孔中心，其半徑均為5像素值。當兩圈重合面積較大時代表預測中心與實際中心更接近。

圖7 不同算法效果展示

從對比結果中可以看出DIAO等提出的方法在鎖孔銹蝕較為嚴重時檢測精度很差，完全偏離了實際的鎖孔中心。即使在檢測表面狀況良好的鎖孔時，其檢測精度也無法達到實際應用需求。此外，由于YOLOv4-tiny算法以檢測框中心近似替代實際鎖孔中心，故中心定位精度較差。而本文方法所擬合的鎖孔中心幾乎與實際中心重疊。

同時，在相同數據集下對不同方法的中心誤差進行計算，統計了3種方法在不同誤差內的圖片比例。對比結果如表2所示。

表2 不同算法結果對比

由表2可以看出，本文提出的方法在中心定位精度方面，遠超另外兩類算法，且在3像素誤差內的檢測成功率可達94%，完全滿足實際應用的需求。

3 結論

快速準確的鎖孔中心定位算法是實現集裝箱裝卸自動化的關鍵技術之一，本文提出了一種基于深度學習的解決方案。該方案包括：1）利用YOLOv4-tiny算法對集裝箱圖片進行檢測，獲取鎖孔區域的矩形框；2）采用基于U-net的圖像分割模型對鎖孔矩形框進行分割，實現鎖孔區域的提取；3）應用圖像處理的方法對分割的鎖孔區域進行計算，從而擬合出鎖孔區域的中心。結果表明：本方案在實際的集裝箱數據集上的識別精度遠高于其他鎖孔中心定位算法，能夠符合集裝箱鎖孔中心實時定位的需求。