零樣本6D物體姿態估計框架SAM-6D，向具身智能更進一步

作者：機器之心 2024-03-25 08:35:00

零樣本 6D 姿態估計是一種更具泛化性的任務設置，給定任意物體的 CAD 模型，旨在場景中檢測出該目標物體，并估計其 6D 姿態。

物體姿態估計在許多現實世界應用中起到至關重要的作用，例如具身智能、機器人靈巧操作和增強現實等。

在這一領域中，最先受到關注的任務是實例級別 6D 姿態估計，其需要關于目標物體的帶標注數據進行模型訓練，使深度模型具有物體特定性，無法遷移應用到新物體上。后來研究熱點逐步轉向類別級別 6D 姿態估計，用于處理未見過的物體，但要求該物體屬于已知的感興趣類別。

而零樣本 6D 姿態估計是一種更具泛化性的任務設置，給定任意物體的 CAD 模型，旨在場景中檢測出該目標物體，并估計其 6D 姿態。盡管其具有重要意義，這種零樣本的任務設置在物體檢測和姿態估計方面都面臨著巨大的挑戰。

圖 1. 零樣本 6D 物體姿態估計任務示意

最近，分割一切模型 SAM [1] 備受關注，其出色的零樣本分割能力令人矚目。SAM 通過各種提示，如像素點、包圍框、文本和掩膜等，實現高精度的分割，這也為零樣本 6D 物體姿態估計任務提供了可靠的支撐，展現了其前景的潛力。

因此，來自跨維智能、香港中文大學（深圳）、華南理工大學的研究人員提出了一個新穎的零樣本 6D 物體姿態估計框架 SAM-6D。該論文目前已被 CVPR 2024 接受。

論文鏈接: https://arxiv.org/pdf/2311.15707.pdf
代碼鏈接: https://github.com/JiehongLin/SAM-6D

SAM-6D 通過兩個步驟來實現零樣本 6D 物體姿態估計，包括實例分割和姿態估計。相應地，給定任意目標物體，SAM-6D 利用兩個專用子網絡，即實例分割模型（ISM）和姿態估計模型（PEM），來從 RGB-D 場景圖像中實現目標；其中，ISM 將 SAM 作為一個優秀的起點，結合精心設計的物體匹配分數來實現對任意物體的實例分割，PEM 通過局部到局部的兩階段點集匹配過程來解決物體姿態問題。SAM-6D 的總覽如圖 2 所示。

圖 2. SAM-6D 總覽圖

總體來說，SAM-6D 的技術貢獻可概括如下：

SAM-6D 是一個創新的零樣本 6D 姿態估計框架，通過給定任意物體的 CAD 模型，實現了從 RGB-D 圖像中對目標物體進行實例分割和姿態估計，并在 BOP [2] 的七個核心數據集上表現優異。
SAM-6D 利用分割一切模型的零樣本分割能力，生成了所有可能的候選對象，并設計了一個新穎的物體匹配分數，以識別與目標物體對應的候選對象。
SAM-6D 將姿態估計視為一個局部到局部的點集匹配問題，采用了一個簡單但有效的 Background Token 設計，并提出了一個針對任意物體的兩階段點集匹配模型；第一階段實現粗糙的點集匹配以獲得初始物體姿態，第二階段使用一個新穎的稀疏到稠密點集變換器以進行精細點集匹配，從而對姿態進一步優化。

實例分割模型 (ISM)

SAM-6D 使用實例分割模型（ISM）來檢測和分割出任意物體的掩膜。

給定一個由 RGB 圖像表征的雜亂場景，ISM 利用分割一切模型（SAM）的零樣本遷移能力生成所有可能的候選對象。對于每個候選對象，ISM 為其計算一個物體匹配分數，以估計其與目標物體之間在語義、外觀和幾何方面的匹配程度。最后通過簡單設置一個匹配閾值，即可識別出與目標物體所匹配的實例。

物體匹配分數的計算通過三個匹配項的加權求和得到:

語義匹配項 —— 針對目標物體，ISM 渲染了多個視角下的物體模板，并利用 DINOv2 [3] 預訓練的 ViT 模型提取候選對象和物體模板的語義特征，計算它們之間的相關性分數。對前 K 個最高的分數進行平均即可得到語義匹配項分數，而最高相關性分數對應的物體模板視為最匹配模板。

外觀匹配項 —— 對于最匹配模板，利用 ViT 模型提取圖像塊特征，并計算其與候選對象的塊特征之間的相關性，從而獲得外觀匹配項分數，用于區分語義相似但外觀不同的物體。

幾何匹配項 —— 鑒于不同物體的形狀和大小差異等因素，ISM 還設計了幾何匹配項分數。最匹配模板對應的旋轉與候選對象點云的平均值可以給出粗略的物體姿態，利用該姿態對物體 CAD 模型進行剛性變換并投影可以得到邊界框。計算該邊界框與候選邊界框的交并比（IoU）則可得幾何匹配項分數。

姿態估計模型 (PEM)

對于每個與目標物體匹配的候選對象，SAM-6D 利用姿態估計模型（PEM）來預測其相對于物體 CAD 模型的 6D 姿態。

將分割的候選對象和物體 CAD 模型的采樣點集分別表示為和，其中 N_m 和 N_o 表示它們點的數量；同時，將這兩個點集的特征表示為和，C 表示特征的通道數。PEM 的目標是得到一個分配矩陣，用于表示從 P_m 到 P_o 之間的局部到局部對應關系；由于遮擋的原因，P_o 只部分與匹配 P_m，而由于分割不準確性和傳感器噪聲，P_m 也只部分與匹配 P_o。

為了解決兩個點集非重疊點的分配問題，ISM 為它們分別配備了 Background Token，記為和，則可以基于特征相似性有效地建立局部到局部對應關系。具體來說，首先可以計算注意力矩陣如下：

接著可得分配矩陣

和分別表示沿著行和列的 softmax 操作，表示一個常數。中的每一行的值（除了首行），表示點集 P_m 中每個點 P_m 與背景及 P_o 中點的匹配概率，通過定位最大分數的索引，則可以找到與 P_m 匹配的點（包括背景）。

一旦計算獲得，則可以聚集所有匹配點對 {(P_m,P_o)} 以及它們的匹配分數，最終利用加權 SVD 計算物體姿態。

圖 3. SAM-6D 中姿態估計模型 (PEM) 的示意圖

利用上述基于 Background Token 的策略，PEM 中設計了兩個點集匹配階段，其模型結構如圖 3 所示，包含了特征提取、粗略點集匹配和精細點集匹配三個模塊。

粗糙點集匹配模塊實現稀疏對應關系，以計算初始物體姿態，隨后利用該姿態來對候選對象的點集進行變換，從而實現位置編碼的學習。

精細點集匹配模塊結合候選對象和目標物體的采樣點集的位置編碼，從而注入第一階段的粗糙對應關系，并進一步建立密集對應關系以得到更精確的物體姿態。為了在這一階段有效地學習密集交互，PEM 引入了一個新穎的稀疏到稠密點集變換器，它實現在密集特征的稀疏版本上的交互，并利用 Linear Transformer [5] 將增強后的稀疏特征擴散回密集特征。

實驗結果

對于 SAM-6D 的兩個子模型，實例分割模型（ISM）是基于 SAM 構建而成的，無需進行網絡的重新訓練和 finetune，而姿態估計模型（PEM）則利用 MegaPose [4] 提供的大規模 ShapeNet-Objects 和 Google-Scanned-Objects 合成數據集進行訓練。

為驗證其零樣本能力，SAM-6D 在 BOP [2] 的七個核心數據集上進行了測試，包括了 LM-O，T-LESS，TUD-L，IC-BIN，ITODD，HB 和 YCB-V。表 1 和表 2 分別展示了不同方法在這七個數據集上的實例分割和姿態估計結果的比較。相較于其他方法，SAM-6D 在兩個方法上的表現均十分優異，充分展現其強大的泛化能力。