SA3D：基于 NeRF 的三维场景分割方法

Paper: Cen J, Zhou Z, Fang J, et al. Segment anything in 3d with nerfs[J]. Advances in Neural Information Processing Systems, 2024, 36.
Introduction: https://jumpat.github.io/SA3D/
Code: https://github.com/Jumpat/SegmentAnythingin3D

SA3D 是一种用于 NeRF 表示的三维场景的分割技术，只需要用户在单一视角的图像上给出粗略的分割点作为 prompt，就可以在三维场景中分割出相应物体，以体素的形式表示。

一. 研究思路

SA3D 根据用户在单一视图上给出的 prompt，使用 SAM 提取 mask 以筛选场景中的体素，从而分割出场景中的三维目标物体。分割三维场景元素的过程如下：在不同视角之间交替进行 mask 逆向渲染 和 跨视角自动 prompt，迭代地完成基于体素网格构建的 3D mask：

• mask 逆向渲染 (mask inverse rendering)：将当前视角 SAM 分割的 2D mask 逆向投影到三维空间，以获取 3D mask；
• 跨视角自动 prompt (cross-view self-prompting)：在其他视角下自动提取 prompt 点传入 SAM 提取 2D mask；

记 NeRF 初始渲染的图像为 ${\mathbf{I}}^{\text{in}}$，用户输入的 prompt 点集为 ${\mathcal{P}}^{\text{in}}$，SAM 根据 prompt 点集分割的 2D mask 为 ${\mathbf{M}}_{\text{SAM }}^{\text{in}}$。先通过 mask 逆向渲染机制将 ${\mathbf{M}}_{\text{SAM}}^{\text{in}}$ 投影到三维空间得到 3D mask $\mathbf{V}$，再将 3D mask 投影到二维空间中新的视角得到粗略的 2D mask ${\mathbf{M}}^{(n)}$，然后通过跨视角自动 prompt 机制从 ${\mathbf{M}}^{(n)}$ 中继续提取 prompt 点 ${\mathcal{P}}^{(n)}$，最后使用 SAM 根据 ${\mathcal{P}}^{(n)}$ 分割出更加精确的 2D mask ${\mathbf{M}}_{\text{SAM}}^{(n)}$。重复上面的操作，直到获得完整的 3D mask。

二. mask 逆向渲染

因为 SA3D 的分割结果是用体素 $\mathbf{V}\in {\mathbb{R}}^{L\times W\times H}$ 表示的，$L\times W\times H$ 表示三维场景的长宽高，所以 mask 逆向渲染就是在三维空间中找到对应的 3D mask。每个体素有一个置信度分数 (soft mask confidence score)，表示属于 3D mask 的概率，初始化时将所有体素的置信度置零，体素 $\mathbf{r}(t)$ 的置信度用 $\mathbf{V}(\mathbf{r}(t))$ 表示。

3D mask 向二维空间投影时，2D mask $\mathbf{M}$ 的置信度为：
$$\mathbf{M}(\mathbf{r})={\int }_{t_n}^{t_f}\omega (\mathbf{r}(t))\mathbf{V}(\mathbf{r}(t))\mathrm{dt}$$

将 $\mathbf{M}(\mathbf{r})$ 和 SAM 分割得到的真实 mask ${\mathbf{M}}_{\text{SAM}}^{(n)}$ 计算损失来优化 3D mask 的置信度：
$${\mathcal{L}}_{\text{proj }}=-\sum _{\mathbf{r}\in \mathcal{R}(\mathbf{I})}{\mathbf{M}}_{\mathrm{SAM}}(\mathbf{r})\cdot \mathbf{M}(\mathbf{r})$$

由于初始阶段 mask 存在较大误差，因此在损失函数中增加一个负的细化项，根据多视图 mask 一致性来优化 3D mask：
$${\mathcal{L}}_{\text{proj }}=-\sum _{\mathbf{r}\in \mathcal{R}(\mathbf{I})}{\mathbf{M}}_{\mathrm{SAM}}(\mathbf{r})\cdot \mathbf{M}(\mathbf{r})+\lambda \sum _{\mathbf{r}\in \mathcal{R}(\mathbf{I})}\left(1-{\mathbf{M}}_{\mathrm{SAM}}(\mathbf{r})\right)\cdot \mathbf{M}(\mathbf{r})$$

三. 跨视角自动 prompt

跨视角自动 prompt 就是要在不同视角上自动采样 prompt 点，从而提供更多的 2D mask 来优化 3D mask。对于每个视角，3D mask 投影到该平面得到粗略的 2D mask ${\mathbf{M}}^{(n)}$，使用 Self-prompting 方法从 ${\mathbf{M}}^{(n)}$ 中提取一组 prompt 点 ${\mathcal{P}}_s^{(n)}$。这里的 ${\mathbf{M}}^{(n)}$ 并不是二位图，而是该像素的置信度，下面使用像素 $\mathbf{p}$ 阐述 Self-prompting 的过程。

${\mathcal{P}}_s^{(n)}$ 初始化为空集，第一个点 ${\mathbf{p}}_0$ 选择具有最高置信度分数的点，即 ${\mathbf{p}}_0={\mathrm{argmax}}_{\mathbf{P}}{\mathbf{M}}^{(n)}(\mathbf{p})$。后续新增的 prompt 点需要同时满足与现有 prompt 点靠近和具有较高的置信度分数，因此引入了一个衰减项来调整每个像素的置信度分数：
$$\Delta {\mathbf{M}}^{(n)}(\mathbf{p})=\min \left\{{\mathbf{M}}^{(n)}(\hat{\mathbf{p}})\cdot d(\mathcal{G}(\mathbf{p}),\mathcal{G}(\hat{\mathbf{p}}))\mid \hat{\mathbf{p}}\in {\mathcal{P}}_s\right\}$$

其中 $d(\cdot ,\cdot )$ 表示 min-max normalized Euclidean distance。计算现有 prompt 点周围点衰减后的置信度分数，选择最高的加入 ${\mathcal{P}}_s^{(n)}$：
$${\tilde{\mathbf{M}}}^{(n)}(\mathbf{p})={\mathbf{M}}^{(n)}(\mathbf{p})-\Delta {\mathbf{M}}^{(n)}(\mathbf{p})$$

四. 实验

五. 总结

因为三维场景是由 NeRF 隐式表达的，所以 SA3D 在表示分割出的物体时需要引入新的表达 —— 体素。SA3D 的思路和 SAM3D ¹ 相似，通过不同视角之间交替进行 mask 逆向渲染和跨视角 Self-prompting，迭代地完成基于体素网格构建的 3D mask。由于是 NeRF 的隐式表达，SA3D 在优化时需要一定时间，从几十秒到几分钟不等，难以做到实时的分割。²