在這篇文章中，來(lái)自清華大學(xué)的團(tuán)隊(duì)提出了 EmbodiedSAM，一個(gè)實(shí)時(shí)在線的 3D 分割一切模型。可以直接讀入 RGB-D 流視頻進(jìn)行 efficient 且 fine-grained 的在線三維實(shí)例分割，這也是對(duì)團(tuán)隊(duì)之前 CVPR 2024 工作 Online3D 的一個(gè)延續(xù)（更快、更細(xì)！），代碼、數(shù)據(jù)和模型權(quán)重均在 github 開源，歡迎訪問。

打開網(wǎng)易新聞 查看精彩圖片

論文標(biāo)題： EmbodiedSAM: Online Segment Any 3D Thing in Real Time 論文地址： https://arxiv.org/pdf/2408.11811 代碼倉(cāng)庫(kù)： https://github.com/xuxw98/ESAM 項(xiàng)目主頁(yè)： https://xuxw98.github.io/ESAM

一、簡(jiǎn)介

EmbodiedSAM 的出發(fā)點(diǎn)來(lái)自于一個(gè)思考：具身智能需要什么形式的場(chǎng)景感知？考慮到最為普遍的具身場(chǎng)景是機(jī)器人在未知環(huán)境中同時(shí)完成對(duì)場(chǎng)景的探索和感知，我們認(rèn)為一個(gè)好的具身感知模型應(yīng)該具有四點(diǎn)性質(zhì)：

  1. 在線。不同于先重建再感知，在線意味著機(jī)器人一邊讀入 RGB-D 視頻流，一邊就要進(jìn)行場(chǎng)景感知，在每一時(shí)刻對(duì)未來(lái)是不可知的；

  2. 實(shí)時(shí)。一個(gè)快速的感知模型才能讓決策規(guī)劃的模塊沒有負(fù)擔(dān)地進(jìn)行調(diào)用，每秒鐘至少需要能夠處理上十幀的數(shù)據(jù)（當(dāng)前水平 10+FPS，未來(lái)目標(biāo) 30+FPS）；

  3. 細(xì)粒度。希望能將場(chǎng)景中一切物體進(jìn)行分割，包括各種小目標(biāo)（未來(lái)可以考慮進(jìn)一步細(xì)化到物體的 part）；

  4. 泛化性。一個(gè)模型即可以處理多樣化的場(chǎng)景，甚至可以適配不同型號(hào)的傳感器，即不同內(nèi)參的 RGB-D 相機(jī)。這是一個(gè)美好的想法，但目前 3D 數(shù)據(jù)的質(zhì)量和規(guī)模還難以滿足僅通過 3D 訓(xùn)練出具備上述能力的模型。

如何解決這個(gè)問題？目前一個(gè)常用的思路是充分利用現(xiàn)有的高質(zhì)量 2D 視覺基礎(chǔ)模型，目前已經(jīng)有很多又快又好的模型可以以很快的速度實(shí)現(xiàn)細(xì)粒度的 2D 分割，如 FastSAM [1] 等。

目前SAM3D [2,3] 之類的 3D 分割工作通過 2D 基礎(chǔ)模型進(jìn)行圖像分割，將結(jié)果映射到深度圖轉(zhuǎn)化的點(diǎn)云中，再在三維空間中設(shè)計(jì)基于點(diǎn)云的合并策略以完成 3D 實(shí)例分割。這一套思路僅利用2D模型就可以完成較高細(xì)粒度和泛化性的 3D 實(shí)例分割，但存在諸多問題：

  1. 只能用于離線數(shù)據(jù)，無(wú)法在線處理流視頻；

  2. 速度極慢，處理一個(gè)場(chǎng)景需要數(shù)分鐘甚至數(shù)十分鐘的時(shí)間；

  3. 缺乏 3D 感知能力，僅獨(dú)立對(duì)每一幀 RGB-D 圖像中的 RGB 信息進(jìn)行分割，沒有利用深度信息，也沒有利用多視角的幾何關(guān)聯(lián)。

為此，我們的 EmbodiedSAM 采用先 2D 后 3D 的分割范式，對(duì)每一張 RGB 圖像利用高效率的基礎(chǔ)模型進(jìn)行分割，得到初始的單幀分割結(jié)果，將此結(jié)果在 3D 空間進(jìn)行可學(xué)習(xí)的優(yōu)化和幀間融合，從而實(shí)現(xiàn)實(shí)時(shí)在線的 3D Segment Anything。

由于分割的任務(wù)交給了 2D 模型，我們的 3D 模型只需要充分利用時(shí)序和幾何信息對(duì)來(lái)自 2D 的 mask 進(jìn)行優(yōu)化，以及對(duì)不同幀的相同物體進(jìn)行融合，學(xué)習(xí)難度大大降低，目前的 3D 數(shù)據(jù)集完全可以滿足需求。

打開網(wǎng)易新聞 查看精彩圖片

以上是一個(gè)簡(jiǎn)單的示例 demo（更詳細(xì)的視頻可以在我們的項(xiàng)目主頁(yè)觀看），和 SAM3D 等離線 3D 分割方法相比，我們的方法可以用于更難的在線任務(wù)，且性能大幅提升，速度快將近 20 倍，處理一個(gè)完整的室內(nèi)場(chǎng)景只需要5s左右的時(shí)間。

二、方法

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

EmbodiedSAM 的輸入為已知相機(jī)內(nèi)外參的 RGB-D 視頻流，我們整體分為三個(gè)環(huán)節(jié)，對(duì)于每一時(shí)刻新讀取的 RGB-D 幀，首先是 Query Lifting 階段，將 2D 基礎(chǔ)模型的分割結(jié)果轉(zhuǎn)換到 3D 空間，并用 query 特征的形式來(lái)表示，一個(gè) query 對(duì)應(yīng)一個(gè) 2D 模型分割出來(lái)的物體。

這種表示形式一方面可以充分將 3D 和時(shí)序特征進(jìn)行融入，另一方面也便于后續(xù)的 Mask 優(yōu)化和幀間融合；之后是 Query Refinement 環(huán)節(jié)，通過物體 query 之間的交互，以及和場(chǎng)景特征的交互，更新 query 特征，并將 3D 中的分割結(jié)果進(jìn)行更新優(yōu)化。注意 Query Lifting 和 Query Refinement 操作都只對(duì)當(dāng)前幀的物體進(jìn)行操作。

最后在 Query Merging 環(huán)節(jié)中，我們將當(dāng)前幀分割出的物體與過去的物體進(jìn)行融合，例如當(dāng)前幀和之前幀都觀測(cè)到了同一把椅子，那它們的全局 instance id 應(yīng)該相同，這也就是所謂的 instance merge，如果是之前幀沒見過的新物體，那便在全局 instance id 中注冊(cè)新物體即可。

打開網(wǎng)易新聞 查看精彩圖片

2.1 Query Lifting

對(duì)于當(dāng)前時(shí)刻新觀測(cè)的 RGB-D，我們首先用 2D 基礎(chǔ)模型（VFM）對(duì) RGB 進(jìn)行分割，并將分割結(jié)果映射到深度圖上轉(zhuǎn)換為點(diǎn)云，從而得到 3D 點(diǎn)云上的分割結(jié)果。不同于之前的方法直接對(duì) 3D 分割結(jié)果進(jìn)行融合，我們認(rèn)為僅通過 2D VFM 分割，雖然細(xì)粒度和泛化性很強(qiáng)，但對(duì)于空間信息的利用不夠充分，容易有較大誤差。

因此我們用一個(gè)具有時(shí)序能力的 3D 骨干網(wǎng)絡(luò) [4] 提取點(diǎn)云特征（可以挖掘當(dāng)前幀和過去幀點(diǎn)云的關(guān)聯(lián)），并將 3D 分割結(jié)果視為 superpoint，引導(dǎo) point-wise 特征池化到 superpoint-wise 特征，從而獲取各物體的 query 表示。

我們額外設(shè)計(jì)了一個(gè)具有幾何感知能力的池化方法（Geo-aware Pooling），對(duì)于落入同一個(gè) superpoint 內(nèi)的點(diǎn)云特征，不僅僅是做 max pooling 或者 average pooling，而是考慮到 superpoint 自身的幾何形狀，給不同空間位置的特征賦予不同的權(quán)重。

2.2 Query Refinement

類似于基于 Mask2Former [5] 等基于 Transformer 的實(shí)例分割架構(gòu)，我們也利用若干層 Transformer 對(duì)每個(gè)物體的 query 信息進(jìn)行更新優(yōu)化，在每層Transformer 中，query 特征分別和當(dāng)前幀點(diǎn)云特征進(jìn)行 cross-attention，和自身進(jìn)行 self-attention，并最終和當(dāng)前幀點(diǎn)云特征交互得到優(yōu)化后的 3D Mask。

該過程有幾個(gè)不同的設(shè)計(jì)，首先是我們采用了不同層次的點(diǎn)云特征，來(lái)獲取最好的 efficiency-accuracy tradeoff。

對(duì)于 cross-attention，我們采用 superpoint-wise 特征進(jìn)行交互，superpoint 數(shù)量少但是包含了場(chǎng)景中的所有信息，有利于加速運(yùn)算。對(duì)于 3D Mask 預(yù)測(cè)，我們則將 query 和 point-wise 特征進(jìn)行交互，從而得到細(xì)粒度的分割結(jié)果。

其次，在每層 Transformer 拿到 point-wise 的分割結(jié)果后，我們會(huì)利用之前 Geo-aware Pooling 的權(quán)重將分割結(jié)果池化為 superpoint-wise，用于下一層中進(jìn)行 masked cross-attention（即通過初步的分割結(jié)果限定 cross-attention 的作用范圍）。

2.3 Query Merging

完成上述階段后，我們拿到了當(dāng)前幀優(yōu)化后的 3D Mask，基于和 3D Mask一一對(duì)應(yīng)的 query 特征，同時(shí)我們還有之前幀的全局 3D Mask。這一階段需要將當(dāng)前幀 3D Mask 與之前的 Mask 進(jìn)行全局融合，判斷當(dāng)前幀每個(gè)物體是融入到過去已存在的某個(gè)物體，還是注冊(cè)為新物體。

這一過程也是時(shí)間上最大的瓶頸，之前如 SAM3D 等手工融合的方法不僅對(duì)超參數(shù)較為敏感，而且速度極慢，能否又快又準(zhǔn)完成 3D Mask 的全局融合是一個(gè)非常有挑戰(zhàn)的問題。

這里得益于我們的 query 表示，我們可以將 3D Mask 的融合問題轉(zhuǎn)換為 query 融合問題。不同物體的 3D Mask 含點(diǎn)云數(shù)量不同，難以高效率并行計(jì)算，只能在循環(huán)中計(jì)算 3D Mask 間的空間距離設(shè)置閾值完成融合。

而不同物體的 query 尺寸是相同的，例如之前幀一共有 M 個(gè)物體，當(dāng)前幀一共 N 個(gè)物體，由于每個(gè)物體都有一一對(duì)應(yīng)的 query，因此這個(gè)問題可以通過矩陣乘法直接完成相似度計(jì)算！

假設(shè) query 特征維度為 C，我們將 MxC 和 NxC 的特征相乘便得到 MxN 的相似度矩陣，再通過二分圖匹配的方法在該矩陣上得到匹配關(guān)系。對(duì)于匹配上的物體，我們將其 3D Mask 合并，并將對(duì)應(yīng)的 query 特征加權(quán)融合；對(duì)于未匹配上的物體則直接注冊(cè)為全局新物體。

實(shí)際中，僅依靠 query 表示雖然高效，卻仍存在表示能力不夠的問題。我們?cè)诓挥绊懰俣鹊那疤嵯略O(shè)計(jì)了三種更細(xì)致的判據(jù)，通過不同的 MLP head 在 query 上預(yù)測(cè)出三種表示：1）物體的外接框；2）物體的對(duì)比表示；3）物體的語(yǔ)義表示。這三種表示具體含義如下圖所示。

打開網(wǎng)易新聞 查看精彩圖片

每次進(jìn)行 3D Mask 的融合后，這三種表示也會(huì)進(jìn)行對(duì)應(yīng)的加權(quán)融合，并作為“過去幀”再次與未來(lái)的“當(dāng)前幀”進(jìn)行高效融合。由于所有操作僅需要矩陣乘法和二分圖匹配，所需時(shí)間在 5ms 以內(nèi)，相比手工融合的策略快數(shù)百倍，且性能更強(qiáng)！

三、實(shí)驗(yàn)結(jié)果

我們進(jìn)行了多種設(shè)置的實(shí)驗(yàn)。首先是 class-agnostic 3D instance segmentation，對(duì)標(biāo) SAM3D 等文章的設(shè)置。由于這些文章大多數(shù)為 zero-shot，而我們需要訓(xùn)練，為了更公平展示泛化性，我們還進(jìn)行了 zero-shot 遷移實(shí)驗(yàn)，從 ScanNet200 遷移到 SceneNN 和 3RScan 上。

打開網(wǎng)易新聞 查看精彩圖片
打開網(wǎng)易新聞 查看精彩圖片

此外我們還進(jìn)行了 online 3D instance segmentation 的實(shí)驗(yàn)，對(duì)標(biāo)之前的在線 3D 實(shí)例分割工作。在兩種任務(wù)上我們都取得了 SOTA 的性能和速度。

打開網(wǎng)易新聞 查看精彩圖片

我們還分析了 EmbodiedSAM 的推理時(shí)間分布，發(fā)現(xiàn)速度的瓶頸主要在 3D 特征提取部分，我們的 merging 策略的開銷幾乎可以忽略不記。如何加速 3D 骨干網(wǎng)絡(luò)的特征提取是一個(gè)非常有價(jià)值的問題，也是未來(lái)突破 30FPS 的最后阻礙。

打開網(wǎng)易新聞 查看精彩圖片

更多實(shí)驗(yàn)、可視化和 Demo 視頻可以參考我們的論文以及項(xiàng)目主頁(yè)。如有問題歡迎大家在 github 上開 issue 討論。

參考文獻(xiàn)

[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.

作者：許修為 來(lái)源：公眾號(hào)【PaperWeekly】

llustration From IconScout By IconScout Store

-The End-

掃碼觀看！

本周上新！

“AI技術(shù)流”原創(chuàng)投稿計(jì)劃

TechBeat是由將門創(chuàng)投建立的AI學(xué)習(xí)社區(qū)（

www.techbeat.net
） 。 社區(qū)上線600+期talk視頻，3000+篇技術(shù)干貨文章，方向覆蓋CV/NLP/ML/Robotis等；每月定期舉辦頂會(huì)及其他線上交流活動(dòng)，不定期舉辦技術(shù)人線下聚會(huì)交流活動(dòng)。我們正在努力成為AI人才喜愛的高質(zhì)量、知識(shí)型交流平臺(tái)，希望為AI人才打造更專業(yè)的服務(wù)和體驗(yàn)，加速并陪伴其成長(zhǎng)。

投稿內(nèi)容

// 最新技術(shù)解讀/系統(tǒng)性知識(shí)分享 //

// 前沿資訊解說(shuō)/心得經(jīng)歷講述 //

投稿須知

稿件需要為原創(chuàng)文章，并標(biāo)明作者信息。

我們會(huì)選擇部分在深度技術(shù)解析及科研心得方向，對(duì)用戶啟發(fā)更大的文章，做原創(chuàng)性內(nèi)容獎(jiǎng)勵(lì)

投稿方式

發(fā)送郵件到

melodybai@thejiangmen.com

或添加工作人員微信（yellowsubbj）投稿，溝通投稿詳情；還可以關(guān)注“將門創(chuàng)投”公眾號(hào)，后臺(tái)回復(fù)“投稿”二字，獲得投稿說(shuō)明。

關(guān)于我“門”

將門是一家以專注于數(shù)智核心科技領(lǐng)域的新型創(chuàng)投機(jī)構(gòu)，也是北京市標(biāo)桿型孵化器。 公司致力于通過連接技術(shù)與商業(yè)，發(fā)掘和培育具有全球影響力的科技創(chuàng)新企業(yè)，推動(dòng)企業(yè)創(chuàng)新發(fā)展與產(chǎn)業(yè)升級(jí)。

將門成立于2015年底，創(chuàng)始團(tuán)隊(duì)由微軟創(chuàng)投在中國(guó)的創(chuàng)始團(tuán)隊(duì)原班人馬構(gòu)建而成，曾為微軟優(yōu)選和深度孵化了126家創(chuàng)新的技術(shù)型創(chuàng)業(yè)公司。

如果您是技術(shù)領(lǐng)域的初創(chuàng)企業(yè)，不僅想獲得投資，還希望獲得一系列持續(xù)性、有價(jià)值的投后服務(wù)，歡迎發(fā)送或者推薦項(xiàng)目給我“門”:

bp@thejiangmen.com

打開網(wǎng)易新聞 查看精彩圖片

點(diǎn)擊右上角，把文章分享到朋友圈