使用NeRFs实现Segment Anything🤖️的3D分割(SA3D)

项目主页 | arXiv论文

使用NeRFs在3D空间中实现Segment Anything
Jiazhong Cen1*, Zanwei Zhou1*, Jiemin Fang2,3†, Chen Yang1, Wei Shen1✉, Lingxi Xie2, Dongsheng Jiang2, Xiaopeng Zhang2, Qi Tian2
1上海交通大学AI研究所 2华为公司 3华中科技大学电子信息与通信学院
*表示贡献相同
†表示项目负责人

给定一个NeRF模型，只需从单一视角输入提示，即可获得3D模型。
<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/9abe479c-9c07-479d-a1ee-f9acd1ecab7f.gif" width="800">

我们提出了一种名为SA3D的新型框架，用于在3D空间中实现Segment Anything。给定一个神经辐射场(NeRF)模型，SA3D允许用户通过在单个渲染视图中一次性手动提示，获得任意目标对象的3D分割结果。获取目标3D模型的整个过程可以在约2分钟内完成，且无需任何工程优化。我们的实验证明了SA3D在不同场景中的有效性，凸显了SAM在3D场景感知中的潜力。

更新

2024/04/16：我们发布了SA3D的3D-GS版本(点击此处)。现在3D分割可以在几秒钟内完成！
2023/11/11：我们发布了SA3D的nerfstudio版本(点击此处)！目前仅支持文本提示作为输入。
2023/06/29：我们现在支持使用MobileSAM作为分割网络。按照MobileSAM中的安装说明进行操作，然后将mobile_sam.pt下载到./dependencies/sam_ckpt文件夹中。您可以使用--mobile_sam切换到MobileSAM。

整体流程

SA3D_pipeline

通过输入提示，SAM从相应视图中分割出目标对象。获得的2D分割掩码通过密度引导的逆渲染投影到3D掩码网格上。然后渲染其他视图的2D掩码，这些掩码大多不完整，但作为跨视图自提示再次输入SAM。可以获得完整的掩码并投影到掩码网格上。这个过程以迭代方式执行，最终可以学习到准确的3D掩码。SA3D可以有效地适应各种辐射场，无需任何额外的重新设计。

安装

git clone https://github.com/Jumpat/SegmentAnythingin3D.git
cd SegmentAnythingin3D

conda create -n sa3d python=3.10
conda activate sa3d
pip install -r requirements.txt

SAM和Grounding-DINO：

# 安装SAM
mkdir dependencies; cd dependencies 
mkdir sam_ckpt; cd sam_ckpt
wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth
git clone git@github.com:facebookresearch/segment-anything.git 
cd segment-anything; pip install -e .

# 安装Grounding-DINO
git clone https://github.com/IDEA-Research/GroundingDINO.git
cd GroundingDINO/; pip install -e .
mkdir weights; cd weights
wget https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth

下载数据

我们现在发布了这些数据集的配置：

前向视角： LLFF
内向视角： mip-NeRF360, LERF

数据结构：

data
├── 360_v2             # 链接: https://jonbarron.info/mipnerf360/
│   └── [bicycle|bonsai|counter|garden|kitchen|room|stump]
│       ├── poses_bounds.npy
│       └── [images|images_2|images_4|images_8]
│
├── nerf_llff_data     # 链接: https://drive.google.com/drive/folders/14boI-o5hGO9srnWaaogTU5_ji7wkX2S7
│   └── [fern|flower|fortress|horns|leaves|orchids|room|trex]
│       ├── poses_bounds.npy
│       └── [images|images_2|images_4|images_8]
│
└── lerf_data               # 链接: https://drive.google.com/drive/folders/1vh0mSl7v29yaGsxleadcj-LCZOE_WEWB
    └── [book_store|bouquet|donuts|...]
        ├── transforms.json
        └── [images|images_2|images_4|images_8]

</details>

使用方法

训练NeRF

python run.py --config=configs/llff/fern.py --stop_at=20000 --render_video --i_weights=10000

在GUI中运行SA3D

python run_seg_gui.py --config=configs/llff/seg/seg_fern.py --segment \
--sp_name=_gui --num_prompts=20 \
--render_opt=train --save_ckpt

渲染并保存飞行视频

python run_seg_gui.py --config=configs/llff/seg/seg_fern.py --segment \
--sp_name=_gui --num_prompts=20 \
--render_only --render_opt=video --dump_images \
--seg_type seg_img seg_density

运行SA3D时的一些提示：

当目标物体极不规则时（如LLFF场景中的蕨类和霸王龙），请增加 --num_prompts 的值；
使用 --seg_poses 指定用于训练3D掩码的相机姿态序列，默认值为'train'，可选项为['train', 'video']。

使用我们基于Dash的图形界面：

选择要使用的提示类型，目前支持：点提示和文本提示；
- 点提示： 在下拉菜单中选择点；点击原始图像添加点提示，然后SAM将生成候选掩码；点击清除点可清除之前的输入；
- 文本提示： 在下拉菜单中选择文本；输入您的文本提示并点击生成以获取候选掩码；注意，不合理的文本输入可能会导致错误。
选择目标掩码；
按下开始训练运行SA3D；我们可视化由跨视图自提示策略生成的渲染掩码和SAM预测结果；
等待几分钟即可看到最终渲染结果。

待办事项

完善图形界面，例如，从任何训练视图开始，添加更多训练超参数选项等；
在图形界面中支持两阶段处理；目前可能存在一些bug。

一些可视化样本

SA3D可以处理各种场景的3D分割。在我们的项目页面上查看更多演示。

前向场景	360°	多目标
<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/48003de5-d228-4a73-b649-350e22a2b3fa.gif" width="200">	<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/deee1c12-4724-450b-9e39-8c004028f1c9.gif" width="200">	<img src="https://yellow-cdn.veclightyear.com/0a4dffa0/52b0c957-989c-46ae-a425-fb4a0344c9c4.gif" width="200">

致谢

感谢以下项目的宝贵贡献：

引用

如果您发现这个项目对您的研究有帮助，请考虑引用该报告并给予一个⭐。

@inproceedings{cen2023segment,
      title={Segment Anything in 3D with NeRFs}, 
      author={Jiazhong Cen and Zanwei Zhou and Jiemin Fang and Chen Yang and Wei Shen and Lingxi Xie and Dongsheng Jiang and Xiaopeng Zhang and Qi Tian},
      booktitle    = {NeurIPS},
      year         = {2023},
}