Pivotal Tuning for Latent-based Editing of Real Images 复现实验 ( 1 )-优快云博客

本文链接：https://blog.youkuaiyun.com/weixin_48981284/article/details/136601664

原论文：Pivotal Tuning for Latent-based Editing of Real Images
论文：https://arxiv.org/abs/2106.05744
代码：https://github.com/danielroich/PTI

本文主要目的在于讲述该篇论文复现流程和代码详解。

在下载相关论文代码并解压后，可以看到在"./notebooks"文件夹中有两个 jupyter notebook 文件。其中 “inference_playground.ipynb” 为作者给出的简易复现流程，本文主要通过完整的运行和讲解一遍该流程来简要说明该论文的实现方法。其他代码的详细解析将会在其他文章中再叙。

如果您的研究方向与生成对抗网络 ( Generative Adversarial Networks , GAN ) 相关，就算不了解这篇论文的具体内容，我相信这篇文章也能对您的实验起到一点点小小的帮助。

0: 事前准备

建议使用 Anaconda 构建一个虚拟环境以运行该实验流程。

我的环境配置：

win11系统
python 3.8.18
pytorch 2.21
pytorch-cuda 12.1
torchvision 0.17.1
ninja 1.11.1.1
dlib 19.24.2
matplotlib 3.7.2
lpips 0.1.4

倘若运行过程中遇到相关库未安装的问题，请自行运行命令"conda install XXX" 或 “pip install XXX” 来解决。(如果是新建的虚拟环境是肯定会遇见的)

1: Step 0~1 Install required packages

from pydrive.auth import GoogleAuth
from pydrive.drive import GoogleDrive
from google.colab import auth
from oauth2client.client import GoogleCredentials

download_with_pydrive = True #@param {type:"boolean"}  

class Downloader(object):
    def __init__(self, use_pydrive):
        self.use_pydrive = use_pydrive

        if self.use_pydrive:
            self.authenticate()
        
    def authenticate(self):
        auth.authenticate_user()
        gauth = GoogleAuth()
        gauth.credentials = GoogleCredentials.get_application_default()
        self.drive = GoogleDrive(gauth)
    
    def download_file(self, file_id, file_dst):
        if self.use_pydrive:
            downloaded = self.drive.CreateFile({'id':file_id})
            downloaded.FetchMetadata(fetch_all=True)
            downloaded.GetContentFile(file_dst)
        else:
            !gdown --id $file_id -O $file_dst

downloader = Downloader(download_with_pydrive)

这一段代码主要是为了下文中下载相关模型加速，但这段代码只能在linux系统上运行，且需要您有个 google drive 账号并进行一系列关联操作。

如果您满足以上条件可以将这段代码保留，否则您也可以像我一样将这段代码中导入包的部分 ( 前四行 ) 注释掉，并修改 ownload_with_pydrive = False。

## Other packages are already builtin in the Colab interpreter
!pip install wandb
!pip install lpips

## Used for faster inference of StyleGAN by enabling C++ code compilation

!wget https://github.com/ninja-build/ninja/releases/download/v1.8.2/ninja-linux.zip
!sudo unzip ninja-linux.zip -d /usr/local/bin/
!sudo update-alternatives --install /usr/bin/ninja ninja /usr/local/bin/ninja 1 --force

如果您已经安装了相关包，这一段也可以注释掉。

Step 2 - Download Pretrained models

import os
os.chdir('/content')
CODE_DIR = 'PTI'

!git clone https://github.com/danielroich/PTI.git $CODE_DIR

os.chdir(f'./{CODE_DIR}')

os.chdir() 函数主要用于更改工作区间。在这里更改工作区间为’.content’。如果您使用跟我一样的Windows系统，在这里指向的是"C:\content"文件夹，如果提示找不到该文件可以自行手动创建。

!git clone https://github.com/danielroich/PTI.git $CODE_DIR

这行代码主要是下载论文其他代码到 “content\PTI” 中。这里下载的代码和之前下载的完整实验代码是一样的，如果您不想再下一次可以直接再次解压之前的实验代码到该目录( 在我这里是 C:\content\PTI ) 下。

import os
import sys
import pickle
import numpy as np
from PIL import Image
import torch
from configs import paths_config, hyperparameters, global_config
from utils.align_data import pre_process_images
from scripts.run_pti import run_PTI
from IPython.display import display
import matplotlib.pyplot as plt
from scripts.latent_editor_wrapper import LatentEditorWrapper

导入相关包。

current_directory = os.getcwd()
save_path = os.path.join(os.path.dirname(current_directory), CODE_DIR, "pretrained_models")
os.makedirs(save_path, exist_ok=True)

os.getcwd() 函数为获取当前的工作区间( C:\content\PTI )，save_path 为下文下载的预训练模型的保存路径。在这里 current_directory = ‘C:\content\PTI’ ; save_path = ‘C:\content\PTI\pretrained_models’. os.makedirs() 在这时创建 save_path 对应的文件夹。

## Download pretrained StyleGAN on FFHQ 1024x1024
downloader.download_file("125OG7SMkXI-Kf2aqiwLLHyCvSW-gZk3M", os.path.join(save_path, 'ffhq.pkl'))

## Download Dlib tool for alingment, used for preprocessing images before PTI
downloader.download_file("1xPmn19T6Bdd-_RfCVlgNBbfYoh1muYxR", os.path.join(save_path, 'align.dat'))

下载预训练的 StyleGAN2 和 Dlib 模型。其中 StyleGAN2 主要是在图像生成阶段使用，Dlib模型主要是在图片与处理阶段使用。在这里是通过 gdown 函数来进行下载的，如果运行时出现乱码说明下载失败，可以模仿【小教程】Ubuntu gdown下载Google Drive文件-优快云博客中下载方法手动进入下载页面下载文件到指定文件夹 ( save_path ) 。

Step 3 - Configuration Setup

image_dir_name = 'image'

## If set to true download desired image from given url. If set to False, assumes you have uploaded personal image to
## 'image_original' dir

use_image_online = True
image_name = 'personal_image'
use_multi_id_training = False

global_config.device = 'cuda'

## 对比基准模型 E4E 的保存路径 
paths_config.e4e = '/content/PTI/pretrained_models/e4e_ffhq_encode.pt'

## 图像名字
paths_config.input_data_id = image_dir_name

## 预处理后图像保存路径
paths_config.input_data_path = f'/content/PTI/{image_dir_name}_processed'

## 预训练的 StyleGAN2 模型的保存路径
paths_config.stylegan2_ada_ffhq = '/content/PTI/pretrained_models/ffhq.pkl'

## 训练过程中模型参数 w 的保存路径 ( 建议修改 ) 
paths_config.checkpoints_dir = '/content/PTI/'

## 其他预训练模型路径
paths_config.style_clip_pretrained_mappers = '/content/PTI/pretrained_models'
hyperparameters.use_locality_regularization = False

设置全局变量。这一步主要是设置各个路径。在以后运行自己的数据集的时候也需要更改相关设置。use_image_online = True 时，后续会下载测试用的图像。

Step 4 -Preproccess Data

这一部分主要是预处理图像。

os.makedirs(f'./{image_dir_name}_original', exist_ok=True)
os.makedirs(f'./{image_dir_name}_processed', exist_ok=True)
os.chdir(f'./{image_dir_name}_original')

创建"content\PTI\image_original"和"content\PTI\image_processed"文件夹。其中 image_original 为原始图像保存地址，image_processed 为预处理后的图像保存地址。

## Download real face image
## If you want to use your own image skip this part and upload an image/images of your choosing to image_original dir
if use_image_online:
  !wget -O personal_image.jpg https://static01.nyt.com/images/2019/09/09/opinion/09Hunter1/09Hunter1-superJumbo.jpg ## Photo of Sarena Wiliams

下载测试用的图片。如果下载不了就点击网页里保存。如果不想用这张图片做测试也可以用自己准备的图片，最好是大于 1024 * 1024 的人脸图像，记得放到 image_original 文件夹中。

original_image = Image.open(f'{image_name}.jpg')
original_image

在这里插入图片描述

image.open() 和 PIL 包的内容可以参考：Python: PIL.Image 基本使用-优快云博客

os.chdir('/content/PTI')

把工作区间改回主要目录，然后接下来是图像预处理的函数调用：

pre_process_images(f'/content/PTI/{image_dir_name}_original')

这一部分详细请看：图像预处理——使用 Dlib 模型预处理图像-优快云博客

看看预处理后的图像是什么样子：

aligned_image = Image.open(f'/content/PTI/{image_dir_name}_processed/{image_name}.jpeg')

aligned_image.resize((512,512))

在这里插入图片描述

在图像预处理完成后接下来就是将这幅图像反演到生成器的 lantent space $W$ 中了，详细请看这个系列的下一篇文章：https://blog.youkuaiyun.com/weixin_48981284/article/details/136629468