Segment Anything 3D

作者：Yunhan Yang、Xiaoyang Wu、Tong He、Hengshuang Zhao、Xihui Liu
机构：上海人工智能实验室、香港大学
技术报告：[arxiv]

我们通过将2D图像的分割信息转移到3D空间，将Segment Anything扩展到3D感知领域。我们期望分割信息能够有助于3D传统感知和开放世界感知。该项目仍在进行中，并将被嵌入到我们的感知代码库Pointcept中。我们非常欢迎任何问题或拉取请求。

结果

示例网格可在此处获取。

安装

conda create -n sam3d python=3.8 -y
conda activate sam3d
# 在此选择你想要的版本：https://pytorch.org/get-started/previous-versions/
conda install pytorch==1.11.0 torchvision==0.12.0 torchaudio==0.11.0 cudatoolkit=11.3 -c pytorch
conda install plyfile -c conda-forge -y
pip install scikit-image opencv-python open3d imageio
pip install git+https://github.com/facebookresearch/segment-anything.git 

cd libs/pointops
# 常规
python setup.py install
# docker和多GPU架构
TORCH_CUDA_ARCH_LIST="ARCH LIST" python  setup.py install
# 例如 7.5: RTX 3000; 8.0: a100 更多信息请参见：https://developer.nvidia.com/cuda-gpus
TORCH_CUDA_ARCH_LIST="7.5 8.0" python  setup.py install
cd ../..

数据准备

ScanNet v2

下载ScanNet v2数据集。
按如下步骤运行原始ScanNet的预处理代码：

准备点云数据（遵循Pointcept）

# RAW_SCANNET_DIR：下载的ScanNet v2原始数据集目录。
# PROCESSED_SCANNET_DIR：处理后的ScanNet数据集目录（输出目录）。
python scannet-preprocess/preprocess_scannet.py --dataset_root ${RAW_SCANNET_DIR} --output_root ${PROCESSED_SCANNET_DIR}

准备RGBD数据（遵循BPNet）

python scannet-preprocess/prepare_2d_data/prepare_2d_data.py --scannet_path data/scannetv2 --output_path data/scannetv2_images --export_label_images

开始使用

请通过sam3d.py尝试

# RGB_PATH：rgb数据的路径
# DATA_PATH：点云数据的路径
# SAVE_PATH：保存pcd结果的位置
# SAVE_2DMASK_PATH：保存SAM生成的2D分割结果的位置
# SAM_CHECKPOINT_PATH：SAM检查点的路径

python sam3d.py --rgb_path $RGB_PATH --data_path $DATA_PATH --save_path $SAVE_PATH --save_2dmask_path $SAVE_2DMASK_PATH --sam_checkpoint_path $SAM_CHECKPOINT_PATH

流程

我们的SAM3D流程如下：

SAM生成掩码
使用SAM在2D帧上获取分割掩码，然后通过深度信息将它们映射到3D空间。

合并两个相邻的点云
使用"双向群组重叠算法"（修改自ContrastiveSceneContexts）合并两个相邻的点云。

区域合并方法
通过区域合并方法合并整个点云。

合并两个分割结果
我们使用默认参数对场景应用Felzenswalb和Huttenlocher的基于图的图像分割算法。详情请参考原始仓库。然后合并两个分割结果得到最终结果（合并代码在sam3d.py/pcd_ensemble中）。

引用

如果您发现_SAM3D_对您的研究有用，请引用我们的工作：

@misc{yang2023sam3d,
      title={SAM3D: Segment Anything in 3D Scenes}, 
      author={Yunhan Yang, Xiaoyang Wu, Tong He, Hengshuang Zhao and Xihui Liu},
      year={2023},
      eprint={2306.03908},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}