基于改進YOLO6D的單目位姿估計算法研究

潘 江，任德均，史雨杭，王淋楠

（四川大學 機械工程學院，四川 成都 610065）

0 引 言

6D位姿估計在自動駕駛、機器人和現實增強等領域有著重要應用［1］。具有代表性的方法有模板匹配［2］、點云配準［3］、深度學習［4］等。基于單目RGB圖像的位姿估計方法可以直接從單張RGB圖像中估計出物體的6D位姿，具有硬件結構簡單、數據易獲取、運算量小、實時性好的特點，這對工程應用十分有利。Rad M等人［5］提出了一種基于單目RGB圖像的位姿估計方法，該方法能在一定程度上解決6D位姿估計問題的難點，如無紋理、背景雜亂、堆疊等，但算法流程復雜、精度低、推理速度慢。

近年來，一些單階段的目標檢測算法如YOLO［6］、SSD［7］等，在算法實時性、準確性等方面表現出優異的性能，這對6D位姿估計具有借鑒意義。Tekin B 等人［8］借鑒YOLO的思想提出了YOLO6D，該網絡是單階段的，且推理速度快，在LINEMOD［9］數據集上達到了較高的精度水平。并且由于遵循YOLO架構，推理速度不受物體數量的影響。

本文借鑒YOLO6D的思想，研究并設計了新的位姿估計網絡，分析并改進了損失函數的構建方法，通過實驗總結了關鍵點數量對透視投影變換PnP（perspective-n-point）算法解算位姿精度的影響規律，并證明了基于3D 最小包圍框（3D bounding box，3D BBox）和PnP解算位姿問題的最佳關鍵點數量，同時為3D BBox-PnP問題提供了一種新的解決思路。

1 算法原理與過程設計

1.1 算法原理及流程

本文算法原理及流程如圖1 所示。第一階段，卷積神經網絡（convolutional neural network，CNN）提取圖像特征并估計物體關鍵點的2D投影像素坐標；第二階段，根據這些關鍵點的2D坐標及其先驗的3D 坐標（相對于物體坐標系）構成2D-3D關鍵點對，利用PnP 算法解算出物體的位姿。選擇和BB8［5］類似的3D BBox 的8 個角點，外加1 個中心點共9 個關鍵點。以物體中心點為原點，分別平行BBox的長、寬、高建立物體坐標系，該坐標系相對于相機坐標系的旋轉矩陣R 和平移向量t 即為物體的6D 位姿，如圖2所示。

圖1 算法原理及流程

圖2 關鍵點和坐標系

1.2 算法網絡設計

本文采用的主干網絡是Liu Z等人［10］提出的純卷積網絡ConvNeXt，基礎結構如圖3 所示。主要具有以下特點：采用了7 ×7大卷積核；使用更少的激活函數并將ReLU替換為GELU；使用Layer Norm替換Batch Norm。相較于結構復雜的Transformer模型，ConvNeXt結構簡單，性能反而更優。

圖3 ConvNeXt的基礎結構ConvNeXt Block

本文網絡完整結構如圖4 所示。主干部分為原Conv-NeXt-T網絡去掉最后的平均池化層、Layer Norm層和全連接層后剩余的部分。Stem 表示先進行1 次核大小為4，步長為4 的卷積下采樣，再經過Layer Norm層歸一化；Downsample所做的操作與Stem相似，不同之處在于交換下采樣和層歸一化處理的順序，并且下采樣的因子變為1／2；Conv-X表示經過X 次“卷積-BN-GELU”操作；Upsample 表示上采樣，倍率為2；Concatenate表示2個特征圖在channel維度上拼接。網絡的輸出為3 個S×S×20 的特征圖，S的含義是將原圖劃分為S×S 個網格（YOLO 中的cell），對應輸出S×S 個特征向量，特征向量的維度為20，分別表示9個關鍵點的18個坐標值，1個類別分量（單目標預測僅需1位表示），1個置信度分數。

圖4 網絡完整結構

1.3 損失函數的構建

本文構建的損失函數如式（1）所示。其中，Lpt，Lcls，Lconf分別為網絡預測的位置損失、分類損失和置信度損失；λpt，λcls，λconf分別為上述3 種損失的權重系數。Lpt用均方誤差計算，Lcls用交叉熵誤差計算損失。本文改進的Lconf用Focal Loss計算，目的是平衡正負樣本以及難易樣本對損失函數的貢獻。除此之外，YOLO6D 通過實驗證明，計算3D BBox的IoU 十分耗時，因此提出式（2）近似替代IoU 的計算。其中，dth為距離閾值，D（x）為第x個BBox各關鍵點誤差的平均值，α為超參數，取值為2；c（x）即該BBox的近似IoU。該近似計算在幾乎不損失精度的前提下，極大減少了耗時。本文在此基礎上，增加對各關鍵點誤差方差的計算，保證BBox的投影在幾何形狀上趨近真實立方體的投影，計算公式如式（3）所示。其中，di為第i 個關鍵點的誤差，為這9個關鍵點的誤差的均值，c′（x）為第x個BBox的近似IoU與誤差方差的綜合損失。式（1）～式（3）表示如下

1.4 改進PnP位姿解算方法

根據相機投影原理，空間中的直線在圖像上的投影依然是直線，空間中兩直線的交點在圖像上的投影依然是對應兩直線的交點。根據3D BBox的幾何特點，將關鍵點數量由9 個增加至15 個，增加的6 個關鍵點為BBox 6 個面的中心點。這些新增關鍵點對應的2D坐標可由對應面的角點的2D坐標計算得到。需要說明的是，由于網絡輸出的4個角點并不總是能構成凸四邊形，從而無法計算中心點的2D坐標。在這種情況下，角點本身具有較大誤差，新增關鍵點意義不大。因此，本文只對可計算面中心點坐標的輸出做上述優化處理。將這15 個關鍵點按照誤差從小到大排列，依次選取前n（n ＝7，8，…，14，15）個誤差最小點作為PnP算法的關鍵點進行實驗，探究精度與關鍵點數量n的關系，實驗結果如表1 所示。其中，粗體表示優于YOLO6D的精度，斜體表示該研究對象的最高精度。實驗結果顯示，適當增加關鍵點數量能顯著提高位姿解算精度。根據實驗結果，本文選擇14作為算法后續研究的PnP關鍵點數量。除此之外，通過實驗證明，將關鍵點數量增加至14時，PnP解算耗時僅略微增加，但仍未超過0.5 ms，這對整個算法來說可忽略不計。

表1 改進PnP位姿解算方法對位姿估計精度的影響

2 實驗與結果分析

2.1 數據集與預處理

本文在LINEMOD數據集上進行了實驗。該數據集是事實上的研究雜亂場景無紋理物體6D 位姿估計的基準數據集。以該數據集的12個物體為研究對象，每個實驗對象隨機抽取85%的圖像作為訓練集，其余圖像作為測試集。為提升網絡的多尺度能力，除了網絡結構上的改進，在訓練時隨機改變輸入圖像的分辨率為32 的倍數（網絡是5 倍下采樣的）。為了減小過擬合，對原始圖像隨機裁剪、縮放，改變色調、飽和度和曝光，利用圖像蒙版隨機變換圖像背景。數據集原始圖像為640 ×480 的RGB 格式，將經過隨機裁剪、縮放和顏色調整的圖像，先以最大邊為基準，將圖像上下填充至長、寬相等，再尺寸變換至416 ×416以適應網絡輸入，從而避免物體變形。

2.2 網絡訓練

基于PyTorch深度學習框架搭建算法網絡部分，實驗平臺搭載CPU 為Intel i9-10900X，內存為64 GB，GPU 為NVIDIA GeForceGTX3080。使用AdamW 優化器，設置初始學習率為0.001，每經過150 個epoch衰減為原來的0.1，權重衰減設置為0.01。采用遷移學習技術，使用在ImageNet-1K上訓練的ConvNeXt-Tiny模型作為預訓練模型載入模型參數。由于設備顯存限制，采用梯度累積策略，設置batch size為4，每16步進行1次參數更新，等效于batch size為64的訓練。訓練輪數設置為500，每經過5 個epoch 對損失值和精度各采樣1次。驗證時，以2D重投影精度作為精度評價指標。值得一提的是，實驗中發現，若以5 像素作為閾值，算法在第280輪時精度就已經達到并穩定在100%，這不利于模型更新。因此，在訓練時設置的閾值為3 像素，訓練過程損失值和驗證精度的變化如圖5所示。

圖5 網絡訓練過程中損失值和驗證精度的變化

2.3 性能測評

為測試本文算法的多尺度能力，以Ape，Benchvise，Can為實驗對象，將輸入圖像的尺寸設置為224，256，…，576，608等32的倍數，閾值設置為5像素，分別統計本文算法和YOLO6D在各個尺度下的精度表現，實驗結果如圖6 所示。結果顯示，在所測試的13 個尺度上，本文算法精度均顯著優于YOLO6D，尤其是在小尺度上，精度提升較大。當尺寸為288時，本文算法的3 個實驗對象精度均超過80%。在尺寸區間416～544，檢測效果最好，精度穩定在99%以上。3個實驗對象在不同尺度下的平均精度，YOLO6D 為78.95%，而本文算法為90.88%，提升11.93個百分點。因此，在不同尺度尤其是在小尺度上，本文算法的精度表現良好，在更廣泛的尺度變化上擁有較高精度水平，擁有更強的多尺度檢測能力。

圖6 不同尺度下的算法精度

為測試本文算法在不同閾值下的精度表現，以Ape，Benchvise，Can 為實驗對象，統計了算法在不同像素閾值下的精度，實驗結果如表2所示。結果顯示，本文算法在不同閾值下的精度均有較大提升，當閾值大于等于5 像素時，3個實驗對象的精度維持在99%以上水平。

表2 不同像素閾值下的算法精度表現 %

本文對全部12個實驗對象進行了精度測試，閾值設置為5像素，結果如表3所示。根據實驗結果，本文算法在其中10個研究對象上的精度表現顯著優于YOLO6D，平均精度提升8.14個百分點。

表3 全部實驗對象的算法精度表現

本文算法在CPU AMD Ryzen7 6800H，GPU NVIDIA Ge-ForceGTX3060環境下，GPU 模式運行平均耗時約60 ms，具有較強的實時性。根據YOLO6D，在檢測多目標時，由于遵循YOLO架構，算法耗時不隨檢測物體數量的增加而增加。本文算法是基于YOLO6D 思想改進的算法，因此同樣具有上述特點。這在某些需要快速、準確、穩定檢測物體位姿的場景中具有較大的應用價值。

本文算法部分物體位姿估計效果可視化結果如圖7 所示。其中，深色表示物體真實3D 包圍框，淺色表示本文算法估計結果。

圖7 部分物體位姿估計可視化結果

3 結束語

本文提出了一種6D位姿估計算法，主要面向的是一類常見的位姿估計應用場景，即能較容易獲取被檢測物體的RGB圖像和尺寸信息（如通過掃描，CAD 模型等），需要快速、準確地獲取物體位姿信息的場景。本文提出的算法具有較高的精度，實時性強，模型簡單易搭建、易訓練，適用于無紋理或缺乏紋理、嚴重遮擋的場景下物體位姿的估計，且僅通過物體的RGB信息，從而簡化了硬件設施，這對位姿估計的工程應用具有實踐意義。