绘制你自己的 GAN

项目 | 论文 | YouTube | 幻灯片

我们的方法接受一张或多张手绘草图，并将现成的 GAN 进行定制，以匹配输入的草图。尽管我们的新模型会改变物体的形状和姿态，但颜色、纹理、背景等其他视觉线索在修改后仍会被忠实地保留。

Sheng-Yu Wang¹, David Bau², Jun-Yan Zhu¹。
CMU¹, MIT CSAIL²
在 ICCV 2021 年会议上发表。

8月16日更新 训练代码、评估代码和数据集已发布。模型权重也已更新，如果您在此更新之前下载过权重，请重新运行 bash weights/download_weights.sh。

结果

我们的方法可以将预训练的 GAN 定制为与输入草图相匹配。

使用我们定制的模型进行插值。 使用我们定制的模型进行潜在空间插值时，过程非常平滑。

图像1

插值

图像2

使用我们定制的模型进行图像编辑。 给定一张真实图像 (a)，我们使用 Huh et al. 将其投影到原始模型的潜在空间 z 中 (b)。然后，我们将投影得到的 z 输入到我们基于草图训练的站立猫模型中 (c)。最后，我们使用 GANSpace 的 添加毛发 操作对图像进行编辑 (d)。

模型插值。 我们可以通过插值 W 潜在空间来实现定制模型之间的过渡。

模型1

W 潜在空间中的插值

模型2

我们观察到，直接插值模型权重也能产生类似的效果。

模型1

模型权重空间中的插值

模型2

失败案例。 我们的方法无法生成与 Attneave 的猫草图或毕加索的马草图相匹配的图像。值得注意的是，Attneave 的猫描绘了一个复杂的姿势，而毕加索的草图则具有独特的风格，这两点都使得我们的方法难以处理。

快速入门

克隆我们的仓库

git clone git@github.com:PeterWang512/GANSketching.git
cd GANSketching

安装依赖包

• 安装 PyTorch（版本 >= 1.6.0）(pytorch.org)

  pip install -r requirements.txt
  

下载模型权重

• 运行 bash weights/download_weights.sh

从定制模型生成样本

此命令会运行由 ckpt 指定的定制模型，并将生成的样本保存到 save_dir 目录中。

# 从“站立猫”模型生成样本。
python generate.py --ckpt weights/photosketch_standing_cat_noaug.pth --save_dir output/samples_standing_cat

# 从论文图1中的猫脸模型生成样本。
python generate.py --ckpt weights/by_author_cat_aug.pth --save_dir output/samples_teaser_cat

# 从定制的 ffhq 模型生成样本。
python generate.py --ckpt weights/by_author_face0_aug.pth --save_dir output/samples_ffhq_face0 --size 1024 --batch_size 20

使用 GANSpace 进行潜在空间编辑

我们的模型保留了原始模型的潜在空间可编辑性。我们可以使用 Härkönen 等人报告的潜在变量（GANSpace），对我们的模型进行同样的编辑操作。

# 为站立的猫添加毛发
python ganspace.py --obj cat --comp_id 27 --scalar 50 --layers 2,4 --ckpt weights/photosketch_standing_cat_noaug.pth --save_dir output/ganspace_fur_standing_cat

# 为站立的猫闭上眼睛
python ganspace.py --obj cat --comp_id 45 --scalar 60 --layers 5,7 --ckpt weights/photosketch_standing_cat_noaug.pth --save_dir output/ganspace_eye_standing_cat

Colab

感谢 Luke Floden 创建了一个 Colab 笔记本，用于快速完成上述入门步骤。Colab 链接：quickstart_colab.ipynb

模型训练

在 PhotoSketch 数据上训练和评估模型时，需要运行 精确率与召回率指标。以下命令将拉取分叉的精确率与召回率 repo 的子模块。

git submodule update --init --recursive

下载数据集和预训练模型

以下脚本将下载我们的草图数据、评估集、LSUN 以及来自 StyleGAN2 和 PhotoSketch 的预训练模型。

# 下载草图数据
bash data/download_sketch_data.sh

# 下载评估集
bash data/download_eval_data.sh

# 从 StyleGAN2 和 PhotoSketch 下载预训练模型
bash pretrained/download_pretrained_models.sh

# 下载 LSUN 猫、马和教堂数据集
bash data/download_lsun.sh

要使用图像正则化训练 FFHQ 模型，请通过此 链接 下载 FFHQ 数据集。这是一个包含 70,000 张 1024x1024 分辨率图像的压缩包。请解压文件，将 images1024x1024 文件夹重命名为 ffhq，并将其放置在 ./data/image/ 目录下。

训练脚本

示例训练配置通过 scripts 文件夹中的脚本指定。使用以下命令启动训练。

# 训练“骑马者”模型
bash scripts/train_photosketch_horse_riders.sh

# 训练论文图 1 中的猫脸模型
bash scripts/train_teaser_cat.sh

# 使用单个 QuickDraw 草图进行训练
bash scripts/train_quickdraw_single_horse0.sh

# 使用人脸草图（1024px）进行训练
bash scripts/train_authorsketch_ffhq0.sh

默认情况下，训练进度会使用 wandb 进行跟踪。若要禁用 wandb 日志记录，请在训练脚本中添加 --no_wandb 标记。

评估

在运行评估之前，请确保已下载评估数据集和模型权重。

# 您可能已在前面的部分中运行过这些脚本
bash weights/download_weights.sh
bash data/download_eval_data.sh

使用以下脚本评估模型，结果将保存到由 --output 标志指定的 CSV 文件中。--models_list 应包含一个元组列表，每个元组对应一个模型权重路径和评估数据。示例请参见 weights/eval_list。

python run_metrics.py --models_list weights/eval_list --output metric_results.csv

相关工作

• R. Gal, O. Patashnik, H. Maron, A. Bermano, G. Chechik, D. Cohen-Or. "StyleGAN-NADA: CLIP 引导的图像生成器领域适应。". ArXiv 出版。（同时期工作）
• D. Bau, S. Liu, T. Wang, J.-Y. Zhu, A. Torralba. "重写深度生成模型". ECCV 2020。
• Y. Wang, A. Gonzalez-Garcia, D. Berga, L. Herranz, F. S. Khan, J. van de Weijer. "MineGAN：从 GAN 向目标域的有效知识迁移，即使只有少量图像". CVPR 2020。
• M. Eitz, J. Hays, M. Alexa. "人类如何绘制物体草图？". SIGGRAPH 2012。

致谢

本仓库部分借鉴了 SPADE、stylegan2-pytorch、PhotoSketch、GANSpace 以及 data-efficient-gans 的代码。

参考文献

如果您认为本项目对您的研究有所帮助，请引用以下文献。

@inproceedings{wang2021sketch,
  title={Sketch Your Own GAN},
  author={Wang, Sheng-Yu and Bau, David and Zhu, Jun-Yan},
  booktitle={IEEE 国际计算机视觉会议论文集},
  year={2021}
}

欢迎随时与我们联系，提出任何意见或反馈。

GANSketching 快速上手指南

GANSketching 是一个允许用户通过手绘草图定制预训练 GAN 模型的开源工具。它能在改变物体形状和姿态的同时，忠实保留颜色、纹理和背景等视觉特征。

环境准备

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

• 操作系统: Linux (推荐) 或 macOS
• Python: 3.6+
• PyTorch: 版本 >= 1.6.0
• GPU: 推荐使用 NVIDIA GPU 以加速推理和训练
• Git: 用于克隆代码库

安装步骤

1. 克隆仓库

首先，将项目代码克隆到本地：

git clone git@github.com:PeterWang512/GANSketching.git
cd GANSketching

2. 安装依赖

安装所需的 Python 包（包含 PyTorch 及其他依赖）：

pip install -r requirements.txt

提示：国内用户若下载缓慢，可指定清华源加速安装： pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

3. 下载模型权重

运行脚本下载预训练的模型权重：

bash weights/download_weights.sh

基本使用

安装完成后，您可以直接使用预训练模型生成图像。以下示例展示如何从定制的“站立猫”模型中生成样本。

生成图像样本

运行 generate.py 脚本，指定模型检查点路径和输出目录：

# 从 "standing cat" (站立的猫) 模型生成样本
python generate.py --ckpt weights/photosketch_standing_cat_noaug.pth --save_dir output/samples_standing_cat

# 从论文图 1 中的 "cat face" (猫脸) 模型生成样本
python generate.py --ckpt weights/by_author_cat_aug.pth --save_dir output/samples_teaser_cat

# 从定制的 FFHQ (人脸) 模型生成高分辨率样本
python generate.py --ckpt weights/by_author_face0_aug.pth --save_dir output/samples_ffhq_face0 --size 1024 --batch_size 20

生成的图片将保存在指定的 --save_dir 文件夹中。

进阶：潜在空间编辑 (可选)

该模型保留了原始 GAN 的潜在空间可编辑性。您可以结合 GANSpace 对生成的图像进行属性编辑（例如给猫添加毛发）：

# 给站立的猫添加毛发
python ganspace.py --obj cat --comp_id 27 --scalar 50 --layers 2,4 --ckpt weights/photosketch_standing_cat_noaug.pth --save_dir output/ganspace_fur_standing_cat

提示：如果您希望快速体验而无需配置本地环境，可以使用官方提供的 Colab 笔记本：quickstart_colab.ipynb。