用交换自编码器进行深度图像 manipulation

Taesung Park、Jun-Yan Zhu、Oliver Wang、Jingwan Lu、Eli Shechtman、Alexei A. Efros、Richard Zhang

加州大学伯克利分校与 Adobe 研究院

NeurIPS 2020

项目页面 | 论文 | 3 分钟视频

概述

交换自编码器由编码（顶部）和交换（底部）两部分组成。 顶部：编码器 E 将输入图像（巴黎圣母院）编码为两个代码。结构代码是一个具有空间维度的张量；纹理代码是一个 2048 维向量。使用生成器 G 解码后应生成一张逼真的图像（由判别器 D 确保与输入匹配，即重建损失）。 底部：使用来自第二张图像（圣巴西尔大教堂）的纹理代码解码时，应在判别器 D 的监督下看起来逼真，并且纹理与目标图像一致。为此，通过训练一个补丁共现判别器 Dpatch 来确保输出图像和参考图像的局部区域难以区分。

安装 / 要求

• 需要 CUDA 10.1 或更高版本，因为该模型使用了 StyleGAN2 的自定义 CUDA 内核，由 @rosinality 移植而来。
• 作者在 Python 3.6 上使用了 PyTorch 1.7.1。
• 使用 pip install dominate torchgeometry func-timeout tqdm matplotlib opencv_python lmdb numpy GPUtil Pillow scikit-learn visdom ninja 安装依赖项。

测试与评估。

我们提供了预训练好的模型以及用于复现论文中各图的若干示例图片。请从 这里 (2.1GB) 下载并解压（注意：这是一个 http（而非 https）地址，某些浏览器如 Chrome 可能需要将链接直接粘贴到地址栏中，或者使用 wget 下载数据集）。脚本假设检查点位于 ./checkpoints/，测试图片位于 ./testphotos/，但可以通过修改 --checkpoints_dir 和 --dataroot 参数来更改路径。

更新：新增了 AFHQ 数据集的预训练模型。请从 这里 (256MB) 下载模型和示例图片（同样，可能需要将链接直接粘贴到地址栏中）。

使用示例图片对山景模型进行交换与插值

要运行简单的交换和插值操作，需指定两幅参考图像，修改 experiments/mountain_pretrained_launcher.py 中的 input_structure_image 和 input_texture_image 字段，然后运行：

python -m experiments mountain_pretrained test simple_swapping
python -m experiments mountain_pretrained test simple_interpolation

提供的脚本，尤其是 experiments/mountain_pretrained_launcher.py 中的 opt.tag("simple_swapping") 和 opt.tag("simple_interpolation")，会调用类似于以下的终端命令：

python test.py --evaluation_metrics simple_swapping \
--preprocess scale_shortside --load_size 512 \
--name mountain_pretrained  \
--input_structure_image [样本图片路径] \
--input_texture_image [样本图片路径] \
--texture_mix_alpha 0.0 0.25 0.5 0.75 1.0

换句话说，如果觉得这样更直观，也可以直接使用此命令。

输出图像将保存在 ./results/mountain_pretrained/simpleswapping/ 目录下。

纹理交换

我们的交换自编码器能够学习将纹理与结构分离，从而实现诸如纹理交换之类的图像编辑任务。每一行展示了最左侧图像的结构代码与上方图像的纹理代码结合后的结果。

要复现这张图（图 4）以及论文中的图 9 和图 12，可运行以下命令：

# 读取 ./experiments/church_pretrained_launcher.py 中的选项
python -m experiments church_pretrained test swapping_grid

# 读取 ./experiments/bedroom_pretrained_launcher.py 中的选项
python -m experiments bedroom_pretrained test swapping_grid

# 读取 ./experiments/mountain_pretrained_launcher.py 中的选项
python -m experiments mountain_pretrained test swapping_grid

# 读取 ./experiments/ffhq512_pretrained_launcher.py 中的选项
python -m experiments ffhq512_pretrained test swapping_grid

请确保在相应的 ./experiments/xx_pretrained_launcher.py 脚本中正确设置了 dataroot 和 checkpoints_dir 路径。

定量评估

要进行表 1、图 5 和表 2 中提到的 FID 等定量评估，首先需要准备输入结构和纹理参考图像的配对。

这些参考图像随机选自 LSUN、FFHQ 和瀑布数据集的验证集。输入结构和纹理图像应分别放在 input_structure/ 和 input_style/ 目录中，且文件名相同。例如，input_structure/001.png 和 input_style/001.png 将被一起加载以进行交换。

将脚本 experiments/church_pretrained_launcher.py 中的 dataroot="./testphotos/church/fig5_tab2/" 字段替换为测试图片的路径，然后运行：

python -m experiments church_pretrained test swapping_for_eval
python -m experiments ffhq1024_pretrained test swapping_for_eval

结果可在 ./results 目录中查看（可通过 --result_dir 参数更改保存路径）。

随后，使用 https://github.com/mseitzer/pytorch-fid 计算交换后的图像与原始结构图像之间的 FID 值。

模型训练。

数据集

• LSUN 教堂和卧室 数据集可以从 这里 下载。下载并解压后，目录应包含 [category]_[train/val]_lmdb/。
• FFHQ 数据集 可以通过此 链接 下载。这是一个包含 70,000 张 1024x1024 分辨率图像的压缩文件。解压后，我们将直接加载这些图像文件。
• Flickr 山脉 数据集和 Flickr 瀑布 数据集由于许可问题无法公开分享。不过，这些图片是从 Mountains Anywhere 和 Waterfalls Around the World 收集而来，使用了 Flickr API 的 Python 封装库。如需更多详情，请以“用于交换自编码器的 Flickr 数据集”为标题联系 Taesung Park。
• AFHQ 数据集 可以从 这里 下载。

训练脚本

训练配置通过 experiments/*_launcher.py 中的脚本指定。可以使用以下命令启动各种训练：

# 修改 experiments/[church,bedroom,ffhq,mountain]_launcher.py 中的 |dataroot| 和 |checkpoints_dir|
python -m experiments church train church_default
python -m experiments bedroom train bedroom_default
python -m experiments ffhq train ffhq512_default
python -m experiments ffhq train ffhq1024_default

# 默认情况下，脚本会使用 GPUtil 检查机器上可用的 GPU，并设置合适的 GPU ID。若要指定特定的 GPU 集合，可使用 |--gpu| 选项。同时请务必在相应脚本中更改 |num_gpus| 选项。
python -m experiments church train church_default --gpu 01234567

训练进度可以通过 visdom 在 --display_port 指定的端口上监控。默认端口是 https://localhost:2004。作为参考，在 LSUN 教堂 256px 数据集上，使用 4 块 V100 GPU 进行训练大约需要 14 天。

此外，还会使用训练集的随机样本生成若干交换网格，并将其保存为网页文件，路径为 [checkpoints_dir]/[expr_name]/snapshots/。网格生成的频率由 --evaluation_freq 控制。

所有可配置参数都会在训练开始时打印出来。这些配置分散在相关类别的代码中，例如 models/swapping_autoencoder_model.py、util/iter_counter.py 或 models/networks/encoder.py 中的 def modify_commandline_options 函数内。若需更改这些配置，只需在训练脚本（如 experiments/church_launcher.py）中的 opt.specify 部分修改相应的选项即可。

解析和配置相关的代码位于 experiments/__init__.py、experiments/__main__.py 和 experiments/tmux_launcher.py 中。

继续训练

默认情况下，训练会从最后一个检查点继续进行，因为 --continue_train 选项默认设置为 True。若想从头开始训练，可以删除检查点文件，或者在训练脚本中将 continue_train 设置为 False（例如在 experiments/church_launcher.py 中）。

代码结构（主要功能）

• models/swapping_autoencoder_model.py：定义损失函数并生成可视化结果的核心文件。
• optimizers/swapping_autoencoder_optimizer.py：定义优化器以及 GAN 的交替训练过程。
• models/networks/：包含模型架构文件，如 generator.py、discriminator.py、encoder.py、patch_discrimiantor.py 和 stylegan2_layers.py。
• options/__init__.py：包含基础选项标志。然而，许多重要选项分散在其他文件中，例如 swapping_autoencoder_model.py 或 generator.py。程序启动时，这些选项会被统一解析。查看已使用选项列表的最佳方式是运行训练脚本，并观察控制台输出的配置信息。
• util/iter_counter.py：包含迭代计数功能。

更改记录

• 2021年4月14日：针对在山脉数据集上训练预训练模型的配置未正确设置，现已更新。
• 2021年10月14日：新增了适用于 AFHQ 数据集的 256x256 预训练模型。请使用 experiments/afhq_pretrained_launcher.py。

BibTeX 引用

如果您在研究中使用了本代码，请引用我们的论文：

@inproceedings{park2020swapping,
  title={Swapping Autoencoder for Deep Image Manipulation},
  author={Park, Taesung and Zhu, Jun-Yan and Wang, Oliver and Lu, Jingwan and Shechtman, Eli and Efros, Alexei A. and Zhang, Richard},
  booktitle={Advances in Neural Information Processing Systems},
  year={2020}
}

致谢

StyleGAN2 的部分层大量借鉴（甚至可以说是直接复制！）了 @rosinality 的 PyTorch 实现。我们感谢 Nicholas Kolkin 在自动化内容与风格评估方面的有益讨论，Jeongo Seo 和 Yoseob Kim 在用户界面设计上的建议，以及 William T. Peebles、Tongzhou Wang 和 Yu Sun 在解耦表示方面的交流与探讨。

Swapping Autoencoder 快速上手指南

Swapping Autoencoder 是一个用于深度图像操作的开源工具，能够将一张图像的结构（Structure）与另一张图像的纹理/风格（Texture）进行解耦和交换。本项目基于 PyTorch 实现，支持图像风格迁移、纹理替换及插值生成。

1. 环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐) 或 macOS
• GPU: NVIDIA GPU (必须)，支持 CUDA 10.1 或更高版本
  • 注：该项目使用了 StyleGAN2 的自定义 CUDA 内核，必须依赖 GPU 运行。
• Python: 推荐 Python 3.6 (作者测试环境为 PyTorch 1.7.1 + Python 3.6)
• 依赖库:
  • PyTorch
  • dominate, torchgeometry, func-timeout, tqdm, matplotlib, opencv_python, lmdb, numpy, GPUtil, Pillow, scikit-learn, visdom, ninja

国内加速建议： 安装 PyTorch 时，建议使用清华或中科大镜像源以加快下载速度。 例如：pip install torch torchvision -i https://pypi.tuna.tsinghua.edu.cn/simple

2. 安装步骤

第一步：安装基础依赖

使用 pip 安装项目所需的所有 Python 包：

pip install dominate torchgeometry func-timeout tqdm matplotlib opencv_python lmdb numpy GPUtil Pillow scikit-learn visdom ninja

第二步：下载预训练模型与测试图片

项目提供了预训练模型（涵盖 Church, Bedroom, Mountain, FFHQ, AFHQ 等数据集）。

下载地址：

• 通用模型与测试图 (2.1GB): http://efrosgans.eecs.berkeley.edu/SwappingAutoencoder/swapping_autoencoder_models_and_test_images.zip
• AFHQ 动物数据集模型 (256MB): http://efrosgans.eecs.berkeley.edu/SwappingAutoencoder/afhq_models_and_test_images.zip

注意：部分浏览器可能无法直接点击 http 链接下载，请将链接复制粘贴到地址栏，或使用 wget 命令下载。

第三步：解压与目录配置

下载完成后，解压文件。默认脚本假设目录结构如下（可根据需要修改）：

• 模型检查点位于：./checkpoints/
• 测试图片位于：./testphotos/

如果路径不同，请在运行命令时通过 --checkpoints_dir 和 --dataroot 参数指定。

3. 基本使用

以下示例展示如何使用预训练的 "Mountain" 模型进行最简单的纹理交换和插值操作。

场景一：简单纹理交换 (Simple Swapping)

将一张图片的结构与另一张图片的纹理结合。

1. 编辑配置文件 experiments/mountain_pretrained_launcher.py，修改以下字段指向你的图片路径：

   • input_structure_image: 提供结构的图片路径
   • input_texture_image: 提供纹理的图片路径

2. 运行交换命令：

python -m experiments mountain_pretrained test simple_swapping

或者直接通过命令行参数指定（无需修改脚本）：

python test.py --evaluation_metrics simple_swapping \
--preprocess scale_shortside --load_size 512 \
--name mountain_pretrained  \
--input_structure_image [path_to_structure_image] \
--input_texture_image [path_to_texture_image] \
--texture_mix_alpha 0.0

结果将保存在：./results/mountain_pretrained/simpleswapping/

场景二：纹理插值 (Simple Interpolation)

在两张不同纹理之间生成平滑过渡的效果。

1. 同样配置好 input_structure_image 和 input_texture_image（通常需要两张不同的纹理图来观察插值效果，具体视脚本逻辑而定，此处沿用上述配置逻辑）。

2. 运行插值命令：

python -m experiments mountain_pretrained test simple_interpolation

或直接使用命令行：

python test.py --evaluation_metrics simple_swapping \
--preprocess scale_shortside --load_size 512 \
--name mountain_pretrained  \
--input_structure_image [path_to_structure_image] \
--input_texture_image [path_to_texture_image] \
--texture_mix_alpha 0.0 0.25 0.5 0.75 1.0

注：--texture_mix_alpha 参数定义了混合比例，多个数值将生成一系列过渡图片。

场景三：批量网格交换 (Swapping Grid)

如果你想复现论文中的交换网格图（多张结构图 x 多张纹理图），可以使用以下命令（以 Church 数据集为例）：

# 确保 experiments/church_pretrained_launcher.py 中的路径配置正确
python -m experiments church_pretrained test swapping_grid

支持的预设模型包括：church, bedroom, mountain, ffhq512, afhq 等。

---

更多高级功能（如定量评估 FID、从头训练模型等）请参考原始 README 文档。