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

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

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

0 引言

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

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

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

1 算法

1.1 算法整體框架

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

圖1 算法流程圖

1.2 基于YOLOv4-tiny的鎖孔檢測

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

圖2 YOLOv4-tiny網(wǎng)絡(luò)結(jié)構(gòu)

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

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

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

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

圖4 U-net模型框架

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

不同行業(yè)中的企業(yè)對人才需求的類別、特質(zhì)不同，同一行業(yè)中的企業(yè)對人才需求的時間、能力也有所不同，如何抽取企業(yè)需求，分解能力指標，區(qū)分能力度，形成可初初始化的需求庫，是構(gòu)建共享平臺的關(guān)鍵問題之一，這需要對行業(yè)和企業(yè)進行深度調(diào)研，有針對性的構(gòu)建人才需求平臺。

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

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

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

2 實驗

2.1 數(shù)據(jù)集及標注

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

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

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

圖6 U-net標注的鎖孔數(shù)據(jù)集

2.2 實驗環(huán)境和實驗參數(shù)

本實驗基于Pytorch深度學(xué)習(xí)平臺進行，操作系統(tǒng)為Ubuntu18.04，CPU為AMD Ryzen 5 3500X，GPU為NVIDIA GTX 1060SUPER 6G。第1階段的模型訓(xùn)練采用ADAM優(yōu)化器，訓(xùn)練批尺寸為16，初始學(xué)習(xí)率設(shè)置為0.001，并采用學(xué)習(xí)率衰減的方式進行訓(xùn)練，一共迭代100次。第2階段的模型訓(xùn)練同樣采用Adam優(yōu)化器，網(wǎng)絡(luò)分辨率為256×256，訓(xùn)練批尺寸為2，學(xué)習(xí)率初始化為10-5，一共循環(huán)100次。

2.3 評價指標

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

2.3.2 鎖孔中心檢測評價指標

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

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

2.4 實驗結(jié)果與分析

2.4.1 鎖孔檢測實驗結(jié)果

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

表1 鎖孔檢測算法對比

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

2.4.2 鎖孔中心定位檢測結(jié)果

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

圖7 不同算法效果展示

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

同時，在相同數(shù)據(jù)集下對不同方法的中心誤差進行計算，統(tǒng)計了3種方法在不同誤差內(nèi)的圖片比例。對比結(jié)果如表2所示。

表2 不同算法結(jié)果對比

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

3 結(jié)論

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