Hyperstyleで、高精度で編集しやすい反転を高速に行う

1.はじめに

　StyleGANの画像編集には画像の潜在空間への反転が必要ですが、画像の正確な表現と潜在変数の編集のし易さはトレードオフの関係にあります。今回ご紹介するのは、このトレードオフの改善を高速処理で実現するHyperstyleという技術です。

＊この論文は、2021.11に提出されました。

2.Hyperstyleとは？

　StyleGANで画像編集を行う場合、まず画像を潜在空間へ反転することが必要で、このための手法は数多く提案されています。しかしながら、このときの画像の正確な表現と潜在変数の編集のし易さはトレードオフの関係にあります。

　最近、このトレードオフを軽減するためにジェネレータを微調整する方法が提案されていますが、新しい画像毎に長いトレーニングが必要です。Hyperstyleは、このジェネレータを微調整するアプローチを高速に行うモデルです。

　下記はHyperstyleのフローです。ポイントは、ジェネレータのパラメータ全体のトレーニングを行うのではなく、パラメータを絞り込んでエンコーダで調整します。

　目的の画像（Input）と近似潜在変数（Initial reconstruction）から求めた画像を、エンコーダ（Hypernetwork H）に入力します。エンコーダでは、Resnetを通して特徴マップに展開し、学習済みジェネレータ（Pre-trained Generator）との差分を計算して、ジェネレータ（modified Generator）に適用します。

3.コード

　コードはGoogle Colabで動かす形にしてGithubに上げてありますので、それに沿って説明して行きます。自分で動かしてみたい方は、この「リンク」をクリックし表示されたノートブックの先頭にある「Colab on Web」ボタンをクリックすると動かせます。

　まず、セットアップを行います。

#@title セットアップ
import os
os.chdir('/content')
CODE_DIR = 'hyperstyle'

# clone repo
!git clone https://github.com/cedro3/hyperstyle.git $CODE_DIR

# install ninja
!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}')


# Import Packages
import time
import sys
import pprint
from tqdm import tqdm
import numpy as np
from PIL import Image
import torch
import torchvision.transforms as transforms

import imageio
from IPython.display import HTML
from base64 import b64encode

sys.path.append(".")
sys.path.append("..")

from notebooks.notebook_utils import Downloader, HYPERSTYLE_PATHS, W_ENCODERS_PATHS, run_alignment
from utils.common import tensor2im
from utils.inference_utils import run_inversion
from utils.domain_adaptation_utils import run_domain_adaptation
from utils.model_utils import load_model, load_generator
from function import *

%load_ext autoreload
%autoreload 2


# download pretrained_models
! pip install --upgrade gdown
import gdown
gdown.download('https://drive.google.com/uc?id=1NxGZfkE3THgEfPHbUoLPjCKfpWTo08V1', 'pretrained_models.zip', quiet=False)
! unzip pretrained_models.zip


# set expeiment data
EXPERIMENT_DATA_ARGS = {
    "faces": {
        "model_path": "./pretrained_models/hyperstyle_ffhq.pt",
        "w_encoder_path": "./pretrained_models/faces_w_encoder.pt",
        "image_path": "./notebooks/images/face_image.jpg",
        "transform": transforms.Compose([
            transforms.Resize((256, 256)),
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
    },
    "cars": {
        "model_path": "./pretrained_models/hyperstyle_cars.pt",
        "w_encoder_path": "./pretrained_models/cars_w_encoder.pt",
        "image_path": "./notebooks/images/car_image.jpg",
        "transform": transforms.Compose([
            transforms.Resize((192, 256)),
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
    },
    "afhq_wild": {
        "model_path": "./pretrained_models/hyperstyle_afhq_wild.pt",
        "w_encoder_path": "./pretrained_models/afhq_wild_w_encoder.pt",
        "image_path": "./notebooks/images/afhq_wild_image.jpg",
        "transform": transforms.Compose([
            transforms.Resize((256, 256)),
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
    }
}

experiment_type = 'faces'
EXPERIMENT_ARGS = EXPERIMENT_DATA_ARGS[experiment_type]


# Load HyperStyle Model
model_path = EXPERIMENT_ARGS['model_path']
net, opts = load_model(model_path, update_opts={"w_encoder_checkpoint_path": EXPERIMENT_ARGS['w_encoder_path']})
print('Model successfully loaded!')


# difine function
def generate_mp4(out_name, images, kwargs):
    writer = imageio.get_writer(out_name + '.mp4', **kwargs)
    for image in tqdm(images):
        writer.append_data(image)
    writer.close()


def get_latent_and_weight_deltas(inputs, net, opts):
    opts.resize_outputs = False
    opts.n_iters_per_batch = 5
    with torch.no_grad():
        _, latent, weights_deltas, _ = run_inversion(inputs.to("cuda").float(), net, opts)
    weights_deltas = [w[0] if w is not None else None for w in weights_deltas]
    return latent, weights_deltas
    

def get_result_from_vecs(vectors_a, vectors_b, weights_deltas_a, weights_deltas_b, alpha):
    results = []
    for i in range(len(vectors_a)):
        with torch.no_grad():
            cur_vec = vectors_b[i] * alpha + vectors_a[i] * (1 - alpha)
            cur_weight_deltas = interpolate_weight_deltas(weights_deltas_a, weights_deltas_b, alpha)
            res = net.decoder([cur_vec],
                              weights_deltas=cur_weight_deltas,
                              randomize_noise=False,
                              input_is_latent=True)[0]
            results.append(res[0])
    return results

def interpolate_weight_deltas(weights_deltas_a, weights_deltas_b, alpha):
    cur_weight_deltas = []
    for weight_idx, w in enumerate(weights_deltas_a):
        if w is not None:
            delta = weights_deltas_b[weight_idx] * alpha + weights_deltas_a[weight_idx] * (1 - alpha)
        else:
            delta = None
        cur_weight_deltas.append(delta)
    return cur_weight_deltas
    
def show_mp4(filename, width):
    mp4 = open(filename + '.mp4', 'rb').read()
    data_url = "data:video/mp4;base64," + b64encode(mp4).decode()
    display(HTML("""
    <video width="%d" controls autoplay loop>
        <source src="%s" type="video/mp4">
    </video>
    """ % (width, data_url)))

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

#@title セットアップ

import os

os.chdir('/content')

CODE_DIR = 'hyperstyle'

# clone repo

!git clone https://github.com/cedro3/hyperstyle.git $CODE_DIR

# install ninja

!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}')

# Import Packages

import time

import sys

import pprint

from tqdm import tqdm

import numpy as np

from PIL import Image

import torch

import torchvision.transforms as transforms

import imageio

from IPython.display import HTML

from base64 import b64encode

sys.path.append(".")

sys.path.append("..")

from notebooks.notebook_utils import Downloader, HYPERSTYLE_PATHS, W_ENCODERS_PATHS, run_alignment

from utils.common import tensor2im

from utils.inference_utils import run_inversion

from utils.domain_adaptation_utils import run_domain_adaptation

from utils.model_utils import load_model, load_generator

from function import *

%load_ext autoreload

%autoreload 2

# download pretrained_models

! pip install --upgrade gdown

import gdown

gdown.download('https://drive.google.com/uc?id=1NxGZfkE3THgEfPHbUoLPjCKfpWTo08V1', 'pretrained_models.zip', quiet=False)

! unzip pretrained_models.zip

# set expeiment data

EXPERIMENT_DATA_ARGS = {

"faces": {

"model_path": "./pretrained_models/hyperstyle_ffhq.pt",

"w_encoder_path": "./pretrained_models/faces_w_encoder.pt",

"image_path": "./notebooks/images/face_image.jpg",

"transform": transforms.Compose([

transforms.Resize((256, 256)),

transforms.ToTensor(),

transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])

"cars": {

"model_path": "./pretrained_models/hyperstyle_cars.pt",

"w_encoder_path": "./pretrained_models/cars_w_encoder.pt",

"image_path": "./notebooks/images/car_image.jpg",

"transform": transforms.Compose([

transforms.Resize((192, 256)),

transforms.ToTensor(),

transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])

"afhq_wild": {

"model_path": "./pretrained_models/hyperstyle_afhq_wild.pt",

"w_encoder_path": "./pretrained_models/afhq_wild_w_encoder.pt",

"image_path": "./notebooks/images/afhq_wild_image.jpg",

"transform": transforms.Compose([

transforms.Resize((256, 256)),

transforms.ToTensor(),

transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])

}

experiment_type = 'faces'

EXPERIMENT_ARGS = EXPERIMENT_DATA_ARGS[experiment_type]

# Load HyperStyle Model

model_path = EXPERIMENT_ARGS['model_path']

net, opts = load_model(model_path, update_opts={"w_encoder_checkpoint_path": EXPERIMENT_ARGS['w_encoder_path']})

print('Model successfully loaded!')

# difine function

def generate_mp4(out_name, images, kwargs):

writer = imageio.get_writer(out_name + '.mp4', **kwargs)

for image in tqdm(images):

writer.append_data(image)

writer.close()

def get_latent_and_weight_deltas(inputs, net, opts):

opts.resize_outputs = False

opts.n_iters_per_batch = 5

with torch.no_grad():

_, latent, weights_deltas, _ = run_inversion(inputs.to("cuda").float(), net, opts)

weights_deltas = [w[0] if w is not None else None for w in weights_deltas]

return latent, weights_deltas

def get_result_from_vecs(vectors_a, vectors_b, weights_deltas_a, weights_deltas_b, alpha):

results = []

for i in range(len(vectors_a)):

with torch.no_grad():

cur_vec = vectors_b[i] * alpha + vectors_a[i] * (1 - alpha)

cur_weight_deltas = interpolate_weight_deltas(weights_deltas_a, weights_deltas_b, alpha)

res = net.decoder([cur_vec],

weights_deltas=cur_weight_deltas,

randomize_noise=False,

input_is_latent=True)[0]

results.append(res[0])

return results

def interpolate_weight_deltas(weights_deltas_a, weights_deltas_b, alpha):

cur_weight_deltas = []

for weight_idx, w in enumerate(weights_deltas_a):

if w is not None:

delta = weights_deltas_b[weight_idx] * alpha + weights_deltas_a[weight_idx] * (1 - alpha)

else:

delta = None

cur_weight_deltas.append(delta)

return cur_weight_deltas

def show_mp4(filename, width):

mp4 = open(filename + '.mp4', 'rb').read()

data_url = "data:video/mp4;base64," + b64encode(mp4).decode()

display(HTML("""

</video>

""" % (width, data_url)))

　サンプル画像を見てみましょう。自分の画像を使用したい場合は、./images/picに画像をアップロードしてください。

#@title サンプル画像表示
display_pic('./images/pic')

1 2	#@title サンプル画像表示 display_pic('./images/pic')

　サンプル画像にalign処理（目、鼻、口、顎などを処置の位置に合わせて切り出す）を行います。

#@title align処理
import glob
from tqdm import tqdm

reset_folder('./images/align')
files = sorted(glob.glob('./images/pic/*.jpg'))
for file in tqdm(files):
    aligned_image = run_alignment(file)
    name = os.path.basename(file)
    aligned_image.save('./images/align/'+name)
    
display_pic('./images/align')

#@title align処理

import glob

from tqdm import tqdm

reset_folder('./images/align')

files = sorted(glob.glob('./images/pic/*.jpg'))

for file in tqdm(files):

aligned_image = run_alignment(file)

name = os.path.basename(file)

aligned_image.save('./images/align/'+name)

display_pic('./images/align')

　それでは、hyperstyleを使って反転処理（実写と同じ画像を出力する潜在変数を求める）を行います。

#@title 反転の実行
import glob
image_paths = sorted(glob.glob('./images/align/*.jpg'))

in_images = []
all_vecs = []
all_weights_deltas = []
img_transforms = EXPERIMENT_ARGS['transform']

if experiment_type == "cars":
    resize_amount = (512, 384)
else:
    resize_amount = (opts.output_size, opts.output_size)

for image_path in image_paths:
    #print(f'Working on {os.path.basename(image_path)}...')
    original_image = Image.open(image_path)
    original_image = original_image.convert("RGB")
    input_image = img_transforms(original_image)
    # get the weight deltas for each image
    result_vec, weights_deltas = get_latent_and_weight_deltas(input_image.unsqueeze(0), net, opts)
    all_vecs.append([result_vec])
    all_weights_deltas.append(weights_deltas)
    in_images.append(original_image.resize(resize_amount))

n_transition = 25
if experiment_type == "cars":
    SIZE = 384
else:
    SIZE = opts.output_size

images = []
image_paths.append(image_paths[0])
all_vecs.append(all_vecs[0])
all_weights_deltas.append(all_weights_deltas[0])
in_images.append(in_images[0])

for i in tqdm(range(1, len(image_paths))):
    if i == 0:
        alpha_vals = [0] * 10 + np.linspace(0, 1, n_transition).tolist() + [1] * 5
    else:
        alpha_vals = [0] * 5 + np.linspace(0, 1, n_transition).tolist() + [1] * 5

    for alpha in alpha_vals:
        image_a = np.array(in_images[i - 1])
        image_b = np.array(in_images[i])
        image_joint = np.zeros_like(image_a)
        up_to_row = int((SIZE - 1) * alpha)
        if up_to_row > 0:
            image_joint[:(up_to_row + 1), :, :] = image_b[((SIZE - 1) - up_to_row):, :, :]
        if up_to_row < (SIZE - 1):
            image_joint[up_to_row:, :, :] = image_a[:(SIZE - up_to_row), :, :]

        result_image = get_result_from_vecs(all_vecs[i - 1], all_vecs[i],
                                            all_weights_deltas[i - 1], all_weights_deltas[i],
                                            alpha)[0]
        if experiment_type == "cars":
            result_image = result_image[:, 64:448, :]

        output_im = tensor2im(result_image)
        res = np.concatenate([image_joint, np.array(output_im)], axis=1)
        images.append(res)

#@title 反転の実行

import glob

image_paths = sorted(glob.glob('./images/align/*.jpg'))

in_images = []

all_vecs = []

all_weights_deltas = []

img_transforms = EXPERIMENT_ARGS['transform']

if experiment_type == "cars":

resize_amount = (512, 384)

else:

resize_amount = (opts.output_size, opts.output_size)

for image_path in image_paths:

#print(f'Working on {os.path.basename(image_path)}...')

original_image = Image.open(image_path)

original_image = original_image.convert("RGB")

input_image = img_transforms(original_image)

# get the weight deltas for each image

result_vec, weights_deltas = get_latent_and_weight_deltas(input_image.unsqueeze(0), net, opts)

all_vecs.append([result_vec])

all_weights_deltas.append(weights_deltas)

in_images.append(original_image.resize(resize_amount))

n_transition = 25

if experiment_type == "cars":

SIZE = 384

else:

SIZE = opts.output_size

images = []

image_paths.append(image_paths[0])

all_vecs.append(all_vecs[0])

all_weights_deltas.append(all_weights_deltas[0])

in_images.append(in_images[0])

for i in tqdm(range(1, len(image_paths))):

if i == 0:

alpha_vals = [0] * 10 + np.linspace(0, 1, n_transition).tolist() + [1] * 5

else:

alpha_vals = [0] * 5 + np.linspace(0, 1, n_transition).tolist() + [1] * 5

for alpha in alpha_vals:

image_a = np.array(in_images[i - 1])

image_b = np.array(in_images[i])

image_joint = np.zeros_like(image_a)

up_to_row = int((SIZE - 1) * alpha)

if up_to_row > 0:

image_joint[:(up_to_row + 1), :, :] = image_b[((SIZE - 1) - up_to_row):, :, :]

if up_to_row < (SIZE - 1):

image_joint[up_to_row:, :, :] = image_a[:(SIZE - up_to_row), :, :]

result_image = get_result_from_vecs(all_vecs[i - 1], all_vecs[i],

all_weights_deltas[i - 1], all_weights_deltas[i],

alpha)[0]

if experiment_type == "cars":

result_image = result_image[:, 64:448, :]

output_im = tensor2im(result_image)

res = np.concatenate([image_joint, np.array(output_im)], axis=1)

images.append(res)

　それでは、求めた潜在変数を使用してモーフィングする動画を作成してみましょう。

#@title 動画の作成
kwargs = {'fps': 15}
save_path = "./notebooks/animations"
os.makedirs(save_path, exist_ok=True)

gif_path = os.path.join(save_path, f"{experiment_type}_gif")
generate_mp4(gif_path, images, kwargs)
show_mp4(gif_path, width=opts.output_size)

#@title 動画の作成

kwargs = {'fps': 15}

save_path = "./notebooks/animations"

os.makedirs(save_path, exist_ok=True)

gif_path = os.path.join(save_path, f"{experiment_type}_gif")

generate_mp4(gif_path, images, kwargs)

show_mp4(gif_path, width=opts.output_size)

　左が実写、右が求めた潜在変数から生成した画像です。ガッキーは経験上反転が難しいキャラクターなのですが、これを見ると中々精度が高いです。

　では、また。

（オリジナルgithub : https://github.com/yuval-alaluf/hyperstyle）

2022.3.21 年齢シミュレーションcolab追加

年齢による顔の変化をシミュレーションするcolabを追加しました。https://github.com/cedro3/hyperstyle/blob/main/hyperstyle_edit.ipynb

（twitter投稿）

StyleGANの画像反転にジェネレータの微調整で精度よく行うHyperstyleを使うと、油絵で描かれた顔の年齢変化の推定も可能になります。

これは、マリー・アントワネットの年齢による顔の変化の推定です。

ブログ：https://t.co/IiADQfWjhh pic.twitter.com/iOd5suYZMI
— cedro (@jun40vn) April 5, 2022

Hyperstyleで、高精度で編集しやすい反転を高速に行う

1.はじめに

2.Hyperstyleとは？

3.コード

2022.3.21 年齢シミュレーションcolab追加

コメントを残すコメントをキャンセル

ABOUTこの記事をかいた人

NEW POSTこのライターの最新記事

Animate Anyoneで、１枚の画像から動画を生成する

SVDで静止画から動画を生成する

DiffMorpherを使って、拡散モデルでモーフィングを行う

Domo AIで、実写動画をアニメ化する

最近の投稿

最近のコメント

アーカイブ

カテゴリー

メタ情報

1.はじめに

2.Hyperstyleとは？

3.コード

2022.3.21 年齢シミュレーションcolab追加

コメントを残す コメントをキャンセル

RECOMMENDこちらの記事も人気です。

PyTorch 画像から文章をサクッと生成してみる

GANs N' Roses で、１つの顔画像から様々なアニメ顔を生成する

SONY Neural Network Console に新しいデータセ…

SONY Neural Network Consoleは、みんなのディー…

学習済みWord2Vec モデルをサクッと使ってみる

ECONで、１枚の写真から３Dモデルを推定する

Keras MLPの文章カテゴリー分類を日本語のデータセットでやってみる

MediaPipeで、顔・手・ポーズのリアルタイム検出を行う

ABOUTこの記事をかいた人

NEW POSTこのライターの最新記事

Animate Anyoneで、１枚の画像から動画を生成する

SVDで静止画から動画を生成する

DiffMorpherを使って、拡散モデルでモーフィングを行う

Domo AIで、実写動画をアニメ化する

最近の投稿

最近のコメント

アーカイブ

カテゴリー

メタ情報

コメントを残すコメントをキャンセル