3D感知条件图像合成 (pix2pix3D)

这是"3D感知条件图像合成"的官方PyTorch实现。Pix2pix3D可以根据2D标签图（如分割图或边缘图）合成3D物体（神经场）。我们还提供了一个交互式3D编辑演示。

https://user-images.githubusercontent.com/28395429/222578030-8bb2c727-397e-44b6-9ab1-9b0b09dd3b5b.mp4

CVPR 2023

卡内基梅隆大学

我们提出了pix2pix3D，这是一个用于可控真实图像合成的3D感知条件生成模型。给定2D标签图（如分割图或边缘图），我们的模型学会从不同视角合成相应的图像。为了实现显式的3D用户控制，我们将条件生成模型扩展为神经辐射场。利用广泛可用的单目图像和标签图对，我们的模型学会为每个3D点分配标签，同时还有颜色和密度，这使得它能够同时渲染图像和像素对齐的标签图。最后，我们构建了一个交互式系统，允许用户从任何视角编辑标签图并相应地生成输出。

入门指南

依赖项

我们提供了一个包含所有其他依赖项的conda环境文件。你可以使用以下命令和Miniconda3创建并激活你的Python环境：

conda env create -f environment.yml
conda activate pix2pix3d

数据

我们提供了预处理过的数据集，包括分割图和边缘图。你可以下载CelebAMask数据集、AFHQ-Cat-Seg数据集和Shapenet-Car-Edge数据集，并将这些zip文件放在data/目录下。

预训练模型

你可以使用以下脚本下载我们的预训练模型：

bash checkpoints/download_models.sh

推理

一旦你下载了预训练模型，我们提供了几个脚本来生成结果。

生成样本

你可以基于数据集中的样本生成结果。

python applications/generate_samples.py --network <network_pkl> --outdir <output_dir> --random_seed <random_seeds list, e.g. 0 1> --cfg <configs, e.g., seg2cat, seg2face, edge2car> --input_id <sample_id in dataset>

例如：

输入标签图	生成的图像	生成的标签图
<img src="https://yellow-cdn.veclightyear.com/835a84d5/52e3ba8a-c5af-48e4-96e2-c71109814552.png" width="256" />	<img src="https://yellow-cdn.veclightyear.com/835a84d5/26396f00-633f-43ca-a018-b4f400570945.png" width="256" />	<img src="https://yellow-cdn.veclightyear.com/835a84d5/cdac3fc1-9c20-4ce4-8234-03db6eb1f2b1.png" width="256" />

你可以通过运行以下命令获得上述结果：

python applications/generate_samples.py --network checkpoints/pix2pix3d_seg2cat.pkl --outdir examples --random_seed 1 --cfg seg2cat --input_id 1666

渲染视频

你可以基于指定的输入标签图渲染视频结果。

python applications/generate_video.py --network <network_pkl> --outdir <output_dir> --random_seed <random_seeds list, e.g. 0 1> --cfg <configs, e.g., seg2cat, seg2face, edge2car> --input <input label map>

例如:

输入标签图	生成的图像	生成的标签图
<img src="https://yellow-cdn.veclightyear.com/835a84d5/52e3ba8a-c5af-48e4-96e2-c71109814552.png" width="256" />	<img src="https://yellow-cdn.veclightyear.com/835a84d5/5967f112-0485-465b-8748-49d1eece5b91.gif" width="256" />	<img src="https://yellow-cdn.veclightyear.com/835a84d5/3af4201a-a9c8-4317-b74d-a343b261bdeb.gif" width="256" />

你可以使用以下命令获得上述结果:

python applications/generate_video.py --network checkpoints/pix2pix3d_seg2cat.pkl --outdir examples --random_seed 1 --cfg seg2cat --input examples/example_input.png

提取语义网格

你还可以提取网格并使用3D语义标签为其着色。网格提取需要一些额外的包(pyrender、trimesh和mcubes)。你可以通过pip安装它们。提取的网格将保存为semantic_mesh.ply。

例如:

输入标签图	语义网格
<img src="https://yellow-cdn.veclightyear.com/835a84d5/52e3ba8a-c5af-48e4-96e2-c71109814552.png" width="256" />	<img src="https://yellow-cdn.veclightyear.com/835a84d5/15a345c1-a0c1-4100-b5ef-b3666ef323e9.gif" width="256" />

你可以使用以下命令获得上述结果:

python applications/extract_mesh.py --network checkpoints/pix2pix3d_seg2cat.pkl --outdir examples --cfg seg2cat --input examples/example_input.png

训练

我们在train_scripts/afhq_seg.sh提供了一个示例训练脚本:

python train.py --outdir=<log_dir> \
                --cfg=afhq --data=data/afhq_v2_train_cat_512.zip \
                --mask_data=data/afhqcat_seg_6c.zip \
                --data_type=seg --semantic_channels=6 \ 
                --render_mask=True --dis_mask=True \ 
                --neural_rendering_resolution_initial=128 \
                --resume=<EG3D-checkpoints>/afhqcats512-128.pkl \
                --gpus=2 --batch=4 --mbstd-group=2 \ 
                --gamma=5 --gen_pose_cond=True \ 
                --random_c_prob=0.5 \
                --lambda_d_semantic=0.1 \
                --lambda_lpips=1 \
                --lambda_cross_view=1e-4 \
                --only_raw_recons=True \
                --wandb_log=False

训练参数:

outdir: 保存检查点和日志的目录。
cfg: 从[afhq, celeba, shapenet]中选择。
data: RGB数据文件。
mask_data: 标签图数据文件。
data_type: 从[seg, edge]中选择。如果使用seg，请指定semantic_channels。
render_mask: 是否与RGB一起渲染标签图。
dis_mask: 是否对渲染的标签图使用GAN损失。
neural_rendering_resolution_initial: NeRF输出的分辨率。
resume: 我们使用EG3D预训练检查点部分初始化我们的网络(在这里下载)。
gpus, batch, mbstd-group: 批量大小和多GPU训练的参数。
gen_pose_cond: 是否根据相机姿态进行条件生成。
random_c_prob: 训练时采样随机姿态的概率。
lambda_d_semantic: 标签图GAN损失的权重。
lambda_lpips: RGB LPIPS损失的权重。
lambda_cross_view: 跨视图一致性损失的权重。
wandb_log: 是否使用wandb日志。

准备你自己的数据集

我们遵循EG3D的数据集格式在这里。你可以通过DINO聚类获取自己数据集的分割掩码，并通过pidinet和informative drawing获取边缘图。

引用

如果你发现这个仓库对你的研究有用，请引用以下工作。

@inproceedings{kangle2023pix2pix3d,
  title={3D-aware Conditional Image Synthesis},
  author={Deng, Kangle and Yang, Gengshan and Ramanan, Deva and Zhu, Jun-Yan},
  booktitle = {CVPR},
  year = {2023}
}

致谢

我们感谢Sheng-Yu Wang、Nupur Kumari、Gaurav Parmer、Ruihan Gao、Muyang Li、George Cazenavette、Andrew Song、Zhipeng Bao、Tamaki Kojima、Krishna Wadhwani、Takuya Narihira和Tatsuo Fujiwara的讨论和帮助。我们感谢索尼公司、新加坡DSTA和CMU Argo AI自动驾驶研究中心的支持。这个代码库大量借鉴了EG3D和StyleNeRF。