KerasでAV女優の類似画像検索機能を実装する

本記事はchainerによるディープラーニングでAV女優の類似画像検索サービスをつくったノウハウを公開するの続編になります。
以前の記事では学習したモデルに画像を入力しクラス分類することでどの女優に似ているかを判定していましたが、本記事では全結合層の特徴ベクトルからコサイン類似度を計算し、類似画像を検索します。
また、今回は単純な画像認識では、ChainerよりもKerasの方が使いやすいと感じたため、Kerasを用いて実装をしています。

データ拡張

KerasではImageDataGeneratorを使うことで簡単にデータ拡張することができます。 ランダムに画像を傾けたり、シフトしたりすることで全く同じデータで学習することが少なくなり、過学習しにくくなるといわれています。
また、ZCA whiteningについては以下の資料を参照してください。

データの白色化 - DEEPTONEWorks
CIFAR-10 と ZCA whitening - まんぼう日記

from keras.preprocessing.image import ImageDataGenerator

# なにかしらのデータとラベルを読み込む
data, label = load_data()

datagen = ImageDataGenerator(
　zca_whitening=True,
    rotation_range=10,
    width_shift_range=0.1,
    height_shift_range=0.1,
    horizontal_flip=True)

datagen.fit(data)

# モデルのfit_generator関数にdatagen.flowを渡し、動的にデータ拡張する
model.fit_generator(datagen.flow(data, label, batch_size=32),
    samples_per_epoch=data.shape[0],
    nb_epoch=100)

顔画像の正面化

前の記事では顔画像の検出にdlibを使いましたが、今回はさらに検出した顔画像に対して顔の特徴点を抽出し、目や口の位置が正面にくるようアフィン変換を行います。 以下のopenfaceやfacenetで実装されているので、これをほぼそのまま使うことができます。

facenet/src/align_dlib.py
openface/util/align-dlib.py

モデル構造

まだまだ学習データセットが少ない段階のため、モデルやハイパーパラメータの調整は全くしていませんが、ひとまず以下のようなモデルを使用しています。

def conv_bn_relu(x, out_ch, name):
    x = Convolution2D(out_ch, 3, 3, border_mode='same', name=name)(x)
    x = BatchNormalization(name='{}_bn'.format(name))(x)
    x = Activation('relu', name='{}_relu'.format(name))(x)
    return x

def face_model(input_shape=(3, 224, 224), nb_classes, weights_path=None):

    inputs = Input(shape=input_shape, name='input')

    x = conv_bn_relu(inputs, 64, name='block1_conv1')
    x = conv_bn_relu(x, 64, name='block1_conv2')
    x = MaxPooling2D((2, 2), strides=(2, 2))(x)

    x = conv_bn_relu(x, 128, name='block2_conv1')
    x = conv_bn_relu(x, 128, name='block2_conv2')
    x = MaxPooling2D((2, 2), strides=(2, 2))(x)

    x = conv_bn_relu(x, 256, name='block3_conv1')
    x = conv_bn_relu(x, 256, name='block3_conv2')
    x = conv_bn_relu(x, 256, name='block3_conv3')
    x = MaxPooling2D((2, 2), strides=(2, 2))(x)

    x = conv_bn_relu(x, 512, name='block4_conv1')
    x = conv_bn_relu(x, 512, name='block4_conv2')
    x = conv_bn_relu(x, 512, name='block4_conv3')
    x = MaxPooling2D((2, 2), strides=(2, 2))(x)

    x = Flatten()(x)
    x = Dense(4096, activation='relu', name='fc1')(x)
    x = Dense(nb_classes, activation='softmax', name='predictions')(x)

    model = Model(input=inputs, output=x)

類似度の計算

今回のモデルでは、最終層の一つ前の全結合層の特徴ベクトルを抽出し、各女優について抽出された特徴ベクトルとのコサイン類似度を計算することで画像の類似度を求めます。

import numpy as np
from scipy.spatial.distance import cosine
from keras.models import Model, model_from_json
from keras.preprocessing import image
from keras.preprocessing.image import img_to_array

def calculate_similarity():
    # Load model
    model_json = open('face_model.json').read()
    base_model = model_from_json(model_json)
    base_model.load_weights('face_model.h5')
    # 全結合層のと4096次元の特徴ベクトルを抽出するモデルを作成
    model = Model(input=base_model.input, output=base_model.get_layer('fc1').output)

    image_file = 'image.jpg'
    img = image.load_img(image_file, target_size=(96, 96))
    x = img_to_array(img)
    x = np.expand_dims(x, axis=0)

    features = model.predict(x)
    features = features.flatten().tolist()

　# actress_featuresは事前に計算しておいた各女優の特徴ベクトル
　# 入力画像から抽出された特徴ベクトルと各女優の特徴ベクトルのコサイン類似度を計算する
    score = 1 - cosine(features, actress_features)

最後に宣伝になりますが、本記事の類似顔画像検索を実装したサイトを作っているので、よかったら見てみてください。
Babelink - あなた好みのAVを検索
※アダルトサイトのため、閲覧には十分注意をしてください。