Infusion - 基于扩散先验的3D高斯体场景修复技术

InFusion: 通过从扩散先验学习深度补全来修复3D高斯

刘志恒* · 欧阳浩* · 王秋雨 · 郑嘉良 · 肖杰 · 朱凯 · 薛楠 · 刘宇 · 沈昱君 · 曹杨


中国科学技术大学 | 香港科技大学 | 蚂蚁集团 | 阿里巴巴集团

新闻

[2024.4.18] 🔥 发布论文、推理代码和预训练检查点。
[进行中] 整理并组织实验中使用的数据集对应的掩码。
[进行中] 使用更多训练数据扩展模型，并发布更强大的模型作为下游任务的基础模型。
[待办] 发布训练代码。

安装

使用conda安装：

conda env create -f environment.yaml
conda activate infusion

🛠️ 对于深度渲染，我们使用来自FSGS的diff-gaussian-rasterization-confidence，感谢他们的工作！:)

下载检查点

下载Infusion检查点并将其放在'checkpoints'文件夹中：

HuggingFace

数据准备

我们的实验在Mip-NeRF、Instruct-NeRF2NeRF和SPIn-NeRF提供的数据集上进行。我们将在几天内上传实验中使用的掩码和我们自己拍摄的挑战场景。

以Mip-NeRF中的"garden"为例，每个场景文件夹应按如下方式组织。

garden
├── images # RGB数据
│   ├── DSC07956.JPG
│   ├── DSC07957.JPG
│   └── ...                   
├── seg # 掩码 
│   ├── DSC07956.JPG
│   ├── DSC07957.JPG
│   └── ... 
│   # 需要修复的部分为白色
└── sparse # Colmap
│   └── 0
│       └── ...

🛠️ 您可以按照这样的结构准备自己的数据。此外，准确的掩码非常重要。在这里，我们推荐两个图像分割工具：Segment and Track Anything和Grounded SAM。
🛠️ 要获取相机参数和初始点云，请参考Gaussian-Splatting中的'convert.py' :)

使用说明

整个流程分为三个阶段：

训练不完整的高斯。
通过扩散先验修复高斯。
结合修复的高斯和不完整的高斯。

🌺 第1阶段

使用预先标注的掩码来训练不完整的高斯。

cd gaussian_splatting
# 训练不完整的高斯
python train.py -s <场景文件夹路径> -m <输出文件夹路径> -u nothing --mask_training
#--color_aug

# 获取c2w矩阵、内参矩阵、不完整深度和rgb渲染图像
python render.py -s <场景文件夹路径> -m <输出文件夹路径> -u nothing

🛠️ 提示：有时，渲染的深度可能会有太多空白点。您可以在训练时使用--color_aug，这将在渲染深度时随机选择背景颜色，可能会使深度图更可靠。
🛠️ 最近，一些工作专注于如何分割高斯。这不是本工作的重点，所以选择了一种相对简单的方法。:)

🌺 第2阶段

使用深度修复模型修复高斯。

📢 您需要在*'输出文件夹路径/train/ours_30000/renders'中选择一个单一图像，并标记需要修复的区域，并将其保存为*'mask.png'**。（它不必精确，但需要覆盖所有缺失的部分。）
📢 接下来,你需要对单张图像进行修复。这里有一些很好的工具可以用来修复单张图像:Stable Diffusion XL Inpainting和Photoroom。以下是一个例子:

# 假设选定的单张图像名为"DSC07956.JPG"。
cd depth_inpainting/run
input_rgb_path=<修复后的单张图像路径>
input_mask_path=<'mask.png'的路径>
input_depth_path=<输出文件夹/train/ours_30000/depth_dis/DSC07956.npy的路径>
c2w=<输出文件夹/train/ours_30000/c2w/DSC07956.npy的路径>
intri=<输出文件夹/train/ours_30000/intri/DSC07956.npy的路径>
model_path=<深度修复模型检查点的路径>  # 绝对路径
output_dir=<输出文件夹的路径>


CUDA_VISIBLE_DEVICES=0 python run_inference_inpainting.py \
            --input_rgb_path $input_rgb_path \
            --input_mask $input_mask_path \
            --input_depth_path $input_depth_path \
            --model_path $model_path \
            --output_dir $output_dir \
            --denoise_steps 20 \
            --intri $intri \
            --c2w $c2w \
            --use_mask\
            --blend  # 是否在推理过程中使用'混合扩散(https://arxiv.org/abs/2111.14818)'。

🛠️ 提示:如果你觉得一次推理获得的深度图不够理想,你可以使用新获得的output_dir/<inpainted_image_name>_depth_dis.npy作为新的$input_depth_path,循环两三次以获得更好的结果。

🌺 第3阶段

合并修复后的高斯点和不完整的高斯点,并在修复后的单张图像上快速微调。

# 假设选定的单张图像名为"DSC07956.JPG"。
origin_ply="输出文件夹/point_cloud/iteration_30000/point_cloud.ply的路径"
supp_ply="输出文件夹/DSC07956_mask.ply的路径"
save_ply="输出文件夹/point_cloud/iteration_30001/point_cloud.ply的路径"
# 合并修复后的高斯点和不完整的高斯点。
python compose.py --original_ply $origin_ply  --supp_ply $supp_ply --save_ply $save_ply --nb_points 100 --threshold 1.0
# 在修复后的单张图像上微调150次迭代。
python train.py -s <场景文件夹路径> -m <输出文件夹路径> -u DSC07956.JPG -n <修复后的单张图像路径> --load_iteration 30001 --iteration 150
# 渲染
python render.py -s <场景文件夹路径> -m <输出文件夹路径> -u nothing --iteration 150

🛠️ 参数--nb_points和--threshold用于去除漂浮点。增加它们的值会移除更多周围的点。去除漂浮点对最终渲染结果非常重要。在这里,我们需要找到最适合从场景中去除漂浮点的参数。
🛠️ 作为显式点,高斯点可以在实际应用中直接编辑和裁剪,例如在KIRI Engine中。

致谢

本项目基于Gaussian-Splatting、Marigold和Magicboomliu的代码库开发。我们感谢他们的出色工作!

引用

如果你发现这个仓库对你的工作有帮助,请考虑引用以下论文,并给公共仓库一个⭐,让更多人发现这个仓库:

@article{liu2024infusion,
  title={InFusion: Inpainting 3D Gaussians via Learning Depth Completion from Diffusion Prior},
  author={Liu, Zhiheng and Ouyang, Hao and Wang, Qiuyu and Cheng, Ka Leong and Xiao, Jie and Zhu, Kai and Xue, Nan and Liu, Yu and Shen, Yujun and Cao, Yang},
  journal={arXiv preprint arXiv:2404.11613},
  year={2024}
}