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，就可以在三维场景中分割出相应物体，以体素的形式表示。
目录
- 一. 研究思路
- 二. mask 逆向渲染
- 三. 跨视角自动 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 初始渲染的图像为 I in \mathbf{I}^{\text {in}} Iin，用户输入的 prompt 点集为 P in \mathcal{P}^{\text {in}} Pin，SAM 根据 prompt 点集分割的 2D mask 为 M SAM in \mathbf{M}_{\text {SAM }}^{\text {in}} MSAM in。先通过 mask 逆向渲染机制将 M SAM in \mathbf{M}_{\text {SAM}}^{\text {in}} MSAMin 投影到三维空间得到 3D mask V \mathbf{V} V，再将 3D mask 投影到二维空间中新的视角得到粗略的 2D mask M ( n ) \mathbf{M}^{(n)} M(n)，然后通过跨视角自动 prompt 机制从 M ( n ) \mathbf{M}^{(n)} M(n) 中继续提取 prompt 点 P ( n ) \mathcal{P}^{(n)} P(n)，最后使用 SAM 根据 P ( n ) \mathcal{P}^{(n)} P(n) 分割出更加精确的 2D mask M SAM ( n ) \mathbf{M}_{\text {SAM}}^{(n)} MSAM(n)。重复上面的操作，直到获得完整的 3D mask。
二. mask 逆向渲染
因为 SA3D 的分割结果是用体素 V ∈ R L × W × H \mathbf{V} \in \mathbb{R}^{L \times W \times H} V∈RL×W×H 表示的， L × W × H L \times W \times H L×W×H 表示三维场景的长宽高，所以 mask 逆向渲染就是在三维空间中找到对应的 3D mask。每个体素有一个置信度分数 (soft mask confidence score)，表示属于 3D mask 的概率，初始化时将所有体素的置信度置零，体素 r ( t ) \mathbf{r}(t) r(t) 的置信度用 V ( r ( t ) ) \mathbf{V}(\mathbf{r}(t)) V(r(t)) 表示。
3D mask 向二维空间投影时，2D mask M \mathbf{M} M 的置信度为：
M ( r ) = ∫ t n t f ω ( r ( t ) ) V ( r ( t ) ) d t \mathbf{M}(\mathbf{r})=\int_{t_n}^{t_f} \omega(\mathbf{r}(t)) \mathbf{V}(\mathbf{r}(t)) \mathrm{dt} M(r)=∫tntfω(r(t))V(r(t))dt
将 M ( r ) \mathbf{M}(\mathbf{r}) M(r) 和 SAM 分割得到的真实 mask M SAM ( n ) \mathbf{M}_{\text {SAM}}^{(n)} MSAM(n) 计算损失来优化 3D mask 的置信度：
L proj = − ∑ r ∈ R ( I ) M S A M ( r ) ⋅ M ( r ) \mathcal{L}_{\text {proj }}=-\sum_{\mathbf{r} \in \mathcal{R}(\mathbf{I})} \mathbf{M}_{\mathrm{SAM}}(\mathbf{r}) \cdot \mathbf{M}(\mathbf{r}) Lproj =−r∈R(I)∑MSAM(r)⋅M(r)
由于初始阶段 mask 存在较大误差，因此在损失函数中增加一个负的细化项，根据多视图 mask 一致性来优化 3D mask：
L proj = − ∑ r ∈ R ( I ) M S A M ( r ) ⋅ M ( r ) + λ ∑ r ∈ R ( I ) ( 1 − M S A M ( r ) ) ⋅ M ( r ) \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}) Lproj =−r∈R(I)∑MSAM(r)⋅M(r)+λr∈R(I)∑(1−MSAM(r))⋅M(r)
三. 跨视角自动 prompt
跨视角自动 prompt 就是要在不同视角上自动采样 prompt 点，从而提供更多的 2D mask 来优化 3D mask。对于每个视角，3D mask 投影到该平面得到粗略的 2D mask M ( n ) \mathbf{M}^{(n)} M(n)，使用 Self-prompting 方法从 M ( n ) \mathbf{M}^{(n)} M(n) 中提取一组 prompt 点 P s ( n ) \mathcal{P}_{s}^{(n)} Ps(n)。这里的 M ( n ) \mathbf{M}^{(n)} M(n) 并不是二位图，而是该像素的置信度，下面使用像素 p \mathbf{p} p 阐述 Self-prompting 的过程。
P s ( n ) \mathcal{P}_{s}^{(n)} Ps(n) 初始化为空集，第一个点 p 0 \mathbf{p}_0 p0 选择具有最高置信度分数的点，即 p 0 = a r g m a x P M ( n ) ( p ) \mathbf{p}_0=\mathrm{arg max}_\mathbf{P} \mathbf{M}^{(n)}(\mathbf{p}) p0=argmaxPM(n)(p)。后续新增的 prompt 点需要同时满足与现有 prompt 点靠近和具有较高的置信度分数，因此引入了一个衰减项来调整每个像素的置信度分数：
Δ M ( n ) ( p ) = min { M ( n ) ( p ^ ) ⋅ d ( G ( p ) , G ( p ^ ) ) ∣ p ^ ∈ P s } \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\} ΔM(n)(p)=min{M(n)(p^)⋅d(G(p),G(p^))∣p^∈Ps}
其中 d ( ⋅ , ⋅ ) d(·,·) d(⋅,⋅) 表示 min-max normalized Euclidean distance。计算现有 prompt 点周围点衰减后的置信度分数，选择最高的加入 P s ( n ) \mathcal{P}_{s}^{(n)} Ps(n)：
M ~ ( n ) ( p ) = M ( n ) ( p ) − Δ M ( n ) ( p ) \tilde{\mathbf{M}}^{(n)}(\mathbf{p})=\mathbf{M}^{(n)}(\mathbf{p})-\Delta \mathbf{M}^{(n)}(\mathbf{p}) M~(n)(p)=M(n)(p)−ΔM(n)(p)
四. 实验
五. 总结
因为三维场景是由 NeRF 隐式表达的，所以 SA3D 在表示分割出的物体时需要引入新的表达 —— 体素。SA3D 的思路和 SAM3D 1 相似，通过不同视角之间交替进行 mask 逆向渲染和跨视角 Self-prompting，迭代地完成基于体素网格构建的 3D mask。由于是 NeRF 的隐式表达，SA3D 在优化时需要一定时间，从几十秒到几分钟不等，难以做到实时的分割。2
Yang Y, Wu X, He T, et al. Sam3d: Segment anything in 3d scenes[J]. arXiv preprint arXiv:2306.03908, 2023. ↩︎
Segment Anything Model 在 3D 视觉方向上有哪些应用？ ↩︎