Logo
Project Icon

HyperInverter

超网络技术实现高质量StyleGAN图像反演

StyleGANGAN反演图像编辑超网络人脸生成Github开源项目

HyperInverter是一种两阶段StyleGAN反演方法,结合编码器和超网络实现高质量图像重建、良好可编辑性和快速推理。实验表明,该方法在保持编码器级推理速度的同时,重建质量显著优于现有编码器方法,接近优化方法。在人脸和建筑图像上均取得出色效果。

访问官网访问官网GithubGithub论文论文
介绍相关项目
目录
  1. 开始使用
  2. 实验
  3. 致谢
  4. 联系方式

HyperInverter: 通过超网络改进StyleGAN反演

Tan M. Dinh, Anh Tran, Rang Nguyen, Binh-Son Hua
VinAI Research, 越南

摘要: 近年来,得益于对GAN潜在空间的探索和利用,真实世界图像操作取得了惊人的进展。GAN反演是这一流程的第一步,旨在忠实地将真实图像映射到潜在代码。不幸的是,大多数现有的GAN反演方法无法同时满足以下三个要求中的至少一个:高重建质量、可编辑性和快速推理。在这项研究中,我们提出了一种新颖的两阶段策略,可以同时满足所有要求。在第一阶段,我们训练一个编码器将输入图像映射到StyleGAN2的W空间,该空间被证明具有出色的可编辑性但较低的重建质量。在第二阶段,我们利用一系列超网络来恢复反演过程中丢失的信息,从而补充初始阶段的重建能力。这两个步骤相互补充,由于超网络分支而产生高重建质量,并由于W空间中完成的反演而具有出色的可编辑性。我们的方法完全基于编码器,从而实现极快的推理。在两个具有挑战性的数据集上进行的大量实验证明了我们方法的优越性。

teaser.png
我们的方法在质量上显著优于其他基于编码器的方法(pSp、e4e、ReStyle),同时推理时间与它们相同。与基于优化的方法(SG2-W+、PTI)相比,我们的工作在质量上与SG2-W+相当,略低于PTI。然而,值得注意的是,我们的方法运行速度非常快,可以用于交互式应用(比SG2-W+和PTI分别快3000倍和1100倍)。

模型架构和实验结果的详细信息可以在我们的论文中找到。

@inproceedings{dinh2021hyperinverter,
    title={HyperInverter: Improving StyleGAN Inversion via Hypernetwork},
    author={Tan M. Dinh and Anh Tuan Tran and Rang Nguyen and Binh-Son Hua},
    booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    year={2022}
}

请在使用我们的模型实现来产生发表结果或将其纳入其他软件时引用我们的论文。

开始使用

代码库在以下环境中测试:

  • Ubuntu
  • CUDA 10.0, CuDNN 7

安装

  • 克隆此仓库:
git clone https://github.com/VinAIResearch/HyperInverter.git
cd HyperInverter
  • 安装依赖:
conda create -p ./envs python=3.7.3
conda activate ./envs
pip install -r requirements.txt

数据集

  • 人脸: 我们使用FFHQ数据集中的70,000张图像进行训练,使用CelebA-HQ数据集中的2,824张图像进行测试。图像分辨率为1024 x 1024,并裁剪对齐到中心。有关预处理步骤的更多详细信息,请参考FFHQ。

  • 教堂: 我们使用LSUN Church官方训练集和测试集中的所有126,227张图像和300张图像来训练和评估我们的模型。图像调整为256 x 256分辨率。

请下载相应的数据集并解压到data文件夹。然后,转到configs/paths_config.py并修改此文件以正确链接数据:

dataset_paths = {
    "ffhq": "/path/to/ffhq/train_img",
    "celeba_test": "/path/to/CelebA-HQ/test_img",
    "church_train": "/path/to/lsun-church/train_img",
    "church_test": "/path/to/lsun-church/test_img",
}

如果您需要尝试自己的数据集,可以在以下文件中进行必要的修改:(i) data_configs.py定义数据路径;(ii) transforms_configs.py定义数据转换。

辅助预训练模型

运行以下命令自动下载实验所需的辅助预训练模型。

python scripts/download_auxiliary_pretrained_models.py

基本上,我们将这些路径正确设置为训练或推理过程。 如果您想更改这些路径,请访问文件configs/paths_config.py并修改model_paths字典中的相应记录。

实验

预训练模型

有关我们官方预训练模型的信息,请参阅模型库。请下载我们的预训练模型(W编码器和HyperInverter)并将它们放入pretrained_models文件夹。

训练

第一阶段: W编码器

请按照pSp或e4e的说明训练W编码器,将图像编码为W空间(512)中的潜在代码。

为了节省时间,我们在模型库中发布了用于教堂和人脸的预训练W编码器。

第二阶段: HyperInverter

我们提供了默认的训练脚本,如下所示。如果您有时间,请进一步调整超参数以获得最佳结果。请注意,参数--hidden_dim对模型质量有很大影响,增加它会导致更大的模型,并倾向于有更好的质量。在论文中,我们使用hidden_dim=256。但是,我们建议使用hidden_dim=128来平衡模型大小和模型性能。此外,值得注意的是,我们的代码也很容易修改以预测StyleGAN2的不同层权重。参考weight_shapes并根据需要进行修改。然后,修改hyper_inverter中的get_target_shapes函数以添加您的选项。

  • 人脸
EXPERIMENT_DIR=""
W_ENCODER_PATH=""
GPU_ID=0
CUDA_VISIBLE_DEVICES="$GPU_ID" \
python scripts/train.py \
--dataset_type=ffhq_encode \
--encoder_type=LayerWiseEncoder \
--w_encoder_path="$W_ENCODER_PATH" \
--output_size=1024 \
--exp_dir="$EXPERIMENT_DIR" \
--batch_size=8 \
--batch_size_used_with_adv_loss=4 \
--workers=4 \
--val_interval=1000 \
--save_interval=5000 \
--encoder_optim_name=adam \
--discriminator_optim_name=adam \
--encoder_learning_rate=1e-4 \
--discriminator_learning_rate=1e-4 \
--hyper_lpips_lambda=0.8 \
--hyper_l2_lambda=1.0 \
--hyper_id_lambda=0.1 \
--hyper_adv_lambda=0.005 \
--hyper_d_reg_every=16 \
--hyper_d_r1_gamma=10.0 \
--step_to_add_adversarial_loss=200000 \
--target_shape_name=conv_without_bias  \
--max_steps=500000 \
--hidden_dim=128 \
--num_cold_steps=20000 \
--save_checkpoint_for_resuming_training \
--use_wandb
  • 教堂
EXPERIMENT_DIR=""
W_ENCODER_PATH=""
GPU_ID=0

CUDA_VISIBLE_DEVICES="$GPU_ID" \
python scripts/train.py \
--dataset_type=church_encode \
--encoder_type=ResNetLayerWiseEncoder \
--w_encoder_path="$W_ENCODER_PATH" \
--output_size=256 \
--exp_dir="$EXPERIMENT_DIR" \
--batch_size=8 \
--batch_size_used_with_adv_loss=4 \
--workers=4 \
--val_interval=1000 \
--save_interval=5000 \
--encoder_optim_name=adam \
--discriminator_optim_name=adam \
--encoder_learning_rate=1e-4 \
--discriminator_learning_rate=1e-4 \
--hyper_lpips_lambda=0.8 \
--hyper_l2_lambda=1.0 \
--hyper_id_lambda=0.5 \
--hyper_adv_lambda=0.15 \
--hyper_d_reg_every=16 \
--hyper_d_r1_gamma=100.0 \
--step_to_add_adversarial_loss=100000 \
--target_shape_name=conv_without_bias \
--max_steps=500000 \
--hidden_dim=128 \
--num_cold_steps=10000 \
--save_checkpoint_for_resuming_training \
--use_wandb

推理

  • 将输入图像放入文件夹中。

  • 预处理(如果需要)。 对于人脸领域,如果输入图像尚未裁剪和对齐,请运行以下脚本预处理数据。

RAW_IMAGE_DIR=""
PROCESSED_IMAGE_DIR=""

python scripts/align_all_parallel.py \
--raw_dir "$RAW_IMAGE_DIR" \
--saved_dir "$PROCESSED_IMAGE_DIR" \
--num_threads 8

参数说明如下。

参数描述
RAW_IMAGE_DIR包含原始输入图像的文件夹路径
PROCESSED_IMAGE_DIR保存处理后输入图像的文件夹路径
  • 运行推理

在运行以下脚本之前正确设置参数。

INPUT_DATA_DIR=""
RESULT_DIR=""
MODEL_PATH=""
GPU_ID=0

CUDA_VISIBLE_DEVICES="$GPU_ID" \
python scripts/inference.py \
--exp_dir="$RESULT_DIR" \
--checkpoint_path="$MODEL_PATH" \
--data_path="$INPUT_DATA_DIR" \
--batch_size=4 \
--workers=4

参数说明如下。

参数描述
RESULT_DIR保存推理结果的目录
MODEL_PATHHyperInverter模型的路径
INPUT_DATA_DIR包含处理后输入图像的文件夹路径

最后,重建的图像可以在RESULT_DIR/inference_results文件夹中找到。

定量评估

我们已经准备了人脸和教堂模型的定量重建评估脚本。请在这些文件中设置与您的模型兼容的参数,并运行以下命令进行评估过程。

  • 人脸
sh sample_scripts/human_faces_reconstruction_quantitative_evaluation.sh
  • 教堂
sh sample_scripts/church_reconstruction_quantitative_evaluation.sh

定性比较

对于以下实验,请访问文件configs/paths_config.py并使用HyperInverter和其他反演方法的预训练模型的路径更新model_paths字典。对于其他反演方法,请访问它们的Github仓库下载预训练权重。

重建

定性重建比较的示例脚本是:

DOMAIN=""
METHODS=""
INPUT_DATA_DIR=""
SAVED_RESULTS_DIR_NAME=""
MAX_NUM_IMAGES=100
SAVED_SIZE=1024
GPU_ID=0


CUDA_VISIBLE_DEVICES="$GPU_ID" \
python evaluation/reconstruction_comparison.py \
--methods="$METHODS" \
--domain="$DOMAIN" \
--input_data_dir="$INPUT_DATA_DIR"  \
--input_data_id="$SAVED_RESULTS_DIR_NAME" \
--output_dir=outputs \
--saved_embedding_dir=embeddings \
--max_num_images="$MAX_NUM_IMAGES" \
--resize="$SAVED_SIZE"

参数说明如下。

参数描述
DOMAIN输入域,选项为{human_faces, churches}
METHODS反演方法,用逗号分隔,支持的方法有{hyper_inverter, psp, e4e, SG2_plus, SG2, w_encoder, hyper_inverter, restyle_e4e}
INPUT_DATA_DIR包含处理后输入图像的文件夹路径
SAVED_RESULTS_DIR_NAME保存结果的文件夹名称
MAX_NUM_IMAGES要处理的最大图像数量
SAVED_SIZE每种方法保存的图像大小

结果可以在outputs/SAVED_RESULTS_DIR_NAME文件夹中找到。

编辑

编辑比较的示例脚本是:

DOMAIN=""
METHODS=""
DIRECTION=""
INPUT_DATA_DIR=""
SAVED_RESULTS_DIR_NAME=""
MAX_NUM_IMAGES=10
SAVED_SIZE=1024
MIN_MAG=-30  
MAX_MAG=30
STEP=5 
GPU_ID=0

CUDA_VISIBLE_DEVICES="$GPU_ID" \
python evaluation/editing_inference.py \
--methods="$METHODS" \
--domain="$DOMAIN" \
--input_data_dir="$INPUT_DATA_DIR" \
--input_data_id="$SAVED_RESULTS_DIR_NAME" \
--output_dir=outputs \
--saved_embedding_dir=embeddings \
--direction="$DIRECTION" \
--min_factor="$MIN_MAG" \
--max_factor="$MAX_MAG" \
--step="$STEP" \
--max_num_images="$MAX_NUM_IMAGES" \
--resize="$SAVED_SIZE" \
--save_edited_images \
--gif_speed=4

结果可以在outputs/SAVED_RESULTS_DIR_NAME文件夹中找到。请尝试不同的MIN_MAG和MAX_MAG值以获得最佳结果。 参数说明如下所示。

参数描述
DOMAIN输入域，选项包括{human_faces, churches}
METHODS反演方法，用逗号分隔，支持的方法有{hyper_inverter, psp, e4e, SG2_plus, SG2, w_encoder, restyle_e4e}，例如：hyper_inverter,psp,e4e
DIRECTION编辑方向，支持的方向见下表
INPUT_DATA_DIR包含已处理输入图像的文件夹路径
MIN_MAG最小编辑幅度，请调整此参数以获得最佳结果
MAX_MAG最大编辑幅度，请调整此参数以获得最佳结果
STEP从最小幅度到最大幅度的步长
SAVED_RESULTS_DIR_NAME保存结果的文件夹名称
MAX_NUM_IMAGES要处理的最大图像数量
SAVED_SIZE每种方法保存图像的大小

支持的编辑方向如下所列。

  • 人脸
方法编辑方向
GANSpaceeye_openness, trimmed_beard, lipstick face_roundness, nose_length, eyebrow_thickness, head_angle_up, displeased
InterFaceGANage, smile, rotation
StyleCLIPsurprised, afro, angry, beyonce, bobcut, bowlcut, curly_hair, hilary_clinton, depp, mohawk, purple_hair, taylor_swift, trump, zuckerberg
  • 教堂
方法编辑方向
GANSpaceclouds, vibrant, blue_skies, trees

应用：真实图像插值

两张真实图像插值的示例脚本如下：

DOMAIN=""
METHODS=""
PATH_TO_INPUT_IMAGE_1=""
PATH_TO_INPUT_IMAGE_2=""
SAVED_RESULTS_DIR=""
SAVED_RESULTS_FILE_NAME=""
SAVED_SIZE=1024
NUM_STEPS=100
GPU_ID=0

CUDA_VISIBLE_DEVICES="$GPU_ID" \
python evaluation/real_image_interpolation.py \
--domain="$DOMAIN" \
--method="$METHODS" \
--left_image_path="$PATH_TO_INPUT_IMAGE_1" \
--right_image_path="$PATH_TO_INPUT_IMAGE_2" \
--steps="$NUM_STEPS" \
--saved_image_size="$SAVED_SIZE" \
--saved_dir="$SAVED_RESULTS_DIR" \
--saved_file_name="$SAVED_RESULTS_FILE_NAME" \
--save_interpolated_images \
--gif_speed=2

参数说明如下所示。

参数描述
DOMAIN输入域，选项包括{human_faces, churches}
METHODS反演方法，用逗号分隔，支持的方法有{hyper_inverter, psp, e4e, SG2_plus, SG2, w_encoder, restyle_e4e}，例如：hyper_inverter,psp,e4e
PATH_TO_INPUT_IMAGE_1输入图像1的路径
PATH_TO_INPUT_IMAGE_2输入图像2的路径
NUM_STEPS插值步数
SAVED_RESULTS_DIR保存结果的文件夹路径
MAX_NUM_IMAGES要处理的最大图像数量
SAVED_RESULTS_FILE_NAMEGIF结果文件的名称
SAVED_SIZE每种方法保存图像的大小

致谢

我们的源代码是基于特拉维夫大学研究小组一系列优秀的StyleGAN反演研究的代码库开发的，包括：pSp、e4e、ReStyle和PTI。

对于辅助预训练模型，我们特别感谢TreB1eN、MoCov2、CurricularFace和MTCNN。对于编辑方向，感谢GANSpace、InterFaceGAN和StyleCLIP的作者。

我们使用了StyleGAN2-ADA的PyTorch实现作为StyleGAN模型。所有预训练的StyleGAN模型来自StyleGAN2的官方发布。我们使用作者的官方脚本将原始TensorFlow代码导出的权重转换为与StyleGAN2-ADA的PyTorch版本兼容。

总的来说，非常感谢这些作者的优秀工作和发布源代码及预训练权重的努力。

联系方式

如果您有任何问题，请发送电子邮件至_tan.m.dinh.vn@gmail.com_或在此存储库中提出问题。

相关项目
项目侧边栏1项目侧边栏2
推荐项目
Project Cover

豆包MarsCode

豆包 MarsCode 是一款革命性的编程助手，通过AI技术提供代码补全、单测生成、代码解释和智能问答等功能，支持100+编程语言，与主流编辑器无缝集成，显著提升开发效率和代码质量。

Project Cover

AI写歌

Suno AI是一个革命性的AI音乐创作平台，能在短短30秒内帮助用户创作出一首完整的歌曲。无论是寻找创作灵感还是需要快速制作音乐，Suno AI都是音乐爱好者和专业人士的理想选择。

Project Cover

有言AI

有言平台提供一站式AIGC视频创作解决方案，通过智能技术简化视频制作流程。无论是企业宣传还是个人分享，有言都能帮助用户快速、轻松地制作出专业级别的视频内容。

Project Cover

Kimi

Kimi AI助手提供多语言对话支持，能够阅读和理解用户上传的文件内容，解析网页信息，并结合搜索结果为用户提供详尽的答案。无论是日常咨询还是专业问题，Kimi都能以友好、专业的方式提供帮助。

Project Cover

阿里绘蛙

绘蛙是阿里巴巴集团推出的革命性AI电商营销平台。利用尖端人工智能技术，为商家提供一键生成商品图和营销文案的服务，显著提升内容创作效率和营销效果。适用于淘宝、天猫等电商平台，让商品第一时间被种草。

Project Cover

吐司

探索Tensor.Art平台的独特AI模型，免费访问各种图像生成与AI训练工具，从Stable Diffusion等基础模型开始，轻松实现创新图像生成。体验前沿的AI技术，推动个人和企业的创新发展。

Project Cover

SubCat字幕猫

SubCat字幕猫APP是一款创新的视频播放器，它将改变您观看视频的方式！SubCat结合了先进的人工智能技术，为您提供即时视频字幕翻译，无论是本地视频还是网络流媒体，让您轻松享受各种语言的内容。

Project Cover

美间AI

美间AI创意设计平台，利用前沿AI技术，为设计师和营销人员提供一站式设计解决方案。从智能海报到3D效果图，再到文案生成，美间让创意设计更简单、更高效。

Project Cover

AIWritePaper论文写作

AIWritePaper论文写作是一站式AI论文写作辅助工具，简化了选题、文献检索至论文撰写的整个过程。通过简单设定，平台可快速生成高质量论文大纲和全文，配合图表、参考文献等一应俱全，同时提供开题报告和答辩PPT等增值服务，保障数据安全，有效提升写作效率和论文质量。

使用协议隐私政策广告服务
投诉举报邮箱: service@vectorlightyear.com
@2024 懂AI·鲁ICP备2024100362号-6·鲁公网安备37021002001498号