技術ブログ

Developers's blog

TensorFlowでVGG19を使ったインスタ映え画像の生成

2019.12.11 岡本 和斗
TensorFlow ニューラルネットワーク 機械学習
TensorFlowでVGG19を使ったインスタ映え画像の生成

概要

DeepArtのようなアーティスティックな画像を作れるサービスをご存知でしょうか?
こういったサービスではディープラーニングが使われており、コンテンツ画像とスタイル画像を元に次のような画像の画風変換を行うことができます。この記事では画風変換の基礎となるGatysらの論文「Image Style Transfer Using Convolutional Neural Networks」[1]の解説と実装を行っていきます。

examples_of_style_transfer

引用元: Gatys et al. (2016)[1]


手法

モデルにはCNN(Convolutional Neural Network) が用いられており、VGG[2]という物体認識のための事前学習済みのモデルをベースとしています。


こちらの図はCNNが各層においてどのようにコンテンツ画像とスタイル画像を表現するか示しています。

image_representaion_on_each_layers_in_CNN

引用元: Gatys et al. (2016)[1]


Content Reconstructions (下段) のa, b, c を見ると入力画像がほぼ完璧に復元されていることがわかります。一方で d, e を見ると詳細な情報は失われているものの、物体認識をする際に重要な情報が抽出されていることがわかります。これは画像からコンテンツが抽出され、画風を表す情報が落とされていると考えられます。よって画風変換のモデルでは画像からコンテンツを捉えるためにCNNの深い層を利用します。
次にStyle Reconstructions (上段) を見てみましょう。a からe はそれまでの各層の特徴マップの相関をもとに復元されています。例えば c は第1層から第3層までの各層における特徴マップの相関から生成されています。こうすることで画像内のコンテンツの配置などによらない画風を抽出することができます[3]。


画風変換のモデルではコンテンツ画像とスタイル画像を入力として受け取り、上記のことを利用してスタイル画像の画風をコンテンツ画像に反映した画像を新たに生成します。


では具体的にモデルの中身を見ていきましょう。


モデルのアーキテクチャ

algorithm_of_style_transfer_model

引用元: Gatys et al. (2016)[1]


まずコンテンツ画像 (\(\vec{p}\)) とスタイル画像 (\(\vec{a}\)) から特徴マップが抽出されます。次にこれらの特徴マップと生成画像 (\(\vec{x}\))の特徴マップとの損失が計算されます。この計算で求められる損失をそれぞれコンテンツ損失、スタイル損失と呼ぶことにします。そしてコンテンツ損失とスタイル損失の合計が最終的な損失となり、これを最小化していきます。
ここで注意しなければならないことが1つあります。通常、ディープラーニングでは重みが最適化の対象になりますが、今回は重みは固定して生成画像のピクセルを最適化します。


コンテンツ損失

コンテンツ損失はコンテンツ画像と生成画像のVGGのある1層から出力された特徴マップの平均二乗誤差によって計算されます。

$$ \mathcal{L}_{content} (\vec{p}, \vec{x}, l) = \frac{1}{2} \sum_{i, j}^{} (F_{ij}^{l} - P_{ij}^{l})^2 $$ ここで\(F_{ij}^{l}\) は生成画像の\(l\)層における\(i\) 番目のフィルターの位置\(j\) でのアクティベーションを表しています。\(P_{ij}^{l}\)についても同様ですが、こちらはコンテンツ画像についてのアクティベーションを表しています。


スタイル損失

スタイル画像から画風を捉えるためにまず各層における特徴マップの相関を計算します。

$$ G_{ij}^{l} = \sum_{k}^{} F_{ik}^{l}F_{jk}^{l} $$ これはグラム行列と呼ばれ、\(l\) 層におけるフィルター間の特徴マップの相関をとっています。このグラム行列を用いることで画風を表現することができます[3]。 そして生成画像のグラム行列とスタイル画像のグラム行列の平均二乗誤差を求めます。 $$ E_{l} = \frac{1}{4N_{l}^{2}M_{l}^{2}} \sum_{i,j}^{} (G_{ij}^{l} - A_{ij}^{l})^2 $$ ここで\(N_{l}\)は特徴マップの数、\(M_{l}\)は特徴マップのサイズを表します。また\(G_{ij}^{l}\)、\(A_{ij}^{l}\) はそれぞれ\(l\) 層における生成画像のグラム行列とスタイル画像のグラム行列を表します。  最後に各層の損失の線形和をとってスタイル損失とします。 $$ \mathcal{L}_{style} (\vec{a}, \vec{x}) = \sum_{l=0}^{L} w_{l}E_{l} $$ このとき\(w_{l}\) は\(l\) 層の損失の重みを表します。


最適化

コンテンツ損失とスタイル損失から合計の損失を求めます。合計損失は\(\alpha\) と\(\beta\) をそれぞれコンテンツ損失とスタイル損失の重みとして $$ \mathcal{L}_{total}(\vec{p}, \vec{a}, \vec{x}) = \alpha\mathcal{L}_{content}(\vec{p}, \vec{x}) + \beta\mathcal{L}_{style}(\vec{a}, \vec{x}) $$ この損失を最小化する形で生成画像の最適化を行っていきます。ですので最適化には\(\frac{\partial L_{totla}}{\partial \vec{x}}\) を用いることになります。 論文中ではL-BFGSで最も良い結果になったと記述されていましたが、今回の実装ではAdam を使って最適化を行いました。 この最適化によって、コンテンツ画像をスタイル画像に合わせて画風変換した新たな画像が生成されます。


実装

実装にはTensorFlowを用いました。また実装に際してTensorFlowのチュートリアル[4]を参考にしました。


まずは必要なライブラリをインストールします。

import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
%matplotlib inline
mpl.rcParams['figure.figsize'] = (12,12)
mpl.rcParams['axes.grid'] = False
import time
import IPython.display as display
import PIL.Image

import tensorflow as tf
from keras.preprocessing import image


画像を読み込んで配列に変換する関数、テンソルを画像に変換する関数を定義します。

# 画像を読み込み配列に変換し正規化する
def load_image(input_path, size):
    image = tf.keras.preprocessing.image.load_img(input_path, target_size=size)
    image = tf.keras.preprocessing.image.img_to_array(image)
    image = np.expand_dims(image, axis=0)
    image /= 255
    return image


# テンソルを画像に戻す
def tensor_to_image(tensor):
    tensor = tensor.numpy()
    tensor *= 255
    tensor = np.array(tensor, dtype=np.uint8)
    if np.ndim(tensor)>3:
        assert tensor.shape[0] == 1
        tensor = tensor[0]
    return PIL.Image.fromarray(tensor)


このクラスは配列に変換された画像を受け取り、VGG19の各層からの出力を返します。この時点でスタイルの表現に使われる特徴マップはグラム行列に変換されます。

class StyleContentModel():
    def __init__(self):
        # VGG19のどの層の出力を使うか指定する
        self.style_layers = ['block1_conv1', 'block2_conv1', 'block3_conv1', 'block4_conv1', 'block5_conv1']
        self.content_layers = ['block5_conv2']

        self.num_style_layers = len(self.style_layers)        
        self.vgg = self.get_vgg_model()
        self.vgg.trainable = False

    def __call__(self, inputs):
        preprocessed_input = tf.keras.applications.vgg19.preprocess_input(inputs * 255)
        vgg_outputs = self.vgg(preprocessed_input)

        style_outputs, content_outputs = (vgg_outputs[:self.num_style_layers], vgg_outputs[self.num_style_layers:])
        style_outputs = [self.gram_matrix(style_output) for style_output in style_outputs]

        style_dict = {style_name:value for style_name, value in zip(self.style_layers, style_outputs)}
        content_dict = {content_name:value  for content_name, value in zip(self.content_layers, content_outputs)}

        return {'style':style_dict, 'content':content_dict}

    # Keras API を利用してVGG19を取得する  
    def get_vgg_model(self):
        vgg = tf.keras.applications.vgg19.VGG19(include_top=False, weights='imagenet')
        vgg.trainable = False

        outputs = [vgg.get_layer(name).output for name in (self.style_layers + self.content_layers)]
        model = tf.keras.Model(vgg.input, outputs)

        return model

    # グラム行列を計算する
    def gram_matrix(self, input_tensor):
        result = tf.linalg.einsum('bijc,bijd->bcd', input_tensor, input_tensor)
        input_shape = tf.shape(input_tensor)
        num_locations = tf.cast(input_shape[1] * input_shape[2], tf.float32)
        return result / num_locations


そして合計損失を計算する関数、計算した損失から勾配を計算する関数を定義します。

# 合計損失を計算する
def compute_loss(model, base_image, style_targets, content_targets, style_weight, content_weight):
    model_outputs = model(base_image)
    style_outputs = model_outputs['style']
    content_outputs = model_outputs['content']

    style_loss = tf.add_n([tf.reduce_mean((style_outputs[name]-style_targets[name])**2) for name in style_outputs.keys()])
    style_loss *= style_weight / len(style_outputs)

    content_loss = tf.add_n([tf.reduce_mean((content_outputs[name]-content_targets[name])**2) for name in content_outputs.keys()])
    content_loss *= content_weight / len(content_outputs)

    loss = style_loss + content_loss
    return loss, style_loss, content_loss

# 損失を元に勾配を計算する
@tf.function()
def compute_grads(params):
    with tf.GradientTape() as tape:
        all_loss = compute_loss(**params)

    grads = tape.gradient(all_loss[0], params['base_image'])
    return grads, all_loss


最後に画像の生成を行う関数を定義していきます。
生成画像のベースとなる画像にはノイズ画像を指定しています。ベース画像にコンテンツ画像やスタイル画像を指定するとまた違った結果が得られます。

def run_style_transfer(style_path, content_path, num_iteration, style_weight, content_weight, display_interval):
    size = image.load_img(content_path).size[::-1]
    noise_image = np.random.uniform(-20, 20, (1, size[0], size[1], 3)).astype(np.float32) / 255
    content_image = load_image(content_path, size)
    style_image = load_image(style_path, size)

    model = StyleContentModel()
    style_targets = model(style_image)['style']
    content_targets = model(content_image)['content']

    # 生成画像のベースとしてノイズ画像を使う
    # ベースにはコンテンツ画像またはスタイル画像を用いることもできる
    base_image = tf.Variable(noise_image)

    opt = tf.optimizers.Adam(learning_rate=0.02, beta_1=0.99, epsilon=1e-1)

    params = {
        'model': model,
        'base_image': base_image,
        'style_targets': style_targets,
        'content_targets': content_targets,
        'style_weight': style_weight,
        'content_weight': content_weight
    }

    best_loss = float('inf')
    best_image = None

    start = time.time()
    for i in range(num_iteration):
        grads, all_loss = compute_grads(params)
        loss, style_loss, content_loss = all_loss

        opt.apply_gradients([(grads, base_image)])
        clipped_image = tf.clip_by_value(base_image, clip_value_min=0., clip_value_max=255.0)
        base_image.assign(clipped_image)

        # 損失が減らなくなったら最適化を終了する        
        if loss < best_loss:
            best_loss = loss
            best_image = base_image
        elif loss > best_loss:
            tensor_to_image(base_image).save('output_' + str(i+1) + '.jpg')
            break

        if (i + 1) % display_interval == 0:
            display.clear_output(wait=True)
            display.display(tensor_to_image(base_image))
            tensor_to_image(base_image).save('output_' + str(i+1) + '.jpg')
            print(f'Train step: {i+1}')
            print('Total loss: {:.4e}, Style loss: {:.4e}, Content loss: {:.4e}'.format(loss, style_loss, content_loss))

    print('Total time: {:.4f}s'.format(time.time() - start))
    display.clear_output(wait=True)
    display.display(tensor_to_image(base_image))

    return best_image


では実際に画風変換を行ってみましょう。styleweightとcontentweightはそれぞれスタイル損失とコンテンツ損失の重みを表します。

style_path = '../input/neural-image-transfer/StarryNight.jpg'
content_path = '../input/neural-image-transfer/FlindersStStation.jpg'
num_iteration = 5000
style_weight = 1e-2
content_weight = 1e4
display_interval = 100

best_image = run_style_transfer(style_path, content_path, num_iteration, style_weight, content_weight, display_interval)


結果

コンテンツ画像とスタイル画像はこれらの画像を使いました。


コンテンツ画像

Flinders_Street_Station_in_Melbourne

引用元: https://commons.wikimedia.org/wiki/File:Flinders_Street_Station_3.jpg


スタイル画像

The_Starry_Night_art_of_Gogh

引用元: https://commons.wikimedia.org/wiki/File:Van_Gogh_-_Starry_Night_-_Google_Art_Project.jpg


上記のコードを走らせて生成した画像がこちらになります。 スタイル損失とコンテンツ損失の重みを変えて4種類の画像を生成しました。

  • \(\alpha = 10, \beta = 10^{-2}\) generated_image_with_content_loss_weight_1e1_style_loss_weight_1e-2



  • \(\alpha = 10^{2}, \beta = 10^{-2}\) generated_image_with_content_loss_weight_1e2_style_loss_weight_1e-2



  • \(\alpha = 10^{3}, \beta = 10^{-2}\) generated_image_with_content_loss_weight_1e3_style_loss_weight_1e-2



  • \(\alpha = 10^{4}, \beta = 10^{-2}\) generated_image_with_content_loss_weight_1e4_style_loss_weight_1e-2



結果からわかるように \(\frac{\alpha}{\beta}\) が大きくなればなるほどコンテンツ画像がはっきりと生成画像に反映されていることがわかります。これは\(\frac{\alpha}{\beta}\)が大きいとコンテンツ損失に対してモデルが敏感になるからです。 逆に\(\alpha = 10, \beta = 10^{-2}\) の場合はスタイル損失に敏感になりすぎて、生成画像からコンテンツを見つけることができなくなっています。


まとめ

今回は画風変換の基礎となる論文の解説と実装を行いました。ベース画像やパラメーター、VGGのどの層の出力を使うかなどによって違った結果が得られるので色々といじってみるのも面白いかもしれません。またこの論文以降にも画風変換の研究が進められていますので、それらを試す助けになればと思います。


参考文献

[1] Image Style Transfer Using Convolutional Neural Networks
[2] Very Deep Convolutional Networks For Large-scale Image Recognition
[3] Texture Synthesis Using Convolutional Neural Networks
[4] TensorFlow Core: Neural style transfer

お問い合わせはこちらから