taichi_3d_gaussian_splatting

由 Taichi 语言实现的论文《用于实时辐射场渲染的 3D 高斯泼溅》（3D Gaussian Splatting for Real-Time Radiance Field Rendering）的非官方版本。

3D 高斯泼溅的作用是什么？

训练：

该算法以多视角图像、稀疏点云和相机位姿作为输入，利用可微分光栅化器对点云进行训练，最终输出带有额外特征（协方差、颜色信息等）的稠密点云。

如果将训练过程视为一个模块，可以这样描述：

graph LR
    A[多视角图像] --> B((训练))
    C[稀疏点云] --> B
    D[相机位姿] --> B
    B --> E[带额外特征的稠密点云]

推理：

该算法以带额外特征的稠密点云和任意相机位姿作为输入，使用相同的光栅化器从给定的相机位姿渲染出图像。

graph LR
    C[带额外特征的稠密点云] --> B((推理))
    D[新相机位姿] --> B
    B --> E[图像]

推理结果示例：

https://github.com/wanmeihuali/taichi_3d_gaussian_splatting/assets/18469933/cc760693-636b-4157-ae85-33813f3da54d

由于点云的良好特性，与其他类似 NeRF 的算法相比，该算法在处理场景或物体合并时更加容易。

https://github.com/wanmeihuali/taichi_3d_gaussian_splatting/assets/18469933/bc38a103-e435-4d35-9239-940e605b4552

为什么选择 Taichi？

• Taichi 是一种用于高性能计算的语言。它旨在弥合以生产力为导向的 Python 语言与以性能和并行性为导向的 C++/CUDA 语言之间的差距。通过使用 Taichi，本项目完全基于 Python 实现，却能达到与 CUDA 实现相当甚至更好的性能。此外，代码也更容易阅读和维护。
• Taichi 提供多种后端，包括 CUDA、OpenGL、Metal 等。我们计划未来切换后端以支持更多平台，但目前该项目仅支持 CUDA 后端。
• Taichi 提供自动微分功能，尽管当前项目尚未使用，但这对于未来的开发来说是一个非常有用的功能。

当前状态

目前，该项目已在官方实现提供的数据集上进行了测试。对于卡车数据集，本项目仅使用官方实现 1/5 到 1/4 的点数，便能获得略高于官方实现的 PSNR 值。然而，训练和推理的速度仍然慢于官方实现。

官方实现和本实现的结果均在同一数据集上测试。我注意到官方实现的结果与其论文中的结果略有不同，这可能是由于测试分辨率的差异所致。

卡车（3万）最新最佳结果：

安装

1. 准备包含 PyTorch 和 torchvision 的环境。
2. 克隆本仓库并进入目录。
3. 运行以下命令：

pip install -r requirements.txt
pip install -e .

所有依赖项均可通过 pip 安装。PyTorch/Torchvision 也可以通过 conda 安装。代码已在 Ubuntu 20.04.2 LTS 上使用 Python 3.10.10 测试过。硬件配置为 RTX 3090 显卡和 CUDA 12.1。虽然尚未在其他平台上测试，但经过少量修改后应该也能在其他平台上运行。

数据集

该算法需要整个场景的点云、相机参数以及真实标签图像。点云以 parquet 格式存储，相机参数和真实标签图像以 JSON 格式存储。运行配置则以 YAML 格式存储。我们提供了一个从 Colmap 输出构建数据集的脚本，也可以直接从原始数据构建数据集。

在 Tank and Temple 卡车场景上训练

├── data
│   ├── tat_truck_every_8_test
│   │   ├── train.json
│   │   ├── val.json
│   │   ├── point_cloud.parquet
├── image
│   ├── 000000.png
│   ├── 000001.png

python gaussian_point_train.py --train_config config/tat_truck_every_8_test.yaml

在示例物体（靴子）上训练

├── data
│   ├── boots_super_sparse
│   │   ├── boots_train.json
│   │   ├── boots_val.json
│   │   ├── point_cloud.parquet
├── image
│   ├── images_train
│   │   ├── COS_Camera.001.png
│   │   ├── COS_Camera.002.png
|   |   ├── ...

使用Colmap生成的数据集进行训练

使用带有额外网格的Instant-NGP格式数据集进行训练

python tools/prepare_InstantNGP_with_mesh.py \
    --transforms_train {训练变换文件路径} \
    --transforms_test {验证变换文件路径，若未提供，则从训练集中采样验证集} \
    --mesh_path {网格文件路径} \
    --mesh_sample_points {在网格上采样的点数} \
    --val_sample {若从训练集中采样验证集，则每隔n帧采样一次} \
    --image_path_prefix {图像路径前缀，通常是包含图像文件夹的目录路径} \
    --output_path {输出文件夹路径}

python gaussian_point_train.py --train_config {配置yaml文件路径}

使用BlenderNerf生成的数据集进行训练

python tools/prepare_InstantNGP_with_mesh.py \
    --transforms_train {transform_train.json文件路径} \
    --mesh_path {stl文件路径} \
    --mesh_sample_points {在网格上采样的点数，默认为500} \
    --val_sample {若从训练集中采样验证集，则每隔n帧采样一次，默认为8} \
    --image_path_prefix {包含train文件夹的目录的绝对路径} \
    --output_path {任意你希望的路径}

python gaussian_point_train.py --train_config {配置yaml文件路径}

使用其他方法生成的数据集进行训练

运行

python gaussian_point_train.py --train_config {配置文件路径}

训练过程的工作流程如下：

stateDiagram-v2
    state WeightToTrain {
        sparsePointCloud
        pointCloudExtraFeatures
    }
    WeightToTrain --> Rasterizer: 输入
    cameraPose --> Rasterizer: 输入
    Rasterizer --> Loss: 光栅化后的图像
    ImageFromMultiViews --> Loss
    Loss --> Rasterizer: 梯度
    Rasterizer --> WeightToTrain: 梯度

结果会在TensorBoard中可视化。TensorBoard日志存储在配置文件中指定的输出目录中。训练后的带特征点云也会以parquet格式存储，其输出目录同样由配置文件指定。

在 Colab 上运行（以利用 Google 提供的 GPU 加速器）

相关笔记本可在以下链接找到：/tools/run_3d_gaussian_splatting_on_colab.ipynb

1. 在 Colab 中设置硬件加速器：“Runtime->Change Runtime Type->Hardware accelerator->选择 GPU->选择 T4”
2. 将此仓库上传至您的 Google Drive 的相应文件夹。
3. 将您的 Google Drive 挂载到笔记本中（参见笔记本）。
4. 安装 condacolab（参见笔记本）。
5. 使用 pip 安装 requirement.txt 中的依赖项（参见笔记本）。
6. 使用 conda 安装 PyTorch、torchvision、pytorch-cuda 等（参见笔记本）。
7. 按照 https://github.com/wanmeihuali/taichi_3d_gaussian_splatting#dataset 中的说明准备数据集。
8. 使用正确的配置运行训练器（参见笔记本）。
9. 通过 TensorBoard 查看训练过程（参见笔记本）。

可视化

提供了一个简单的可视化工具。该工具由 Taichi GUI 实现，帧率被限制在 60 FPS（如果有人知道如何解除这一限制，请联系我）。该可视化工具可以接收一个或多个 Parquet 格式的输出结果。示例 Parquet 文件可在此下载：https://drive.google.com/file/d/12-kZZay8RFlDk7hJQysG_Cr4-oxDp37l/view?usp=sharing。

python3 visualizer --parquet_path_list <parquet_path_0> <parquet_path_1> ...

该可视化工具会将多个点云合并，并在同一场景中显示。

• 按 0 键选择所有点云（默认状态）。
• 按 1 到 9 键选择其中一个点云。
• 当所有点云都被选中时，使用 WASD 和 - 键移动相机，使用 QE 键绕 Y 轴旋转，或直接拖动鼠标进行自由旋转。
• 当仅选择一个点云时，使用 WASD 和 - 键移动对象或场景，使用 QE 键绕 Y 轴旋转对象或场景，或直接拖动鼠标以对象中心为轴进行自由旋转。

如何贡献/使用 CI 在云端训练

我现在已经启用了 CI 和基于云的训练功能。该功能目前还不太稳定，但它允许即使没有 GPU 的人也能为这个仓库做出贡献。 通常的工作流程如下：

1. 对于任何算法改进，请创建一个新的分支并提交拉取请求。
2. 请在拉取请求中 @wanmeihuali，我会检查代码并为拉取请求添加 need_experiment、need_experiment_garden 或 need_experiment_tat_truck 标签。
3. CI 会自动构建 Docker 镜像并将其上传到 AWS ECR，随后触发云端训练。训练结果会以评论的形式上传到拉取请求中，例如 this PR。数据集是使用 colmap 的默认配置生成的。训练在 g4dn.xlarge Spot 实例（NVIDIA T4，性能弱于 3090/A6000）上进行，通常需要 2-3 小时。
4. 目前 README.md 中的最佳训练结果是由人工更新的。未来我会尝试自动化这一过程。

当前的实现基于我对论文的理解，与论文或官方实现可能会存在一些差异（他们计划在七月发布代码）。作为一个个人项目，参数尚未经过充分调优。我将在未来继续改进性能。如果您有任何问题，欢迎随时提出 issue；也欢迎提交 PR，尤其是针对性能提升的贡献。

待办事项

算法部分

• 修复自适应控制器部分，目前致密化过程存在问题，且论文中的描述非常模糊。需要进一步实验来确定正确或更好的实现方式。
  • 确定致密化是否应应用于所有点，还是仅限于当前帧中的点。
  • 弄清楚“视空间位置梯度的平均 magnitude”具体指什么，是跨帧平均，还是跨像素平均？
  • 确定正确的分割策略。新点的位置应该放在哪里？目前的位置是优化前的位置。是否应该将其放置在原始椭球体的焦点上更好？ 对于过度重建的情况，使用概率密度函数采样；对于欠重建的情况，则使用优化前的位置。
• 在 README.md 中添加结果评分和图像
  • 尝试使用论文中相同的数据集。
  • 修复当前 Blender 插件中的问题，并将插件开源。
• 相机位姿优化：获取相机位姿的梯度，并在训练过程中对其进行优化。
• 支持动态刚体对象。目前的实现已经支持在一个场景中使用多个相机位姿，因此刚体对象的运动可以转换为相机的运动。需要寻找一种能够为不同物体提供 6 自由度位姿估计的 SfM 解决方案，并修改数据集代码以进行测试。

工程部分

• 修复 bug：当相机中没有点时程序会崩溃。
• 添加仅推理框架，以支持在场景中添加/移动物体、场景合并、场景编辑等功能。
• 添加安装脚本/Docker 镜像。
• 支持批量训练。目前代码仅支持单张图像训练，且只使用了 GPU 内存的一小部分。
• 使用 Taichi 实现基数排序/Cumsum，而不是依赖 PyTorch。PyTorch-Taichi 的张量转换似乎仅在 CUDA 设备上可用。如果要切换到其他设备，就需要摆脱 PyTorch。
• 实现仅使用 Taichi 场的纯 Taichi 推理光栅化器，并迁移到 macOS、Android 和 iOS 平台。

taichi_3d_gaussian_splatting 快速上手指南

taichi_3d_gaussian_splatting 是一个基于 Taichi Lang 实现的 3D Gaussian Splatting 非官方版本。它利用纯 Python 代码实现了高性能的可微分光栅化，用于从多视角图像训练稠密点云并进行实时渲染。相比官方 CUDA 实现，该版本代码更易读且易于维护，同时在部分场景下能以更少的点数量达到更高的 PSNR。

环境准备

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

• 操作系统: 推荐 Ubuntu 20.04.2 LTS (其他 Linux 发行版或 Windows 可能需要微调)。
• Python 版本: 推荐 Python 3.10.10。
• 硬件要求: NVIDIA GPU (测试环境为 RTX 3090)，需安装对应的 CUDA 驱动 (测试环境为 CUDA 12.1)。
  • 注意：目前主要支持 CUDA 后端。
• 前置依赖:
  • PyTorch
  • TorchVision

建议先通过 Conda 创建虚拟环境并安装 PyTorch：

conda create -n taichi_splatting python=3.10
conda activate taichi_splatting
# 请访问 pytorch.org 获取适合你 CUDA 版本的安装命令，例如：
# pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121

国内加速提示：如果下载 PyTorch 或 pip 包速度较慢，推荐使用清华源或阿里源。

export PIP_INDEX_URL=https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

克隆仓库并安装项目依赖：

1. 克隆代码库并进入目录：

   git clone https://github.com/wanmeihuali/taichi_3d_gaussian_splatting.git
   cd taichi_3d_gaussian_splatting
   

2. 安装 Python 依赖包及项目本身：

   pip install -r requirements.txt
   pip install -e .
   

基本使用

本项目最核心的功能是训练和推理。以下以官方提供的 Tank and Temple Truck 数据集为例，展示最简单的训练流程。

1. 准备数据

你需要准备图像、相机姿态（JSON 格式）和稀疏点云（Parquet 格式）。

• 下载预处理数据：可以从 Google Drive 链接 下载预处理的 Truck 场景数据。

• 目录结构：将下载的图片放入根目录下的 image 文件夹，将相机姿态和点云文件放入 data/tat_truck_every_8_test 文件夹。最终结构应如下所示：

  ├── data
  │   ├── tat_truck_every_8_test
  │   │   ├── train.json
  │   │   ├── val.json
  │   │   ├── point_cloud.parquet
  ├── image
  │   ├── 000000.png
  │   ├── 000001.png
  │   └── ...
  

  (注：如果你有自己的 COLMAP 输出数据，可以使用 tools/prepare_colmap.py 脚本转换为上述格式)

2. 开始训练

使用提供的配置文件启动训练过程。该配置文件已预设好数据集路径和网络参数。

python gaussian_point_train.py --train_config config/tat_truck_every_8_test.yaml

训练过程中，程序会读取多视角图像和稀疏点云，通过可微分光栅器优化点云属性（协方差、颜色等），最终输出带有额外特征的稠密点云。

3. 推理与渲染

训练完成后，模型会自动保存（具体路径在 yaml 配置中指定，通常在 output-model-dir）。推理阶段只需加载训练好的稠密点云和新的相机姿态即可渲染图像。具体的推理脚本逻辑包含在训练代码的验证环节中，或者你可以参考仓库中的测试脚本来加载模型并生成新视角的图像。

提示：对于其他数据集（如 BlenderNerf 生成的数据或 Instant-NGP 格式数据），请参考仓库 tools 目录下的转换脚本进行数据预处理，并修改对应的 yaml 配置文件路径后运行相同的训练命令。