SCube：使用 VoxSplats 实现实时大规模场景重建

SCube：使用 VoxSplats 实现实时大规模场景重建
Xuanchi Ren$^* $, Yifan Lu$^* $, Hanxue Liang, Zhangjie Wu, Huan Ling, Mike Chen, Sanja Fidler, Francis Williams, Jiahui Huang
* 表示共同第一作者
[项目页面]

最新消息

• 2024-12-11：也请查看我们的最新研究 InfiniCube。
• 2024-12-11：代码已发布！我们还在此仓库中集成了 XCube，以便更方便地使用。有关详细说明，请参阅 XCube 使用说明。

安装

mkdir SCube-release # 将其包裹在一个文件夹中
cd SCube-release
git clone https://github.com/nv-tlabs/SCube.git 
mkdir wandb
mkdir waymo_tfrecords 
cd SCube

(可选) 安装 libMamba，以大幅提升使用 Conda 时的体验

conda update -n base conda
conda install -n base conda-libmamba-solver
conda config --set solver libmamba

安装 Conda 环境

conda env create -f environment.yml
conda activate scube
mim install "mmcv>=2.0.0"
pip install "mmsegmentation>=1.0.0"

设置 WandB

我们的训练脚本高度依赖于 WandB。请先注册一个 WandB 账户并获取您的 API_key。然后您可以在终端中运行以下命令来为您的机器进行设置：

wandb login # 需要您的 API 密钥

数据准备

数据拆分下载

首先下载我们的 Waymo 数据集拆分。

gdown https://drive.google.com/drive/folders/1d8V4aoHsDf5U0-SH_VY2MOZJdYvOG8TM -O ../waymo_split --folder

Waymo 数据集下载

从 Waymo 感知数据集 v1.4.2 下载所有训练和验证片段，并将其保存到 SCube-release/waymo_tfrecords。

如果您有 sudo 权限，可以使用 gcloud 从终端下载它们。

sudo apt-get update
sudo apt-get install apt-transport-https ca-certificates gnupg curl
curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo gpg --dearmor -o /usr/share/keyrings/cloud.google.gpg
echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" | sudo tee -a /etc/apt/sources.list.d/google-cloud-sdk.list
sudo apt-get update && sudo apt-get install google-cloud-cli

# 或者使用 `gcloud init --no-launch-browser` 如果您在远程终端会话中
gcloud init 
bash datagen/download_waymo.sh datagen/waymo_all.json ../waymo_tfrecords

下载 tfrecord 文件后，我们期望的文件夹结构如下：

SCube-release
|-- SCube
|   `-- ...
|-- waymo_split
|   `-- ...
`-- waymo_tfrecords
    |-- segment-10247954040621004675_2180_000_2200_000_with_camera_labels.tfrecord
    |-- segment-11379226583756500423_6230_810_6250_810_with_camera_labels.tfrecord
    |-- ...
    `-- segment-1172406780360799916_1660_000_1680_000_with_camera_labels.tfrecord

[!NOTE] 如果您从控制台下载 tfrecord 文件，文件名前可能会带有 individual_files_training_ 或 individual_files_validation_ 等前缀。请确保在进一步处理之前移除这些前缀。

数据处理

我们首先从 tfrecord 文件中提取图像、LiDAR、位姿、内参和标注文件，并将它们保存到 ../waymo_ns。如果您有多核 CPU，可以设置 --num_workers 为您拥有的核心数以实现并行化。

python datagen/waymo2nerfstudio.py -i ../waymo_tfrecords -o ../waymo_ns [--num_workers YOUR_CPU_CORE]

然后我们将原始数据处理成 SCube 训练所需的 webdataset 格式，同时使用 SegFormer 推断天空掩码以及 Metric3Dv2 获取 GT 深度。结果将保存到 ../waymo_webdataset。请注意，这一步对于 1000 个片段 × 198 帧 × 5 个摄像头来说非常耗时，在 8 张 A100 GPU 上可能需要超过一天的时间。

[!NOTE] 在使用 mmcv 时，您可能会遇到 AssertionError: MMCV==2.2.0 is used but incompatible. Please install mmcv>=2.0.0rc4. 的错误。解决方案在这里。

[!NOTE] 在使用 Metric3D v2 时，您可能会遇到 RuntimeError: Failed to find function: mono.model.backbones.vit_large_reg 的错误。解决方案在这里。

# 单 GPU
python datagen/nerfstudio2webdataset.py -i ../waymo_ns -o ../waymo_webdataset

# 多 GPU
torchrun --nnodes=1 --nproc-per-node=8 'datagen/nerfstudio2webdataset.py' -i ../waymo_ns -o ../waymo_webdataset

现在 ../waymo_ns 和 ../waymo_tfrecords 将不再使用。我们期望的文件夹结构如下：

SCube-release
|-- SCube
|   `-- ...
|-- waymo_split
|   `-- ...
`-- waymo_webdataset
    `-- ...

GT 体素下载

我们在 Hugging Face 上提供了地面真值体素的 webdataset 格式文件。请将其下载并解压缩到 SCube-release/waymo_webdataset 中，与其他属性一起存放。

要可视化以点云格式存储的地面真值体素数据，您可以运行以下命令：

python inference/visualize_gt_pc.py -p <GT_VOXEL_TAR_FILE>

模型训练

VAE 训练

# 粗略阶段 VAE
python train.py configs/waymo_scube/train_vae_64x64x64_dense_height_down2_residual.yaml --wname train_vae_64x64x64_dense_height_down2_residual --max_epochs 10 --gpus 8 --eval_interval 1

# 精细阶段 VAE
python train.py configs/waymo_scube/train_vae_256x256x128_sparse.yaml --wname train_vae_256x256x128_sparse --max_epochs 10 --gpus 8 --eval_interval 1

几何重建训练（潜在体素扩散模型）

[!IMPORTANT]
您需要在 configs/waymo_scube/train_diffusion_64x64x64_image_cond.yaml 和 configs/waymo_scube/train_diffusion_256x256x128_sparse.yaml 中将 vae_checkpoint 修改为您自己的实验。在您的 WandB 页面上，您可以找到一个唯一的“run id”，并将其替换为 yaml 文件中的 [YOUR_RUN_ID]。

# 图像条件下的体素扩散
python train.py configs/waymo_scube/train_diffusion_64x64x64_image_cond.yaml --wname train_diffusion_64x64x64_image_cond --max_epochs 40 --gpus 8 --eval_interval 1

# 语义条件下的上采样体素扩散
python train.py configs/waymo_scube/train_diffusion_256x256x128_sparse.yaml --wname train_diffusion_256x256x128_sparse --max_epochs 20 --gpus 8 --eval_interval 1

外观重建训练（GSM，Gaussian Splatting Model）

python train.py configs/waymo_scube/train_gsm_unet3d_view3.yaml --wname train_gsm_unet3d_view3 --max_epochs 30 --gpus 8 --eval_interval 1

模型推理

VAE 推理

# 将 nvidia-toronto 替换为您的 W&B 账号！
python inference/vae.py none --ckpt_vae wdb:nvidia-toronto/scube-scene-recon/waymo_wds/train_vae_64x64x64_dense_height_down2_residual
python inference/vae.py none --ckpt_vae wdb:nvidia-toronto/scube-scene-recon/waymo_wds/train_vae_256x256x128_sparse

这将生成一个文件夹 ../vae_output_waymo_wds/，存储预测的和真实的网格及语义信息。请通过 --ckpt_vae 指定实验名称。

扩散模型推理

# 将 nvidia-toronto 替换为您的 W&B 账号！
python inference/diffusion.py none --use_ema --use_ddim --ddim_step 100 --ckpt_dm wdb:nvidia-toronto/scube-scene-recon/waymo_wds/train_diffusion_64x64x64_image_cond
python inference/diffusion.py none --use_ema --use_ddim --ddim_step 100 --ckpt_dm wdb:nvidia-toronto/scube-scene-recon/waymo_wds/train_diffusion_256x256x128_sparse

这将生成一个文件夹 ../diffusion_output_waymo_wds/，存储预测的和真实的带有语义信息的体素网格。请通过 --ckpt_dm 指定实验名称。

GSM 推理

# 将 nvidia-toronto 替换为您的 W&B 账号！
python inference/gaussian_feedforward.py none --ckpt_gsm wdb:nvidia-toronto/scube-scene-recon/waymo_wds/train_gsm_unet3d_view3 [--save_gs]

这将生成一个文件夹 ../diffusion_output_waymo_wds/，存储预测的和真实的网格及语义信息。请通过 --ckpt_gsm 指定实验名称。添加 [--save_gs] 可以保存 3D 高斯分布。

完整流水线推理

# 将 nvidia-toronto 替换为您的 W&B 账号！
python inference/diffusion_cascading_gsm.py none \
    --use_ema --use_ddim --ddim_step 100 \
    --ckpt_dm_c wdb:nvidia-toronto/scube-scene-recon/waymo_wds/train_diffusion_64x64x64_image_cond \
    --ckpt_dm_f wdb:nvidia-toronto/scube-scene-recon/waymo_wds/train_diffusion_256x256x128_sparse \
    --ckpt_gsm wdb:nvidia-toronto/scube-scene-recon/waymo_wds/train_gsm_unet3d_view3 \
    --input_frame_offsets [0] \
    --sup_frame_offsets [0,5,10] \
    --split val \
    --val_starting_frame 0

可视化推理结果

网格可视化

在 ../vae_output_waymo_wds/、../diffusion_output_waymo_wds/ 和 ../cascading_diffusion_output_waymo_wds/ 中，您可以找到一些包含 x.pt 和 x_gt.pt 的文件夹，其中 x 是您推理的样本索引。运行以下命令，即可在网页浏览器中启动交互式可视化工具（即使是在远程 SSH 机器上）。

python inference/visualize_grid_compare.py -p <PT_FILES_FOLDER> [-t voxel]

其中 -t 可以是 voxel 或 pc。请勿在精细阶段使用 voxel，因为渲染负担过重。在使用 VSCode 连接远程机器时，可以轻松完成此操作；VSCode 会自动为您转发端口。

3DGS 可视化

在 ../splat_output_waymo_wds 和 cascading_diffusion_output_waymo_wds 中，您可以找到一些包含 x_splat.pkl 的文件夹，其中 x 是您推理的样本索引。您可以通过以下命令可视化 3DGS：

python inference/visualize_3dgs_pkl.py -p <3DGS_PKL_FILE>

快捷方式

您可以安装该项目以创建可视化快捷方式。

poetry install 

安装后，您将拥有 vis-pair 命令，等同于 python inference/visualize_grid_compare.py；vis-gs 等同于 python inference/visualize_3dgs_pkl.py；vis-gt-pc 等同于 inference/visualize_gt_pc.py。您可以这样运行可视化：

vis-pair -p <PT_FILES_FOLDER>
vis-gs -p <3DGS_PKL_FILE>
vis-gt-pc -p <GT_VOXEL_TAR_FILE>

许可证

版权所有 © 2024，NVIDIA Corporation 及其附属公司。保留所有权利。 本作品根据 Nvidia 源代码许可证 提供。

相关工作

• Lu 等，2024 年。InfiniCube：基于世界引导视频模型的无界且可控的动态 3D 驾驶场景生成。
• Ren 等，2024 年。XCube：利用稀疏体素层次结构进行大规模 3D 生成建模。
• Williams 等，2024 年。𝑓VDB：用于稀疏、大规模和高性能空间智能的深度学习框架。

致谢

本仓库基于 https://github.com/nv-tlabs/XCube。

引用

@inproceedings{
    ren2024scube,
    title={SCube：使用 VoxSplats 实现即时的大规模场景重建},
    author={Ren, Xuanchi and Lu, Yifan and Liang, Hanxue and Wu, Jay Zhangjie and Ling, Huan and Chen, Mike and Fidler, Sanja annd Williams, Francis and Huang, Jiahui},
    booktitle={第 38 届神经信息处理系统年度会议},
    year={2024},
}

SCube 快速上手指南

SCube 是一个基于 VoxSplats 技术，用于即时大规模场景重建的开源工具。本指南将帮助中国开发者快速完成环境配置并运行基础推理。

1. 环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐 Ubuntu 20.04+)
• GPU: NVIDIA GPU (建议显存 ≥ 24GB，如 A100/A6000/3090/4090)，已安装 CUDA 驱动。
• 包管理器: Conda (推荐安装 libmamba 求解器以加速依赖解析)。
• 账号: Weights & Biases (WandB) 账号（用于模型权重管理和实验追踪）。
• 网络: 由于项目依赖 Google Drive、Hugging Face 和 Waymo 数据集，国内用户建议配置科学网络环境或使用代理加速下载。

2. 安装步骤

2.1 克隆代码与目录初始化

mkdir SCube-release
cd SCube-release
git clone https://github.com/nv-tlabs/SCube.git 
mkdir wandb
mkdir waymo_tfrecords 
cd SCube

2.2 配置 Conda 环境

为提升依赖安装速度，强烈建议启用 libmamba 求解器：

conda update -n base conda
conda install -n base conda-libmamba-solver
conda config --set solver libmamba

创建并激活虚拟环境，安装核心依赖：

conda env create -f environment.yml
conda activate scube
mim install "mmcv>=2.0.0"
pip install "mmsegmentation>=1.0.0"

注意: 若遇到 mmcv 版本兼容性问题，请参考官方文档或相关 Issue 进行微调。

2.3 配置 WandB

登录您的 WandB 账号以获取 API Key 并完成本地认证：

wandb login
# 按提示输入您的 API Key

2.4 数据准备（简述）

完整训练需要下载 Waymo 数据集并进行复杂预处理。若仅需快速测试推理，您可以跳过繁琐的数据采集步骤，直接从 Hugging Face 下载预处理的 Ground Truth Voxel 数据用于验证可视化流程：

1. 访问 SCube-Data Hugging Face 页面。
2. 下载文件并解压至 SCube-release/waymo_webdataset 目录。

(注：完整的数据集下载与 TFRecord 转换流程较长，涉及 Waymo 官网注册及大量磁盘 IO，初次上手建议先熟悉推理流程。)

3. 基本使用

以下示例展示如何使用官方预训练模型进行全链路推理（从体素生成到高斯泼溅重建）。

3.1 运行全链路推理

请确保将命令中的 nvidia-toronto 替换为您自己的 WandB 用户名（如果使用了官方公开权重且权限允许，可直接尝试；否则需先在 WandB 上复现或申请权限）。

python inference/diffusion_cascading_gsm.py none \
    --use_ema --use_ddim --ddim_step 100 \
    --ckpt_dm_c wdb:nvidia-toronto/scube-scene-recon/waymo_wds/train_diffusion_64x64x64_image_cond \
    --ckpt_dm_f wdb:nvidia-toronto/scube-scene-recon/waymo_wds/train_diffusion_256x256x128_sparse \
    --ckpt_gsm wdb:nvidia-toronto/scube-scene-recon/waymo_wds/train_gsm_unet3d_view3 \
    --input_frame_offsets [0] \
    --sup_frame_offsets [0,5,10] \
    --split val \
    --val_starting_frame 0

执行成功后，结果将保存在 ../cascading_diffusion_output_waymo_wds/ 目录下。

3.2 可视化结果

项目提供了便捷的可视化工具。推荐安装 Poetry 以使用快捷命令：

poetry install

查看体素网格对比 (Voxel Grid):

vis-pair -p <PT_FILES_FOLDER>
# 或者直接使用 python:
# python inference/visualize_grid_compare.py -p <PT_FILES_FOLDER> [-t voxel]

提示：在远程服务器使用时，VSCode 会自动转发端口，可直接在浏览器中交互查看。

查看 3D 高斯泼溅 (3DGS):

vis-gs -p <3DGS_PKL_FILE>
# 或者直接使用 python:
# python inference/visualize_3dgs_pkl.py -p <3DGS_PKL_FILE>

查看 Ground Truth 点云:

vis-gt-pc -p <GT_VOXEL_TAR_FILE>

通过以上步骤，您即可完成 SCube 的基础部署并观察到重建效果。如需进行模型训练，请参考原文档中的 "Model Training" 章节配置相应的 YAML 文件。