具身万物？EmbodiedSAM：实时在线分割任意3D物体

最新推荐文章于 2024-10-09 18:00:00 发布

自动驾驶之心

最新推荐文章于 2024-10-09 18:00:00 发布

阅读量535

点赞数

文章标签： 3d

原文链接：https://mp.weixin.qq.com/s?__biz=Mzg2NzUxNTU1OA==&mid=2247627080&idx=3&sn=9c9cb77128f48f3357ca76c782698dcd&chksm=cf5448c6eb3d92d72441b58ac7b88e0dc75b1043443b83df9503e6c5629831628521bc789f9f&scene=126&sessionid=0

版权

点击下方卡片，关注“自动驾驶之心”公众号

戳我-> 领取自动驾驶近15个方向学习路线

今天自动驾驶之心为大家分享清华&南洋理工的工作EmbodiedSAM！如果您有相关工作需要分享，请在文末联系我们！

自动驾驶课程学习与技术交流群事宜，也欢迎添加小助理微信AIDriver004做进一步咨询

>>点击进入→自动驾驶之心『具身智能』技术交流群

编辑 | 自动驾驶之心

写在前面 & 笔者的个人理解

在具身任务中，代理需要在探索过程中全面理解3D场景，这迫切要求开发一种在线、实时、细粒度且泛化性强的3D感知模型。然而由于高质量的3D数据相对匮乏，直接在3D空间中训练此类模型几乎是不可能的任务。幸运的是，视觉基础模型（Vision Foundation Models, VFM）已经在2D计算机视觉领域取得了革命性的进展，这为使用VFM辅助3D感知成为可能。但是，目前大多数VFM辅助的3D感知方法要么无法在线使用，要么速度太慢，难以应用于实际的具身任务中。为此本文提出了一种名为EmbodiedSAM（ESAM）的VFM辅助3D实例分割框架，旨在利用Segment Anything Model（SAM）实现实时的3D实例分割。这一挑战主要在于输入的RGB-D视频流中不包含未来的帧，且实例可能在多个帧中被观察到，因此需要在帧之间进行目标匹配。为解决这些问题，作者首先提出了一个几何感知查询提升模块，该模块将由SAM生成的2D掩码通过3D感知查询进行表示，并由一个双层查询解码器迭代细化，从而将2D掩码转换为3D点云上的细粒度形状。得益于3D掩码的查询表示，作者可以通过高效的矩阵运算计算不同视图下的3D掩码之间的相似性矩阵，实现实时推理。作者在ScanNet、ScanNet200、SceneNN和3RScan数据集上的实验表明，作者的方法在性能上达到了领先水平，即使与离线方法相比也毫不逊色。此外EmbodiedSAM在多个零样本数据集迁移实验中展现了出色的泛化能力，并在开放词汇和数据高效的场景中显示出巨大潜力。训练和评测过程仅需一块RTX 3090 GPU。

原文链接：https://arxiv.org/abs/2408.11811

图 1：与传统的3D SAM方法相比，ESAM采用2D掩码生成3D查询，并通过迭代细化来预测精确且几何上一致的3D掩码。利用3D查询，ESAM能够通过简单的矩阵操作快速合并不同帧中的3D掩码。以SAM3D为例，作者的ESAM在平均精度上超越了其15.7%，并且速度提升了20倍以上。

EmbodiedSAM方法详解

给定一系列RGB-D图像及其对应的已知姿态，作者的目标是在相应的3D场景中对任何实例进行分割。具体来说，对于每一帧，其中代表彩色图像，代表由深度图像通过姿态参数投影到3D空间得到的点云，作者的方法旨在预测对应于观察到的3D场景的实例掩码。此外，作者希望这一过程是在线进行的，即在任意时间点，作者无法预知未来的帧，但需要在每个时间点提供时间上一致的3D实例掩码。

概述：作者的方法概述如图2所示。作者以增量方式解决在线3D实例分割问题，以实现实时处理。在时间点，作者只预测当前帧的实例掩码。然后作者将与的先前实例掩码合并，得到的更新实例掩码。

图 2：展示了ESAM的工作原理。在每个新的时间点t，作者首先使用SAM生成2D实例掩码M_2d^t。作者设计了一个几何感知的查询提升模块，将2D掩码提升为3D查询，同时保留细节形状信息。然后通过双级解码器细化这些查询，以生成精细的点掩码M_cur^t。最后，将当前掩码合并到先前的掩码中，采用快速的查询合并策略。

查询提升与细化

在接收到第个 RGB-D 帧时，作者首先采用 SAM 来自动生成来自的 2D 实例掩码。在本节中，为了描述的简洁性，作者将省略下标。

几何感知的查询提升。 SAM 不使用前一帧的信息，也不利用深度图像中的 3D 信息，直接将投影到上会导致不准确的 3D 掩码，并且这些掩码在时间上也不一致。作者的目标是将每个 2D 掩码提升为一个 3D 查询特征，这样作者就可以进一步细化这些查询，以生成 3D 实例掩码。由于 2D 掩码的信息较少，作者首先从场景中提取点云特征，然后使用 2D 掩码作为索引来将点云聚类成超点，查询可以直接从超点特征中选择。设点云并且中有个掩码，作者首先根据颜色-深度对应关系将映射到上，得到超点索引，其中中的每个元素都在范围内。然后作者将点云输入到一个带有基于记忆的适配器的 3D 稀疏 U-Net 中，以提取时间感知的 3D 特征。有了和，作者可以将点特征池化到超点特征。

为了在每个超点内更好地保留点特征，作者考虑了每个超点的几何形状。对于超点，作者计算所有点相对于超点中心的归一化相对位置。这样，集合表示了这个超点的归一化形状，其直径为 1，中心位于原点。然后作者计算每个点的局部和全局特征：

其中，MLP 对每个单独的点执行操作，Agg 是通过通道最大池化实现的聚合函数。局部和全局特征表示点与形状的相关性，因此作者将这两个特征结合起来，并通过另一个 MLP 来预测点权重：

最后，作者使用加权平均池化将点特征聚合到第个超点：

通过增强聚合后的超点特征，以完全结合形状级几何特征和场景级 3D U-Net 特征。每个超点的计算可以并行化，使用逐点 MLP 和 Scatter 函数，使得这种几何感知池化实际上是高效的。

双级查询解码器。 在池化之后，中的个 2D 实例掩码被提升为 3D 超点特征。然后作者从初始化一系列 3D 实例查询，并通过几个基于 Transformer 的查询解码器层迭代细化，用于预测 3D 掩码。在训练期间，作者随机采样的一部分（0.5 到 1 之间）来构建进行数据增强。而在推理时，作者简单地设置。

每个查询解码器都采用蒙版交叉注意力机制，在查询和场景表示之间聚合每个查询的实例信息：

其中，点乘表示矩阵乘法，是查询解码器层的线性投影，和是场景表示的线性投影。可以是点特征或超点特征。是根据预测的 3D 实例掩码派生的注意力掩码。表示第个查询关注第个点或超点。然后作者将输入到自注意力层和前馈网络中，得到，接着通过掩码生成模块预测每个查询的实例掩码：

其中是线性层。如果，是点掩码；否则，它是超点掩码。

查询解码器通常采用相同级别的场景表示来进行交叉注意力和掩码生成。但由于 SAM 已经输出了高级别语义感知掩码，作者观察到。如果采用点特征进行查询解码，交叉注意力操作将因点的数量大而消耗大量内存。而如果作者使用超点特征，预测的 3D 实例掩码将只是超点的组合，因此无法细化到更细粒度。为了两者兼得，作者的查询解码器设计为双级。在交叉注意力中，作者设置以实现高效的交互；而在掩码预测时，作者设置以实现细粒度掩码生成。为了支持蒙版注意力，作者在交叉注意力之前将点掩码池化到超点掩码：

其中是几何感知池化。作者可以重用预先计算的权重来减少计算。通过 3 次查询解码器迭代后，作者获得了准确的点掩码以及相应的查询，在后续小节中表示为和。作者在上执行掩码 NMS 以过滤掉冗余掩码及其相应的查询。

高效在线查询合并

在完成了2D掩码到3D掩码的提升之后，接下来作者需要将这些掩码合并到先前的实例掩码中，以生成更新的实例掩码。对于的初始情况，作者有等于。然而，以往的实例掩码合并方法通常需要遍历中的所有掩码，并且与中的每个掩码进行逐一比较，这是一个非常耗时的过程。为了准确判断新掩码是否应该合并到先前的掩码中，这些方法需要计算两个掩码点云之间的几何相似性，例如交并比(IoU)或最近点距离，这些计算通常涉及每个掩码的所有点，导致计算复杂度高且难以并行化。

为了解决这个问题，作者提出了一种新的方法，通过固定大小的向量来表示每个掩码，并通过高效的矩阵运算来计算相似性。得益于作者的架构，和中的每个掩码都有对应的查询特征。查询特征本身是固定大小的向量表示，但直接计算它们之间的相似性可能信息量不足。因此，作者设计了几项辅助任务，基于查询特征学习不同度量的向量表示，用于计算几何、对比和语义相似性。

首先，对于几何相似性，作者发现模型能够通过部分观察学习到整个几何结构。但由于分割的限制，模型无法完全表达对整个几何形状的理解。因此，作者引入了一个辅助任务，使用多层感知器(MLP)基于每个查询的中心（即对应超点的中心）来预测边界框回归，得到边界框。然后，可以通过计算两个边界框的IoU来得到两个掩码之间的几何相似性。作者忽略了边界框的方向，因为两个轴对齐边界框集合之间的IoU矩阵可以通过简单的矩阵运算来计算。

第二，对于对比相似性，作者的目标是学习一种实例特定的表示，使得同一实例的特征应该被拉近，不同实例的特征则相互远离。这种表示可以通过相邻帧之间的对比训练来学习：作者使用MLP将查询特征映射到对比特征。对于在第帧和第帧中出现的实例，作者选择该实例的两个特征作为正例对，并从其他实例中采样特征作为负例对。详细的损失函数将在下一小节中展示。

最后，对于语义相似性，作者简单地采用MLP预测每个预定义类别的类别概率分布，其中是预定义类别的数量。也有其他的选择，例如，如果作者采用语义SAM代替SAM，作者可以直接使用2D掩码的语义预测作为相应查询的。

通过上述方式，作者可以高效地计算和之间的相似性矩阵：

其中，表示两组轴对齐边界框之间的IoU矩阵。作者通过将小于阈值的元素设置为来剪枝矩阵C。然后在和上执行代价为的二分图匹配，该匹配将中的每个掩码分配给中的一个掩码。如果一个新掩码未能与任何先前掩码匹配，作者将为该掩码注册一个新的实例。否则，作者将两个掩码及其对应的边界框B、对比特征f和语义分布S进行合并。掩码合并可以通过简单的并集实现，而其他表示则通过加权平均进行合并：

作者假设第个新掩码合并到第个先前掩码中。这里是合并计数，表示已合并到的掩码数量。

图 3：详细介绍了作者高效的查询合并策略。作者设计了三种代表性的辅助任务，分别生成几何、对比和语义的向量表示。然后，相似性矩阵可以通过矩阵乘法高效地计算。作者进一步剪枝相似性矩阵，并采用二分图匹配来合并实例。

损失函数

在作者的模型训练中，每个RGB-D帧都配备了语义和实例的标签。在一个RGB-D视频序列中，不同帧之间对应的实例标签是一致的。基于这些标签，作者针对每帧计算损失，这些损失是基于每个查询的预测结果来确定的。由于查询是直接从2D SAM掩码中一对一提升而来的，作者可以直接使用2D掩码的标签来监督对应查询的预测，省去了复杂的标签分配步骤。作者假设每个2D SAM掩码仅对应一个实例，因此可以为每个查询获取对应的真实语义标签和2D实例掩码。利用深度图像的像素对应关系，作者将2D实例掩码映射到3D点云上，并据此计算出真实的轴对齐边界框。利用这些标注信息，作者定义了以下损失函数：

分类损失：通过交叉熵计算，用以区分前景和背景实例的二元分类损失。
3D掩码损失：包括二元交叉熵损失和Dice损失，用于指导预测的3D掩码。
边界框损失：基于预测的边界框与真实边界框之间的IoU来计算损失。
语义预测损失：采用二元交叉熵损失来预测每个实例的语义类别。

除了每帧内的损失外，作者还引入了相邻帧之间的对比损失：

其中，表示余弦相似度，是帧中实例的总数。

最终，总损失函数被定义为：

这里，和被设置为0，以保证序列两端的边界条件。

实验结果和分析

在这一节中，作者将首先描述作者的实验设置和实现细节。然后，作者将比较作者的方法与现有的基于视觉基础模型(VFM)辅助的3D实例分割方法和在线3D分割方法，以验证其有效性。作者还将展示EmbodiedSAM(ESAM)在开放词汇和数据高效设置中的应用潜力。最后，通过消融研究，作者将提供对作者设计的综合分析。

基准测试与实现细节

作者的方法在以下四个数据集上进行了评估：ScanNet、ScanNet200、SceneNN和3RScan。ScanNet数据集包含1513个室内场景的3D重建，作者选取了其中的1201个序列用于训练，余下的312个用于测试。ScanNet200提供了更为细致的注释，涵盖了200多个类别。SceneNN由50个高质量的3D扫描场景组成，每个场景都带有实例和语义标签。依照文献，作者挑选了12个清洁序列用于测试。3RScan数据集更具挑战性，其RGB-D序列是通过快速移动的相机捕获的，作者使用了该数据集的测试分割部分，共包含46个场景。这些数据集都提供了带有标注的姿态RGB-D序列和重建的点云。

基准测试：作者首先在表1中与基于VFM的3D实例分割方法进行了比较。作者在ScanNet200的训练集上训练了不同的方法（如果需要），并在ScanNet200的验证集上以类别不可知的方式进行了评估。对于离线方法，输入是整个重建的点云和RGB-D视频，预测是在点云上直接进行的。对于在线方法，输入是连续的RGB-D视频流，作者将最终预测结果映射到重建的点云上，以便进行比较。

表 1：展示了不同方法在ScanNet200数据集上进行类别不可知的3D实例分割的结果。作者与常规聚类方法和基于VFM的3D场景感知方法进行了比较。速度单位为每帧的毫秒数，其中VFM和其他部分的速度分别报告。

由于某些方法如SAM3D无需训练即可使用，为了与可学习的方法进行公平的比较，作者还进一步在表2中评估了表1中的方法在SceneNN和3RScan数据集上的表现，且未进行微调。这一基准测试展示了不同方法的泛化能力。

表 2：展示了不同方法从ScanNet200迁移到SceneNN和3RScan数据集的结果。作者直接在这些数据集上评估了表1中的模型，以展示它们的泛化能力。

作者也在表3中与在线3D实例分割方法进行了比较。按照先前的研究，作者在ScanNet训练集上训练了不同的方法，并在ScanNet的验证集和SceneNN上进行了评估。

表 3：展示了不同方法在ScanNet和SceneNN数据集上的3D实例分割结果。

最后，作者在表4中评估了ESAM在ScanNet200数据集上198个类别的开放词汇3D实例分割能力。由于ESAM输出的是类别不可知的3D掩码，作者可以通过两种方式将其扩展到开放词汇的3D分割。第一种是将这些掩码输入到像OpenMask3D这样的开放词汇掩码分类模型中。第二种是使用开放词汇的2D分割模型来为每个2D掩码获取类别标签。由于ESAM中3D掩码和2D掩码之间存在一一对应关系，作者可以相应地为每个3D掩码获取类别标签。这里作者遵循SAI3D的方法，采用了第一种方式，并与之作了比较。

表 4：展示了在ScanNet200数据集上进行开放词汇3D实例分割的结果。

实现细节：作者按照文献的描述，分两个阶段训练ESAM。首先，在ScanNet(200)-25k子集上训练单视图感知模型，该子集包含单独的RGB-D帧，不包括基于记忆的适配器和三个辅助任务的损失。然后，在RGB-D序列上微调单视图感知模型，加入了适配器和全部损失。为了降低内存占用，作者在每次迭代中随机采样每个场景的8个相邻RGB-D帧。在超参数设置上，作者选择了，，，，和。

与SOTA比较

作者对比了三种版本的ESAM：标准版ESAM、采用FastSAM以提高推理速度的ESAM-E，以及ESAM-E+FF，后者不仅使用了FastSAM的2D掩码，还将其特征与点云数据进行了融合。作者也提供了一些可视化结果来定性评估不同方法的性能。

根据表1的结果，在3D实例分割的“segment anything task”中，ESAM在与先前方法的比较中表现出色，甚至超越了一些离线处理方法。值得注意的是，在线方法相较于离线方法在感知3D场景时面临更大的挑战，因为它们需要处理的是部分和噪声更大的视频帧。尽管如此，ESAM不仅在准确性上取得了领先，而且在推理速度上也有显著优势，每帧的处理时间仅为80毫秒，远快于需要超过1秒的传统方法。当使用FastSAM替代SAM后，ESAM-E能够在保持较高准确性的同时，实现约10 FPS的实时在线3D实例分割。

在泛化能力方面，ESAM同样表现优异。如表2所示，ESAM在未进行微调的情况下直接应用于其他数据集，其准确性依然领先于其他零样本方法。此外，作者还注意到，SAI3D在3RScan数据集上的表现不如SAM3D，这可能是因为SAI3D高度依赖于清晰且精确对齐的3D网格和RGB图像，而3RScan数据集中由于相机移动速度快，图像和姿态信息较为模糊。

图4展示了不同3D实例分割方法在ScanNet200数据集上的可视化结果。ESAM能够生成准确且细节丰富的3D实例分割掩码，并且能够实时处理连续的RGB-D视频流。作者还提供了在线可视化结果（见图5），进一步证明ESAM在实际应用中的有效性。更多细节请参考作者提供的视频演示。

图 4：展示了不同3D实例分割方法在ScanNet200数据集上的可视化结果。在红色框中突出显示了SAM3D预测的噪声掩码，以及SAI3D倾向于将单一实例过度分割为多个部分的问题。

图 5：展示了ESAM在ScanNet200数据集上的在线可视化结果。更多细节可以参考作者项目页面上的视频演示。

在表3和表4中，ESAM与现有的在线3D实例分割方法和开放词汇3D实例分割方法相比，也展现出了领先的性能。

ESAM分析

数据高效的学习。 在表5中，作者展示了ESAM在仅使用部分训练集时的性能表现。即使训练样本数量减少，ESAM的分类不可知性能在ScanNet200上的下降也相对有限。这表明，得益于2D VFM提供的优良初始化，ESAM的学习部分能够快速收敛。

表 5：展示了ESAM在仅使用部分训练集时的性能。作者报告了在ScanNet200上，使用不同比例的训练集大小训练时的类别不可知平均精度。

推理时间的分解。 表6进一步分解了ESAM的推理时间（不包括VFM部分）。结果显示，由于设计高效，ESAM的解码器和合并操作仅占用了总推理时间的一小部分。

表 6：展示了ESAM的推理时间分解（不包括VFM）。详细列出了3D U-Net适配器、解码器、相似性计算、匹配和更新各部分的时间消耗。

消融研究。 作者还对ESAM进行了消融研究，以验证所提出方法的有效性。在架构设计方面，作者在ScanNet200-25k上进行了实验，并在表7中报告了每种方法的类别不可知平均精度（AP）和平均推理延迟（不包括SAM）。结果显示，几何感知池化能够将性能提升1.4%，同时几乎不增加计算开销。作者还发现，ESAM的双级设计在保持与使用点特征FP策略相当的准确性的同时，相比使用超点特征FS的策略，仅略微增加了推理延迟。在合并策略方面，作者在表8中展示了不同设计对ScanNet200上平均精度的影响。结果表明，每个辅助任务对于掩码合并的质量都至关重要。特别是，几何相似性对最终性能有显著影响，因为大多数掩码对可以根据它们之间的距离被快速排除。

表 7：展示了ESAM架构设计的影响。在ScanNet200-25k上进行实验，并报告了每种设计选择的类别不可知平均精度和每帧的平均推理延迟。

表 8：展示了ESAM合并策略的影响。在ScanNet200上进行实验，并报告了不同设计选择对平均精度的影响。

辅助任务的可视化。 图6中的可视化结果进一步展示了作者的辅助任务预测。图6(a)显示了模型如何基于部分观察预测物体的完整几何形状。图6(b)中的t-SNE可视化验证了模型成功学习了用于实例匹配的区分性查询表示。最后，图6(c)中的语义分割结果表明，ESAM能够学习到令人满意的语义表示，并且可以扩展到3D语义分割任务。

图 6：展示了作者合并策略中辅助任务的可视化结果。(a) 展示了不同时间点上物体的3D边界框预测，用于几何相似性。(b) 展示了实例特定表示的t-SNE可视化，用于对比相似性，不同颜色代表不同的实例，不同点代表不同帧中的实例特征。

结论

在本项工作中，作者成功地提出了ESAM，这是一个高效的框架，它借助视觉基础模型（VFM）来实现在线、实时、细粒度、高度泛化，以及开放词汇的3D实例分割。作者的方法通过几何感知池化将2D掩码提升为3D查询，并利用双路径查询解码器进一步细化这些查询，以生成精确的3D实例掩码。此外，利用查询与掩码之间的对应关系，作者设计了三种辅助任务，用以在三种不同的向量空间中表示每个3D掩码的特征，从而实现高效的掩码合并。通过在四个不同的数据集上进行广泛的实验，作者证明了ESAM在性能、在线实时推理能力以及泛化能力方面均达到了领先水平。ESAM还显示出了在开放词汇识别和数据高效学习场景中的应用潜力。作者认为，ESAM为有效利用2D VFM进行具身感知任务提供了新的视角和方法。

ESAM的潜在局限性。 尽管ESAM在多个方面表现出色，但仍存在一些限制。首先，ESAM能否实现实时处理取决于所选用的VFM。目前，作者采用了SAM和FastSAM作为VFM，但只有FastSAM能够达到实时处理的要求。作者预期，随着未来2D VFM在效率和性能上的进一步提升，ESAM也将随之得到改进。其次，ESAM中的3D U-Net和基于记忆的适配器在特征提取方面相对较为复杂，占据了大部分的推理时间。如果能够优化这些组件，提高后端的效率，ESAM的处理速度将有潜力达到更高水平，这是作者未来研究工作的重点方向。

参考

[1] EmbodiedSAM: Online Segment Any 3D Thing in Real Time

投稿作者为『自动驾驶之心知识星球』特邀嘉宾，欢迎加入交流！重磅，自动驾驶之心科研论文辅导来啦，申博、CCF系列、SCI、EI、毕业论文、比赛辅导等多个方向，欢迎联系我们！

① 全网独家视频课程

BEV感知、BEV模型部署、BEV目标跟踪、毫米波雷达视觉融合、多传感器标定、多传感器融合、多模态3D目标检测、车道线检测、轨迹预测、在线高精地图、世界模型、点云3D目标检测、目标跟踪、Occupancy、cuda与TensorRT模型部署、大模型与自动驾驶、Nerf、语义分割、自动驾驶仿真、传感器部署、决策规划、轨迹预测等多个方向学习视频（扫码即可学习）

网页端官网：www.zdjszx.com

② 国内首个自动驾驶学习社区

国内最大最专业，近3000人的交流社区，已得到大多数自动驾驶公司的认可！涉及30+自动驾驶技术栈学习路线，从0到一带你入门自动驾驶感知（2D/3D检测、语义分割、车道线、BEV感知、Occupancy、多传感器融合、多传感器标定、目标跟踪）、自动驾驶定位建图（SLAM、高精地图、局部在线地图）、自动驾驶规划控制/轨迹预测等领域技术方案、大模型、端到端等，更有行业动态和岗位发布！欢迎扫描下方二维码，加入自动驾驶之心知识星球，这是一个真正有干货的地方，与领域大佬交流入门、学习、工作、跳槽上的各类难题，日常分享论文+代码+视频