[Arxiv 2024] StyleGaussian: 基于高斯散射的即时3D风格迁移

项目主页 | 论文

本仓库包含了论文StyleGaussian: 基于高斯散射的即时3D风格迁移的PyTorch实现。StyleGaussian是一种新颖的3D风格迁移管道，能够实现即时风格迁移，同时保持实时渲染和严格的多视图一致性。

概览图

1 安装

我们使用Mamba来管理环境，因为它速度快，但也可以使用Conda。

mamba env create -f environment.yml -n stylegaussian

2 快速开始

数据集和检查点（Google Drive）

请下载预处理的数据集并将它们放在datasets文件夹中。我们还提供了预训练的检查点，应放在output文件夹中。如果您更改了数据集或检查点的位置，请相应地修改检查点文件夹中cfg_args里的model_path或source_path。

2.1 交互式远程查看器

我们的交互式远程查看器基于Viser。要启动查看器，运行以下命令：

python viewer.py -m [model_path] --style_folder [style_folder] --viewer_port [viewer_port]

# 例如：
python viewer.py -m output/train/artistic/default --style_folder images --viewer_port 8080

其中model_path是预训练模型的路径，命名为output/[scene_name]/artistic/[exp_name]，style_folder是包含风格图像的文件夹，viewer_port是查看器的端口号。--style_folder和--viewer_port是可选参数。

使用方法：

渲染设置。 使用鼠标旋转相机，使用鼠标滚轮放大缩小。使用WASD移动相机，QE上下移动。您可以从训练视图中选择视角。点击Reset View重置上方向。您还可以更改渲染分辨率和图像质量，这会影响渲染速度，因为重构速度受限于网络速度。
风格迁移。 在Style Image文本框中指定风格图像路径后，渲染将立即开始，风格也会立即迁移。Random Style按钮将从风格文件夹中随机选择一张风格图像。点击Display Style Image按钮在右下角显示风格图像。
风格插值。 您可以通过在Style 1和Style 2中指定两张风格图像的路径来在两种风格之间进行插值。可以通过滑块调整插值比例。Style 1图像显示在左下角，Style 2图像显示在右下角。

2.2 推理渲染

1. 风格迁移

您可以通过以下方式对场景进行单一风格图像的风格迁移：

python render.py -m [model_path] --style [style_image_path] 

# 例如：
python render.py -m output/train/artistic/default --style images/0.jpg

其中model_path是预训练模型的路径，命名为output/[scene_name]/artistic/[exp_name]，style_image_path是风格图像的路径。渲染的风格化多视图图像将保存在output/[scene_name]/artistic/[exp_name]/train文件夹中。

2. 风格插值

您可以通过以下方式对场景进行四张风格图像的风格插值：

python render.py -m [model_path] --style [style_image_path1] [style_image_path2] [style_image_path3] [style_image_path4]

# 例如：
python render.py -m output/train/artistic/default --style images/0.jpg images/1.jpg images/2.jpg images/3.jpg

其中model_path是预训练模型的路径，命名为output/[scene_name]/artistic/[exp_name]，style_image_path1、style_image_path2、style_image_path3、style_image_path4是四张风格图像的路径。渲染的插值风格化可视化结果将保存在output/[scene_name]/artistic/[exp_name]/style_interpolation文件夹中。

3 训练

我们使用WikiArt数据集作为风格图像数据集。

3.1 从头开始训练

您可以通过以下方式从头开始训练模型：

python train.py --data [数据集路径] --wikiartdir [wikiart路径] --exp_name [实验名称]

# 例如：
python train.py --data datasets/train --wikiartdir datasets/wikiart --exp_name default

其中数据集路径是训练数据集的路径，wikiart路径是WikiArt数据集的路径，实验名称是实验的名称。训练过程将保存在output/[场景名称]/artistic/[实验名称]文件夹中。--exp_name是一个可选参数。

3.2 逐步训练

或者，您可以逐步训练模型。StyleGaussian分三个步骤进行训练：重建训练、特征嵌入训练和风格迁移训练。这三个步骤可以分别通过以下方式进行训练：

1. 重建训练

python train_reconstruction.py -s [数据集路径]

# 例如：
python train_reconstruction.py -s datasets/train

训练好的重建模型将保存在output/[场景名称]/reconstruction文件夹中。

2. 特征嵌入训练

python train_feature.py -s [数据集路径] --ply_path [ply文件路径]

# 例如：
python train_feature.py -s datasets/train --ply_path output/train/reconstruction/default/point_cloud/iteration_30000/point_cloud.ply

其中数据集路径是训练数据集的路径，ply文件路径是重建训练阶段重建的3D高斯点云的路径，命名为output/[场景名称]/reconstruction/[实验名称]/point_cloud/iteration_30000/point_cloud.ply。训练好的特征嵌入模型将保存在output/[场景名称]/feature文件夹中。

3. 风格迁移训练

python train_artistic.py -s [数据集路径] --wikiartdir [wikiart路径] --ckpt_path [特征检查点路径] --style_weight [风格权重]

# 例如：
python train_artistic.py -s datasets/train --wikiartdir datasets/wikiart --ckpt_path output/train/feature/default/chkpnt/feature.pth --style_weight 10

其中数据集路径是训练数据集的路径，wikiart路径是WikiArt数据集的路径，特征检查点路径是特征嵌入模型的检查点路径，命名为output/[场景名称]/feature/[实验名称]/chkpnt/feature.pth，风格权重是风格损失的权重。训练好的风格迁移模型将保存在output/[场景名称]/artistic/[实验名称]文件夹中。--style_weight是一个可选参数。

3.3 使用初始化进行训练

您可以通过使用另一个场景的训练模型初始化模型来节省风格迁移训练阶段的时间：

python train_artistic.py -s [数据集路径] --wikiartdir [wikiart路径] --ckpt_path [特征检查点路径] --style_weight [风格权重] --decoder_path [初始化检查点路径]

# 例如：
python train_artistic.py -s datasets/train --wikiartdir datasets/wikiart --ckpt_path output/train/feature/default/chkpnt/feature.pth --style_weight 10 --decoder_path output/truck/artistic/default/chkpnt/gaussians.pth

其中数据集路径是训练数据集的路径，wikiart路径是WikiArt数据集的路径，特征检查点路径是特征嵌入模型的检查点路径，命名为output/[场景名称]/feature/[实验名称]/chkpnt/feature.pth，风格权重是风格损失的权重，初始化检查点路径是另一个场景的训练模型的检查点路径，命名为output/[另一个场景]/artistic/[实验名称]/chkpnt/gaussians.pth。训练好的风格迁移模型将保存在output/[场景名称]/artistic/[实验名称]文件夹中。

4 致谢

我们的工作基于3D Gaussian Splatting和StyleRF。我们感谢这些作者的出色工作和开源代码。

5 引用

@article{liu2023stylegaussian,
  title={StyleGaussian: Instant 3D Style Transfer with Gaussian Splatting},
  author={Liu, Kunhao and Zhan, Fangneng and Xu, Muyu and Theobalt, Christian and Shao, Ling and Lu, Shijian},
  journal={arXiv preprint arXiv:2403.07807},
  year={2024},
}