ICLR2025|即時線上！清華、NTU提出適用於具身場景的3D分割一切模型

©PaperWeekly 原創 · 作者 | 許修為

單位 | 清華大學博士生

研究方向 | 計算機視覺，具身智慧

在這篇文章中，我們提出了 EmbodiedSAM，一個即時線上的 3D 分割一切模型。可以直接讀入 RGB-D 流影片進行 efficient 且 fine-grained 的線上三維例項分割，這也是對我們之前 CVPR 2024 工作 Online3D 的一個延續（更快、更細！），程式碼、資料和模型權重均在 github 開源，歡迎訪問。

論文標題：

EmbodiedSAM: Online Segment Any 3D Thing in Real Time

論文地址：

https://arxiv.org/pdf/2408.11811

程式碼倉庫：

https://github.com/xuxw98/ESAM

專案主頁：

https://xuxw98.github.io/ESAM

簡介

EmbodiedSAM 的出發點來自於一個思考：具身智慧需要什麼形式的場景感知？考慮到最為普遍的具身場景是機器人在未知環境中同時完成對場景的探索和感知，我們認為一個好的具身感知模型應該具有四點性質：

1. 線上。不同於先重建再感知，線上意味著機器人一邊讀入 RGB-D 影片流，一邊就要進行場景感知，在每一時刻對未來是不可知的；

2. 即時。一個快速的感知模型才能讓決策規劃的模組沒有負擔地進行呼叫，每秒鐘至少需要能夠處理上十幀的資料（當前水平 10+FPS，未來目標 30+FPS）；

3. 細粒度。希望能將場景中一切物體進行分割，包括各種小目標（未來可以考慮進一步細化到物體的 part）；

4. 泛化性。一個模型即可以處理多樣化的場景，甚至可以適配不同型號的感測器，即不同內參的 RGB-D 相機。這是一個美好的想法，但目前 3D 資料的質量和規模還難以滿足僅透過 3D 訓練出具備上述能力的模型。

如何解決這個問題？目前一個常用的思路是充分利用現有的高質量 2D 視覺基礎模型，目前已經有很多又快又好的模型可以以很快的速度實現細粒度的 2D 分割，如 FastSAM [1] 等。

目前SAM3D [2,3] 之類的 3D 分割工作透過 2D 基礎模型進行影像分割，將結果對映到深度圖轉化的點雲中，再在三維空間中設計基於點雲的合併策略以完成 3D 例項分割。這一套思路僅利用2D模型就可以完成較高細粒度和泛化性的 3D 例項分割，但存在諸多問題：

1. 只能用於離線資料，無法線上處理流影片；

2. 速度極慢，處理一個場景需要數分鐘甚至數十分鐘的時間；

3. 缺乏 3D 感知能力，僅獨立對每一幀 RGB-D 影像中的 RGB 資訊進行分割，沒有利用深度資訊，也沒有利用多視角的幾何關聯。

為此，我們的 EmbodiedSAM 採用先 2D 後 3D 的分割正規化，對每一張 RGB 影像利用高效率的基礎模型進行分割，得到初始的單幀分割結果，將此結果在 3D 空間進行可學習的最佳化和幀間融合，從而實現即時線上的 3D Segment Anything。

由於分割的任務交給了 2D 模型，我們的 3D 模型只需要充分利用時序和幾何資訊對來自 2D 的 mask 進行最佳化，以及對不同幀的相同物體進行融合，學習難度大大降低，目前的 3D 資料集完全可以滿足需求。

‍

以上是一個簡單的示例 demo（更詳細的影片可以在我們的專案主頁觀看），和 SAM3D 等離線 3D 分割方法相比，我們的方法可以用於更難的線上任務，且效能大幅提升，速度快將近 20 倍，處理一個完整的室內場景只需要5s左右的時間。

方法

下面是我們的整體框架，個人一直覺得方法部分看原文就好，中文解讀更重要的是 motivation 和整體思路，因此這次也是從比較 high-level 的角度介紹一下方法（其實是翻譯太累了）。

EmbodiedSAM 的輸入為已知相機內外參的 RGB-D 影片流，我們整體分為三個環節，對於每一時刻新讀取的 RGB-D 幀，首先是 Query Lifting 階段，將 2D 基礎模型的分割結果轉換到 3D 空間，並用 query 特徵的形式來表示，一個 query 對應一個 2D 模型分割出來的物體。

這種表示形式一方面可以充分將 3D 和時序特徵進行融入，另一方面也便於後續的 Mask 最佳化和幀間融合；之後是 Query Refinement 環節，透過物體 query 之間的互動，以及和場景特徵的互動，更新 query 特徵，並將 3D 中的分割結果進行更新最佳化。注意 Query Lifting 和 Query Refinement 操作都只對當前幀的物體進行操作。

最後在 Query Merging 環節中，我們將當前幀分割出的物體與過去的物體進行融合，例如當前幀和之前幀都觀測到了同一把椅子，那它們的全域性 instance id 應該相同，這也就是所謂的 instance merge，如果是之前幀沒見過的新物體，那便在全域性 instance id 中註冊新物體即可。

2.1 Query Lifting

對於當前時刻新觀測的 RGB-D，我們首先用 2D 基礎模型（VFM）對 RGB 進行分割，並將分割結果對映到深度圖上轉換為點雲，從而得到 3D 點雲上的分割結果。不同於之前的方法直接對 3D 分割結果進行融合，我們認為僅透過 2D VFM 分割，雖然細粒度和泛化性很強，但對於空間資訊的利用不夠充分，容易有較大誤差。

因此我們用一個具有時序能力的 3D 骨幹網路 [4] 提取點雲特徵（可以挖掘當前幀和過去幀點雲的關聯），並將 3D 分割結果視為 superpoint，引導 point-wise 特徵池化到 superpoint-wise 特徵，從而獲取各物體的 query 表示。

我們額外設計了一個具有幾何感知能力的池化方法（Geo-aware Pooling），對於落入同一個 superpoint 內的點雲特徵，不僅僅是做 max pooling 或者 average pooling，而是考慮到 superpoint 自身的幾何形狀，給不同空間位置的特徵賦予不同的權重。

2.2 Query Refinement

類似於基於 Mask2Former [5] 等基於 Transformer 的例項分割架構，我們也利用若干層 Transformer 對每個物體的 query 資訊進行更新最佳化，在每層Transformer 中，query 特徵分別和當前幀點雲特徵進行 cross-attention，和自身進行 self-attention，並最終和當前幀點雲特徵互動得到最佳化後的 3D Mask。

該過程有幾個不同的設計，首先是我們採用了不同層次的點雲特徵，來獲取最好的 efficiency-accuracy tradeoff。

對於 cross-attention，我們採用 superpoint-wise 特徵進行互動，superpoint 數量少但是包含了場景中的所有資訊，有利於加速運算。對於 3D Mask 預測，我們則將 query 和 point-wise 特徵進行互動，從而得到細粒度的分割結果。

其次，在每層 Transformer 拿到 point-wise 的分割結果後，我們會利用之前 Geo-aware Pooling 的權重將分割結果池化為 superpoint-wise，用於下一層中進行 masked cross-attention（即透過初步的分割結果限定 cross-attention 的作用範圍）。

2.3 Query Merging

完成上述階段後，我們拿到了當前幀最佳化後的 3D Mask，基於和 3D Mask一一對應的 query 特徵，同時我們還有之前幀的全域性 3D Mask。這一階段需要將當前幀 3D Mask 與之前的 Mask 進行全域性融合，判斷當前幀每個物體是融入到過去已存在的某個物體，還是註冊為新物體。

這一過程也是時間上最大的瓶頸，之前如 SAM3D 等手工融合的方法不僅對超引數較為敏感，而且速度極慢，能否又快又準完成 3D Mask 的全域性融合是一個非常有挑戰的問題。

這裡得益於我們的 query 表示，我們可以將 3D Mask 的融合問題轉換為 query 融合問題。不同物體的 3D Mask 含點雲數量不同，難以高效率平行計算，只能在迴圈中計算 3D Mask 間的空間距離設定閾值完成融合。

而不同物體的 query 尺寸是相同的，例如之前幀一共有 M 個物體，當前幀一共 N 個物體，由於每個物體都有一一對應的 query，因此這個問題可以透過矩陣乘法直接完成相似度計算！

假設 query 特徵維度為 C，我們將 MxC 和 NxC 的特徵相乘便得到 MxN 的相似度矩陣，再透過二分圖匹配的方法在該矩陣上得到匹配關係。對於匹配上的物體，我們將其 3D Mask 合併，並將對應的 query 特徵加權融合；對於未匹配上的物體則直接註冊為全域性新物體。

實際中，僅依靠 query 表示雖然高效，卻仍存在表示能力不夠的問題。我們在不影響速度的前提下設計了三種更細緻的判據，透過不同的 MLP head 在 query 上預測出三種表示：1）物體的外接框；2）物體的對比表示；3）物體的語義表示。這三種表示具體含義如下圖所示。

每次進行 3D Mask 的融合後，這三種表示也會進行對應的加權融合，並作為“過去幀”再次與未來的“當前幀”進行高效融合。由於所有操作僅需要矩陣乘法和二分圖匹配，所需時間在 5ms 以內，相比手工融合的策略快數百倍，且效能更強！

實驗結果

我們進行了多種設定的實驗。首先是 class-agnostic 3D instance segmentation，對標 SAM3D 等文章的設定。由於這些文章大多數為 zero-shot，而我們需要訓練，為了更公平展示泛化性，我們還進行了 zero-shot 遷移實驗，從 ScanNet200 遷移到 SceneNN 和 3RScan 上。

此外我們還進行了 online 3D instance segmentation 的實驗，對標之前的線上 3D 例項分割工作。在兩種任務上我們都取得了 SOTA 的效能和速度。

我們還分析了 EmbodiedSAM 的推理時間分佈，發現速度的瓶頸主要在 3D 特徵提取部分，我們的 merging 策略的開銷幾乎可以忽略不記。如何加速 3D 骨幹網路的特徵提取是一個非常有價值的問題，也是未來突破 30FPS 的最後阻礙。

更多實驗、視覺化和 Demo 影片可以參考我們的論文以及專案主頁。如有問題歡迎大家在 github 上開 issue 討論。

參考文獻

[1] Zhao, Xu, et al. 'Fast segment anything.'arXiv preprint arXiv:2306.12156(2023).

[2] Yang, Yunhan, et al. 'Sam3d: Segment anything in 3d scenes.'arXiv preprint arXiv:2306.03908(2023).

[3] Yin, Yingda, et al. 'Sai3d: Segment any instance in 3d scenes.'Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.

[4] Xu, Xiuwei, et al. 'Memory-based Adapters for Online 3D Scene Perception.'Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.

[5] Cheng, Bowen, et al. 'Masked-attention mask transformer for universal image segmentation.'Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

更多閱讀