SAGA
SAGA (Segment Any 3D GAussians)的官方实现。
安装
SAGA的安装类似于3D Gaussian Splatting。
git clone git@github.com:Jumpat/SegAnyGAussians.git
或
git clone https://github.com/Jumpat/SegAnyGAussians.git
然后安装依赖项:
conda env create --file environment.yml
conda activate gaussian_splatting
默认情况下,我们使用SAM的公开ViT-H模型。您可以从这里下载预训练模型,并将其放在./third_party/segment-anything/sam_ckpt目录下。
准备数据
使用的数据集包括360_v2、nerf_llff_data和LERF。
SAGA的数据结构如下所示:
./data
/360_v2
/garden
/images
/images_2
/images_4
/images_8
/sparse
/features
/sam_masks
/mask_scales
...
/nerf_llff_data
/fern
/images
/poses_bounds.npy
/sparse
/features
/sam_masks
/mask_scales
/horns
...
...
/lerf_data
...
由于我们需要预训练的3D-GS模型来提取掩码尺度,第一步是训练3D高斯:
预训练3D高斯
我们继承了3DGS的所有属性,有关训练高斯的更多信息可以在他们的仓库中找到。
python train_scene.py -s <COLMAP或NeRF Synthetic数据集的路径>
准备数据
然后,要获取sam_masks和相应的掩码尺度,运行以下命令:
python extract_segment_everything_masks.py --image_root <场景数据的路径> --sam_checkpoint_path <预训练SAM模型的路径> --downsample <1/2/4/8>
python get_scale.py --image_root <场景数据的路径> --model_path <预训练3DGS模型的路径>
注意,由于GPU内存限制,有时降采样是必要的。
如果您想尝试开放词汇分割,首先提取CLIP特征:
python get_clip_features.py --image_root <场景数据的路径>
训练3D高斯亲和力特征
python train_contrastive_feature.py -m <预训练3DGS模型的路径> --iterations 10000 --num_sampled_rays 1000
3D分割
目前,SAGA提供了一个使用dearpygui实现的交互式GUI(saga_gui.py)和一个jupyter笔记本(prompt_segmenting.ipynb)。要运行GUI:
python saga_gui.py --model_path <预训练3DGS模型的路径>
暂时,开放词汇分割仅在jupyter笔记本中实现。请参考prompt_segmenting.ipynb获取详细说明。
GUI使用:
设置GUI后,您可以看到以下界面:
视角控制:
左键拖动
:旋转。中键拖动
:平移。右键点击
:输入点提示(需要先选中分割模式)。
分割控制:
超参数选项:
scale
:3D尺度(用于分割和聚类)。score thresh
:分割相似度阈值(用于分割)。
渲染选项:
RGB
:在特定视角显示当前3D-GS模型的原始RGB。PCA
:在特定视角显示当前3D-GS模型3D特征的PCA分解结果。SIMILARITY
:显示给定点提示的相似度图(需要先输入提示)。3D CLUSTER
:显示当前3D-GS模型的3D聚类结果。
分割模式选项:
click mode
:在此模式下只能输入一个点。multi-click mode
:在此模式下可以输入多个点以同时选择多个对象。preview_segmentation_in_2d
:显示当前输入提示(点、尺度和分数阈值)的2D分割结果。注意,2D分割结果可能与3D结果不一致。
分割选项:
选择感兴趣的目标后,您可以点击segment3D
获取3D分割结果。如果结果不满意,您可以点击roll back
撤销此分割,或点击clear
回到未分割状态,或者您可以继续输入提示,基于临时分割结果进行分割。您可以点击save as
将当前分割结果保存在./segmentation_res/your_name.pt
中,这是3D-GS模型中所有3D高斯的二进制掩码。
聚类选项:
您可以随时点击cluster3d
获取当前3D-GS模型的聚类结果。例如,您可以直接对整个场景进行聚类,或在临时分割的对象中进行聚类以进行分解。点击reshuffle_cluster_color
可以重新随机分配聚类的渲染颜色。
注意,直接对整个场景进行聚类可能需要一些时间,因为为了方便起见,我们使用没有GPU支持的HDBSCAN。
渲染
在交互式GUI中保存分割结果或运行prompt_segmenting.ipynb中的脚本后,高斯的位图将保存在./segmentation_res/your_name.pt
中(您可以自行设置名称)。要在训练视图上渲染分割结果(通过移除背景获得分割对象),运行以下命令:
python render.py -m <预训练3DGS模型的路径> --precomputed_mask <分割结果的路径> --target scene --segment
要获取2D渲染掩码,运行以下命令:
python render.py -m <预训练3DGS模型的路径> --precomputed_mask <分割结果的路径> --target seg
您也可以在不进行分割的情况下渲染预训练的3DGS模型:
python render.py -m <预训练3DGS模型的路径> --target scene
引用
如果您发现这个项目对您的研究有帮助,请考虑引用报告并给予⭐。
@article{cen2023saga,
title={Segment Any 3D Gaussians},
author={Jiazhong Cen and Jiemin Fang and Chen Yang and Lingxi Xie and Xiaopeng Zhang and Wei Shen and Qi Tian},
year={2023},
journal={arXiv preprint arXiv:2312.00860},
}
致谢
SAGA的实现参考了GARField、OmniSeg3D和Gaussian Splatting,我们衷心感谢他们对社区的贡献。