私的AI研究会 > GAN-e4e
「StyleGAN」で画像編集する。
StyleGAN を使った画像編集では「潜在変数の推定」プロセスに時間が掛かる。
これをを高速に実行できる e4e（encoder4editing）という技術について、上記サイトに従い検証してみる。
e4e は StyleGAN の学習済みモデルを使用して、画像を入力すると目的の潜在変数を直接出力する専用エンコーダを作る手法。Encoderに画像を入力すると１つの潜在変数 w とN個のオフセット△が出力され、これらを合成し N個の潜在変数とし Pretrained StyleGAN に入力する。このとき、元画像と出力画像の誤差を表すロスに加えて、オフセットの分散を表すロスを設定し、これら２つのロスの合計を最小化するように、Encoder のパラメータを学習する。
# --- セットアップ ---
import os
os.chdir('/content')
CODE_DIR = 'encoder4editing'
!git clone https://github.com/cedro3/encoder4editing.git $CODE_DIR
!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
os.chdir(f'./{CODE_DIR}')
from argparse import Namespace
import time
import os
import sys
import numpy as np
from PIL import Image
import torch
import torchvision.transforms as transforms
sys.path.append(".")
sys.path.append("..")
from utils.common import tensor2im
from models.psp import pSp # we use the pSp framework to load the e4e encoder.
%load_ext autoreload
%autoreload 2
# 学習済みパラメータのダウンロード
! pip install --upgrade gdown
import os
import gdown
os.makedirs('pretrained_models', exist_ok=True)
gdown.download('https://drive.google.com/u/1/uc?id=1Du_8FzOPKJhk6aJmiOBhAWVe3_6vAyET', 'pretrained_models/e4e_ffhq_encode.pt', quiet=False)
# ランドマークデータのダウンロード
! wget http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2
! bzip2 -dk shape_predictor_68_face_landmarks.dat.bz2
# モデルに学習済みパラメータをロード
model_path = 'pretrained_models/e4e_ffhq_encode.pt' ####
ckpt = torch.load(model_path, map_location='cpu')
opts = ckpt['opts']
opts['checkpoint_path'] = model_path
opts= Namespace(**opts)
net = pSp(opts)
net.eval()
net.cuda()
print('Model successfully loaded!')
# --- 顔画像の切り出し ---
import os
import shutil
from tqdm import tqdm
if os.path.isdir('align'):
shutil.rmtree('align')
os.makedirs('align', exist_ok=True)
def run_alignment(image_path):
import dlib
from utils.alignment import align_face
predictor = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")
aligned_image = align_face(filepath=image_path, predictor=predictor)
return aligned_image
path = './images'
files = sorted(os.listdir(path))
for i, file in enumerate(tqdm(files)):
if file=='.ipynb_checkpoints':
continue
input_image = run_alignment(path+'/'+file)
input_image.resize((256,256))
input_image.save('./align/'+file)
# --- 潜在変数の推定 ---
if os.path.isdir('vec_pic'):
shutil.rmtree('vec_pic')
os.makedirs('vec_pic', exist_ok=True)
if os.path.isdir('vec'):
shutil.rmtree('vec')
os.makedirs('vec', exist_ok=True)
img_transforms = transforms.Compose([
transforms.Resize((256, 256)),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
path = './align'
files = sorted(os.listdir(path))
for i, file in enumerate(tqdm(files)):
if file=='.ipynb_checkpoints':
continue
input_image = Image.open(path+'/'+file)
transformed_image = img_transforms(input_image)
with torch.no_grad():
images, latents = net(transformed_image.unsqueeze(0).to('cuda').float(), randomize_noise=False, return_latents=True)
result_image, latent = images[0], latents[0]
tensor2im(result_image).save('./vec_pic/'+file) # vec_pic 保存
torch.save(latents, './vec/'+file[:-4]+'.pt') # vec 保存
%matplotlib inline・変更した以下のセルを実行する
# --- 元画像と生成画像の表示 ---
%matplotlib inline
import matplotlib.pyplot as plt
from PIL import Image
import os
def display_pic(folder):
fig = plt.figure(figsize=(30, 40))
files = os.listdir(folder)
files.sort()
for i, file in enumerate(files):
img = Image.open(folder+'/'+file)
images = np.asarray(img)
ax = fig.add_subplot(10, 10, i+1, xticks=[], yticks=[])
image_plt = np.array(images)
ax.imshow(image_plt)
ax.set_xlabel(folder+'/'+file, fontsize=15)
plt.show()
plt.close()
display_pic('align')
display_pic('vec_pic')
・上段が切り出した顔画像、下段が潜在変数から生成した画像~
#@title 設定
latent = "03.pt"#@param {type:"string"}
direction = "pose" #@param ["age", "pose", "smile", "age+pose"] {allow-input: true}
min = -50 #@param {type:"slider", min:-50, max:0, step:10}
max = 50 #@param {type:"slider", min:0, max:50, step:10}
# --- 静止画の生成 ---
import os
import shutil
if os.path.isdir('pic'):
shutil.rmtree('pic')
os.makedirs('pic', exist_ok=True)
from editings import latent_editor
from tqdm import trange
folder = 'vec'
latents = torch.load(folder+'/'+latent)
editor = latent_editor.LatentEditor(net.decoder, False)
interfacegan_directions = {
'age': 'editings/interfacegan_directions/age.pt',
'smile': 'editings/interfacegan_directions/smile.pt',
'pose': 'editings/interfacegan_directions/pose.pt',
'age+pose': 'editings/interfacegan_directions/age+pose.pt'
}
interfacegan_direction = torch.load(interfacegan_directions[direction]).cuda()
cnt = 0
for i in trange(0, min, -1, desc='0 -> min'):
result = editor.apply_interfacegan(latents, interfacegan_direction, factor=i).resize((512,512))
result.save('./pic/'+str(cnt).zfill(6)+'.jpg')
cnt +=1
for i in trange(min, max, desc='min -> max'):
result = editor.apply_interfacegan(latents, interfacegan_direction, factor=i).resize((512,512))
result.save('./pic/'+str(cnt).zfill(6)+'.jpg')
cnt +=1
for i in trange(max, 0, -1, desc='max -> 0'):
result = editor.apply_interfacegan(latents, interfacegan_direction, factor=i).resize((512,512))
result.save('./pic/'+str(cnt).zfill(6)+'.jpg')
cnt +=1
# --- mp4動画の作成 ---
# 既に output.mp4 があれば削除する
import os
if os.path.exists('./output.mp4'):
os.remove('./output.mp4')
# pic フォルダーの静止画から動画を作成
! ffmpeg -r 30 -i pic/%6d.jpg\
-vcodec libx264 -pix_fmt yuv420p output.mp4
# movieフォルダへ名前を付けてコピー
import shutil
os.makedirs('movie', exist_ok=True)
shutil.copy('output.mp4', 'movie/'+direction+'_'+latent[:-3]+'.mp4')
# --- mp4動画の再生 ---
from IPython.display import HTML
from base64 import b64encode
mp4 = open('./output.mp4', 'rb').read()
data_url = 'data:video/mp4;base64,' + b64encode(mp4).decode()
HTML(f"""
<video width="50%" height="50%" controls>
<source src="{data_url}" type="video/mp4">
</video>""")