illustration2vecをkerasモデルに変換するコードを作ってみた

概要

https://github.com/rezoo/illustration2vec
illustration2vecはイラストの特徴やタグを検出できるモデルです。モデルの構造はほぼVGGモデルです。オリジナルのVGGとは変更されている箇所もあるので、詳細は上記リンクに解説論文があるのでそちらを。caffeとchainerのモデルです。面白いモデルですがchainerが開発終了するようなので、ぜひ再利用したいと思いkerasへの変換コードを書きました。torchは使ったことがないのでスルーで。

実行はgoogle colaboratoryで行ったので以下にリンクを。
https://colab.research.google.com/drive/1UZN7pn4UzU5s501dwSIA2IHGmjnAmouY

リンク先が開けない場合は以下をcolabにコピペすれば動くと思います。

全コード

!git clone https://github.com/rezoo/illustration2vec.git
!sh illustration2vec/get_models.sh
!pip install -r /content/illustration2vec/requirements.txt
!mv /content/illustration2vec/i2v /content/
import i2v
illust2vec_tag = i2v.make_i2v_with_chainer('/content/illustration2vec/illust2vec_tag_ver200.caffemodel', '/content/illustration2vec/tag_list.json')
%tensorflow_version 2.x
import tensorflow as tf
import numpy as np

# tag estimater model
model_tag = tf.keras.Sequential(name='illustration2vec_tag')
model_tag.add(tf.keras.layers.Input(shape=(224, 224, 3)))
pool_idx = [0, 1, 3, 5, 7]
for i, chainer_layer in enumerate(illust2vec_tag.net.children()):
    kernel, bias = tuple(chainer_layer.params())
    k_kernel = np.transpose(kernel.data, axes=[3, 2, 1, 0])
    bias = bias.data
    if i == 0:
        k_kernel = k_kernel[:,:,::-1,:]
    channel = bias.shape[0]
    keras_layer = tf.keras.layers.Conv2D(channel, 3, padding='SAME', activation='relu', kernel_initializer=tf.keras.initializers.constant(k_kernel), bias_initializer=tf.keras.initializers.constant(bias), name='Conv_%d'%i)
    model_tag.add(keras_layer)
    if i in pool_idx:
        model_tag.add(tf.keras.layers.MaxPooling2D())
    del kernel, bias
model_tag.add(tf.keras.layers.AveragePooling2D(pool_size=(7, 7)))
model_tag.add(tf.keras.layers.Lambda(lambda x : tf.nn.sigmoid(tf.squeeze(x, axis=[1, 2])), name='sigmoid'))
model_tag.save('illust2vec_tag_ver200.h5')
del model_tag, illust2vec_tag

# feature vector model
illust2vec = i2v.make_i2v_with_chainer('/content/illustration2vec/illust2vec_ver200.caffemodel')
model = tf.keras.Sequential(name='illustration2vec')
model.add(tf.keras.layers.Input(shape=(224, 224, 3)))
pool_idx = [0, 1, 3, 5, 7]
for i, chainer_layer in enumerate(illust2vec.net.children()):
    if i == 12:
        break
    kernel, bias = tuple(chainer_layer.params())
    if len(kernel.data.shape) == 4:
        k_kernel = np.transpose(kernel.data, axes=[3, 2, 1, 0])
        bias = bias.data
        if i == 0:
            k_kernel = k_kernel[:,:,::-1,:]
        channel = bias.shape[0]
        keras_layer = tf.keras.layers.Conv2D(channel, 3, padding='SAME', activation='relu', kernel_initializer=tf.keras.initializers.constant(k_kernel), bias_initializer=tf.keras.initializers.constant(bias), name='Conv_%d'%i)
        model.add(keras_layer)
        if i in pool_idx:
            model.add(tf.keras.layers.MaxPooling2D())
        elif i == 10:
            model.add(tf.keras.layers.Flatten())
    elif len(kernel.data.shape) == 2:
        model.add(tf.keras.layers.Dense(4096, kernel_initializer=tf.keras.initializers.constant(kernel.data), bias_initializer=tf.keras.initializers.constant(bias.data), name='encode1'))
    del kernel, bias
model.save('illust2vec_ver200.h5')
del model, illust2vec
def resize(imgs):
    mean = tf.constant(np.array([181.13838569, 167.47864617, 164.76139251]).reshape((1, 1, 3)), dtype=tf.float32)
    resized = []
    for img in imgs:
        img = tf.cast(img, tf.float32)
        im_max = tf.reduce_max(img, keepdims=True)
        im_min = tf.reduce_min(img, keepdims=True)
        im_std = (img - im_min) / (im_max - im_min + 1e-10)
        resized_std = tf.image.resize(im_std, (224, 224))
        resized_im = resized_std*(im_max - im_min) + im_min
        resized_im = resized_im - mean
        resized.append(tf.expand_dims(resized_im, 0))
    return tf.concat(resized, 0)

解説

やっていることはシンプルでkerasで同じ構造のモデルを定義してchainerモデルのweightで初期化しているだけです。
変更点は以下の3つ。
・weightのtranspose
　このchainerモデルのconvolutionのカーネルは(out_channel, in_channel, k_size, k_size)になっていますが、kerasモデル用に(k_size, k_size, in_channel, out_channel)に転置しています。
・入力画像をBGRからRGBに変更
　元のモデルは入力画像はBGRになっていますが、RGBでの入力に変更しました。そのため第1層の畳み込みのカーネルをin_channelの軸で逆順にしています。
・入力画像をchannel_lastに変更
　chainerのconvolutionは入力は(N, C, H, W)になっていますが、kerasのデフォルトの(N, H, W, C)になっています。

変換後のモデルの入力は(batch, 224, 224, 3)のサイズです。上記のcolabリンク先か全コードの中のresize関数に画像のnumpy.arrayのリストを渡せばリサイズ+正規化を行ってくれます。illust2vec_tag_ver200.h5とillust2vec_ver200.h5の2つのモデルがありますが、1つ目の出力はhttps://github.com/rezoo/illustration2vecにあるtag_list.jsonの1539個のタグそれぞれの確率になります。2つ目のモデルの出力は画像のfeature vectorになっています。

何か不備があれば指摘していただければと。