【PyTorch】TargetCLIPというものを使わせて頂きました

PyTorch GAN

はじめに

TargetCLIPというのを使わせて頂いた時の手順を記録しておきます。

こちらを使わせて頂きました。
github.com
TargetイメージとSourceイメージを与えると、SourceイメージがTargetイメージに近づくように変換されます。

ベースにはStyleGANを使っているようです。
学習済みモデルを公開してくれているので使わせて頂きました。

動作環境

Windowsでは動かなかったのでWSL2上のUbuntuで動かしました。
またGPUがないPCでも動きませんでした。

Ubuntu 20.04 LTS on WSL2（Windows 11）
Python 3.8.10

方法の概略

TargetイメージとSourceイメージの2つの画像を使用するのですがどちらもそのままでは使えません。

それぞれの潜在変数（Latent variable）をあらかじめ求めておく必要があります。

その時に使用するのがこちらです。
github.com

Python環境の構築

CMakeのインストール

こちらはPythonにdlibをインストールするために必要です。

sudo apt install cmake

これでダメならこちらもインストールして下さい。

 sudo apt install build-essential

必要なモジュールのインストール

venvで仮想環境を作って以下をインストールしました。

pip install torch==1.10.0+cu111 torchvision==0.11.1+cu111 -f https://download.pytorch.org/whl/cu111/torch_stable.html
pip install matplotlib
pip install scipy
pip install dlib
pip install ninja

方法

使用する画像の潜在変数を求める

リポジトリのクローンとdlibデータのダウンロード

まずはWSL2上のUbuntuを立ち上げて以下を実行します。

git clone https://github.com/omertov/encoder4editing.git
cd encoder4editing
wget http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2
bzip2 -dk shape_predictor_68_face_landmarks.dat.bz2

学習済みモデルのダウンロード

こちらから「e4e_ffhq_encode.pt」というファイルをダウンロードします。
このファイルは「encoder4editing」フォルダ内（直下）に保存します。

画像の準備

TargetとSourceの二つの画像を用意します。アニメ画像はうまく顔のパーツが識別されずにエラーが起きることがあります。
二つの画像はそれぞれ「target.jpg」と「source.jpg」という名前で「encoder4editing」フォルダ内（直下）に保存します。

Pythonスクリプトの実行

以下のスクリプトを「make_latent.py」という名前で「encoder4editing」フォルダ内（直下）に保存します。

from argparse import Namespace
import numpy as np
import sys

import torch
from torchvision import transforms

from models.psp import pSp

image_path = sys.argv[1]
out_path = sys.argv[2]

model_path = 'e4e_ffhq_encode.pt'
ckpt = torch.load(model_path)

opts = ckpt['opts']
opts['checkpoint_path'] = model_path

opts= Namespace(**opts)
net = pSp(opts)
net.eval().to('cuda')

import dlib
from utils.alignment import align_face

def run_alignment(image_path):
  predictor = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")
  aligned_image = align_face(filepath=image_path, predictor=predictor) 
  print("Aligned image has shape: {}".format(aligned_image.size))
  return aligned_image 

input_image = run_alignment(image_path)

transform_fn = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])

transformed_image = transform_fn(input_image)

with torch.no_grad():
    images, latents = net(transformed_image.unsqueeze(0).to('cuda').float(), randomize_noise=False, return_latents=True)

latent_array = latents.to('cpu').numpy()

np.save(out_path, latent_array)

そして次のように実行します。

python make_latent.py target.jpg target
python make_latent.py source.jpg source

これで「target.npy」と「source.npy」が作成されます。これらは後から使用します。

TargetCLIPを実行する

リポジトリのクローン

いったん「encoder4editing」フォルダから抜けて新たにリポジトリをクローンします。

git clone https://github.com/hila-chefer/TargetCLIP.git
cd TargetCLIP

学習済みモデルのダウンロード

こちらから「stylegan2-ffhq-config-f.pt」というファイルをダウンロードします。
このファイルは「TargetCLIP」フォルダ内（直下）に保存します。

潜在変数の移動

先ほど作った「target.npy」と「source.npy」を「TargetCLIP」フォルダ内（直下）に移動します。

Pythonスクリプトの実行

以下のスクリプトを「transfer.py」という名前で「TargetCLIP」フォルダ内（直下）に保存します。

import os

import numpy as np
import torch
import torchvision

from models.stylegan2.model import Generator
import math
import copy

checkpoint = 'stylegan2-ffhq-config-f.pt'

g_ema = Generator(1024, 512, 8)
g_ema.load_state_dict(torch.load(checkpoint)["g_ema"], strict=False)
g_ema.eval().to('cuda')

source = torch.from_numpy(np.load('source.npy'))
target = torch.from_numpy(np.load('target.npy'))

os.makedirs('result', exist_ok = True)
for alpha in [x / 10 for x in range(0, 21)]:
    
    source_amp, _ = g_ema([(source + target * alpha).to('cuda')], 
                            input_is_latent=True,
                            randomize_noise=False)

    torchvision.utils.save_image(source_amp, "result/results_manipulated_%f.png"%alpha, normalize=True, range=(-1, 1))

そして次のように実行します。

python transfer.py