2024年-火爆全网的DragGAN 最新图片编辑项目超基础教学

2024年-火爆全网的DragGAN 最新图片编辑项目超基础教学

不同于 Stable Diffusion 与 Midjourney 使用的扩散模型，DragGAN 是一款基于生成对抗网络（GANs）技术的图片编辑项目。追根溯源，其实还是 NVIDIA 的 StyleGAN 的衍生项目。

DragGAN 由两个主要部分组成： 第一部分是基于特征的运动监督，通过控制起始点向目标点运动，实现图像的变形； 第二部分是新的点跟踪方法，即最近邻检索，通过在相同的特征空间上进行点跟踪，提高了效率，并避免了累计误差损失。

简单来说，它能够让你轻松掌控图片中的人物和动物，通过拖拽控制点和目标点来调整他们的位置、形状、姿势等等。目前这个项目在 Github 上已经斩获了 20K star 了，趋势榜第一，足见其火爆程度！

部署要求

• 支持 Linux 和 Windows，建议使用 Linux
• 内存 12 GB 以上
• 1–8 NVIDIA 高端 GPU
• Python >= 3.8，PyTorch >= 1.9.0
• CUDA toolkit >= 11.1

安装部署

安装过程，真是一把鼻涕一把泪，到处都是坑，先是在阿里 PAI 平台栽跟头，折腾好久才发现是基础环境问题，后来又在安装方法上翻了车，官方的安装指导也是各种报错。

一路艰难摸索过来，现在将实测可行的一套部署方法带给大家。

这里我们选择 AutoDL 云平台，使用 Python 3.8，CUDA 11.8 的镜像，这个镜像的环境就已经满足项目要求了。

下载源码

代码语言：javascript

复制

git clone https://github.com/XingangPan/DragGAN.git
1

安装依赖

代码语言：javascript

复制

cd DragGAN
pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
1
2

如果出现下面的报错，代表这个 pip 源里没有对应的版本。

修改 requirements.txt 文件，去除特定版本依赖，再次执行即可。

代码语言：javascript

复制

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
1
2
3
4
5
6
7
8
9
10
11
12
13

下载模型

代码语言：javascript

复制

sh scripts/download_model.sh
1

运行

这里选择运行 WebUI，看到 URL 就代表启动成功了。

代码语言：javascript

复制

# 运行 WebUI
python visualizer_drag_gradio.py
1
2

效果展示

使用上非常简单，鼠标在图中点点就行，红点是起始点，蓝点是目标点，实现图像的拖拽效果。

还可以控制区域图像，避免整体图像跟着变化。

很多人都以为这个项目能直接选一张照片，就能像 PS 一样进行 P 图了，其实并不是，至少目前还不支持，上面你看到的这些图片都是经过训练而生成的模型，所以本质上，你是在操作一个模型，而非图片！

如果你想操作自己的图片，那么按官网说的，你需要使用到 PTI 项目，将你自定义的图片训练成 StyleGAN 模型。

项目自由编辑修改任何图片

这里主要使用到 PTI 项目，可以将你自定义的图片训练成 StyleGAN 潜空间模型，从而实现任何图片的编辑修改。

环境准备

这里我们还是继续选择 AutoDL 云平台，使用 Python 3.8，CUDA 11.8 的镜像，这个镜像的环境满足项目要求。

下载源码

代码语言：javascript

复制

git clone https://github.com/danielroich/PTI.git
# 安装依赖
cd PTI && pip install -r requirements.txt
1
2
3

你肯定好奇，项目里没有 requirements.txt 文件呀？放心，我给你准备好了~

代码语言：javascript

复制

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
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

下载预训练模型

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

图片预处理

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

代码语言：javascript

复制

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

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

代码语言：javascript

复制

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

然后运行。

代码语言：javascript

复制

export PYTHONPATH=$PYTHONPATH:/root/auto-tmp/PTI
python align_data.py
1
2

使用 PTI 进行 GAN 反演

反演是指将一个图像映射到生成模型的潜空间中，然后通过调整潜空间向量来修改图像的外观。通过这种方式，可以实现对图像的各种编辑操作，例如改变姿势、修改外貌特征或添加不同的风格。通过编辑潜空间，可以实现对图像的高级编辑，同时保持图像的真实性和准确性。

本项目中通过如下命令即可完成图像反演工作。

代码语言：javascript

复制

python run_pti.py
1

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

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

代码语言：javascript

复制

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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

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

效果展示

一键整合包

想必你看了上面这一通操作，心里已经默默打了退堂鼓，为了方便大家直接使用，我已经将上述步骤整合成 Jupyter Notebook 文档，一键即可运行，完成图像反演！

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

)]

[外链图片转存中…(img-ICCZAI44-1723803867494)]

[外链图片转存中…(img-7KgxxKw7-1723803867495)]

一键整合包

想必你看了上面这一通操作，心里已经默默打了退堂鼓，为了方便大家直接使用，我已经将上述步骤整合成 Jupyter Notebook 文档，一键即可运行，完成图像反演！

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

博客原文一键整合包地址：专业人工智能技术开源社区