DragGAN 完全自由 P 图指南

在上篇中，树先生教我们怎么正确部署 DragGAN 项目，完成自在拖拽式 P 图。

超炫酷项目来袭！DragGAN 正式开源了！

但可惜只能运用项目预置的一些图片，本篇教我们怎么利用该项目自在修正修正任何图片。

这儿主要运用到 PTI 项目，能够将你自定义的图片训练成 StyleGAN 潜空间模型，从而完成任何图片的修正修正。

环境预备

这儿我们还是继续选择 AutoDL 云渠道，运用 Python 3.8，CUDA 11.8 的镜像，这个镜像的环境满足项目要求。

下载源码

git clone https://github.com/danielroich/PTI.git
# 装置依靠
cd PTI && pip install -r requirements.txt

你必定猎奇，项目里没有 requirements.txt 文件呀？定心，我给你预备好了~

torch>=2.0.0
scipy
Ninja==1.10.2
gradio>=3.35.2
imageio-ffmpeg>=0.4.3
huggingface_hub
hf_transfer
pyopengl
imgui
glfw==2.6.1
pillow>=9.4.0
torchvision>=0.15.2
imageio>=2.9.0
dlib
wandb
lpips

下载预训练模型

这儿我们只要下载前 2 个即可，分别是 StyleGAN 的生成器文件 ffhq.pkl 和图片的预处理器文件 align.dat，下载完成后放到项目的 pretrained_models 目录下。

图片预处理

主要是完成原始图片人脸关键点检测作业，将你想要修正的图片上传到项目的 image_original 目录下，然后将该目录的绝对路径写入 utils/align_data.py 文件中。

if __name__ == "__main__":
    pre_process_images(f'/root/autodl-tmp/PTI/image_original')

同时修正 configs/paths_config.py 文件参数。

### Input dir, where the images reside
input_data_path = '/root/autodl-tmp/PTI/image_processed'

然后运转。

cp utils/align_data.py .
python align_data.py

运用 PTI 进行 GAN 反演

反演是指将一个图画映射到生成模型的潜空间中，然后经过调整潜空间向量来修正图画的外观。经过这种方法，能够完成对图画的各种修正操作，例如改动姿势、修正外貌特征或添加不同的风格。经过修正潜空间，能够完成对图画的高级修正，同时保持图画的真实性和准确性。

本项目中经过如下指令即可完成图画反演作业。

cp run_pti.py .
python run_pti.py

保存为 DragGAN 可识别的模型文件

上述经过 PTI 反演后的文件不是 DragGAN 可识别的模型文件格局，所以这儿额外处理一下，将 pt 文件转换成 pkl 文件格局，转换脚本递上。

import os
import sys
import pickle
import numpy as np
from PIL import Image
import torch
from configs import paths_config, hyperparameters, global_config
def load_generators(model_id, image_name):
  with open(paths_config.stylegan2_ada_ffhq, 'rb') as f:
    old_G = pickle.load(f)['G_ema'].cuda()
  with open(f'{paths_config.checkpoints_dir}/model_{model_id}_{image_name}.pt', 'rb') as f_new:
    new_G = torch.load(f_new).cuda()
  return old_G, new_G
def export_updated_pickle(new_G,model_id):
  print("Exporting large updated pickle based off new generator and ffhq.pkl")
  with open(paths_config.stylegan2_ada_ffhq, 'rb') as f:
    d = pickle.load(f)
    old_G = d['G_ema'].cuda() ## tensor
    old_D = d['D'].eval().requires_grad_(False).cpu()
  tmp = {}
  tmp['G'] = old_G.eval().requires_grad_(False).cpu()# copy.deepcopy(new_G).eval().requires_grad_(False).cpu()
  tmp['G_ema'] = new_G.eval().requires_grad_(False).cpu() # copy.deepcopy(new_G).eval().requires_grad_(False).cpu()
  tmp['D'] = old_D
  tmp['training_set_kwargs'] = None
  tmp['augment_pipe'] = None
  with open(f'{paths_config.checkpoints_dir}/stylegan2_custom_512_pytorch.pkl', 'wb') as f:
      pickle.dump(tmp, f)
if __name__ == "__main__": 
  # checkpoints 目录下 pt 文件名的一部分
  model_id = "BWISZTGIKPZT"
  # 图片名
  image_name = "myimg"   
  generator_type = image_name
  old_G, new_G = load_generators(model_id, generator_type)
  export_updated_pickle(new_G,model_id)

最后将 checkpoints 目录下生成的模型文件和对应的 embeddings 目录下的文件放入 DragGAN 项目的 checkpoints 目录下，然后重启 DragGAN，功德圆满！

作用展现

一键整合包

想必你看了上面这一通操作，心里现已静静打了退堂鼓，为了便利我们直接运用，我现已将上述步骤整合成 Jupyter Notebook 文档，一键即可运转，完成图画反演！

食用宝典： align.dat 放入项目 pretrained_models 目录下，visualizer_drag_gradio_custom.py 放入项目根目录下，ipynb 文件直接运转即可。

获取方法： 原文结尾处