Edited at

Chainer とAnnoyを使った 類似画像検索 【入門】


はじめに

今回は以前の書いた記事(類似画像検索のための、Pythonを使った近似最近傍探索【入門】)の続きで、Deep Learningの特徴量抽出と近似最近傍探索を使って類似画像検索を実装します。

類似画像検索とは特徴ベクトル $\boldsymbol{x}$で表現された画像に近い画像をデータベースの中から探すことです。データベースのすべての画像と類似度を比較していると膨大な計算量になり時間がかかるので、効率よく似ている画像を探しだすことが求められます。

類似画像検索はSNSのPinterestのズームイン検索という機能に使われています。以下の例でいうと、左の画像で囲ったキーボードの画像を使って、類似画像を検索し表示しています。リモコンのようなものを検索されていますが、色や形状などが似ている画像を表示することができます。

スクリーンショット 2018-08-27 7.21.30.png

Deep Learningを使って類似画像検索を実装する方法は基本的には以下の2つの方法があります。


  • 画像特徴量の類似度計算

    画像からCNNなどで特徴量を抽出し、コサイン類似度などの関数によって類似度計算を行うことで画像の類似度を求める方法


  • Deep Learningを使った類似度学習

    似ている画像であれば、類似度が大きなるように損失を設定して学習を行う方法


今回は一つ目の「画像特徴量の類似度計算」を実装します。

構成としては画像から特徴量抽出を行う部分をDeep LearningフレームワークのChainer、特徴量の類似度計算 を行う部分を近似最近傍探索(ANN)ライブラリのAnnoyを使って実装します。


環境

Google ColaboratoryのGPU環境を使います。無料なのでGoogleアカウントがあれば誰でもGPU環境を手に入れることができます。


実装

以下はすべてGoogle Colaboratory上での実装になります。


インストール

ChainerをGPU環境で使えるようにインストールを行います。 ChainerCVは画像処理に特化したChainerのライブラリで、便利なので一緒にインストールしておきます。 

!apt -y install libcusparse8.0 libnvrtc8.0 libnvtoolsext1

!ln -snf /usr/lib/x86_64-linux-gnu/libnvrtc-builtins.so.8.0 /usr/lib/x86_64-linux-gnu/libnvrtc-builtins.so
!pip install 'cupy-cuda80==4.0.0b4' 'chainer==4.0.0b4'
!pip install chainercv

import chainer
print('GPU availability:', chainer.cuda.available)
print('cuDNN availablility:', chainer.cuda.cudnn_enabled) 

# GPU availability: True
# cuDNN availablility: True


画像分類モデルと学習の定義

今回はcifar-10という小規模なデータセットを使います。特徴量抽出を行うモデル得るために通常の画像分類と同様方法で学習を行います。

%matplotlib inline

import matplotlib.pyplot as plt
import numpy as np
from chainercv.links import VGG16
import random
from chainer import training
from chainer import iterators
from chainer import optimizers
from chainer.training import extensions
from chainer.datasets import split_dataset_random
from chainer import iterators
import chainer.links as L
import chainer.functions as F
from chainer.links import Linear

class ConvBlock(chainer.Chain):

def __init__(self, n_ch, pool_drop=False):
w = chainer.initializers.HeNormal()
super(ConvBlock, self).__init__()
with self.init_scope():
self.conv = L.Convolution2D(None, n_ch, 3, 1, 1, nobias=True, initialW=w)
self.bn = L.BatchNormalization(n_ch)
self.pool_drop = pool_drop

def __call__(self, x):
h = F.relu(self.bn(self.conv(x)))
if self.pool_drop:
h = F.max_pooling_2d(h, 2, 2)
h = F.dropout(h, ratio=0.25)
return h

class LinearBlock(chainer.Chain):

def __init__(self, drop=False):
w = chainer.initializers.HeNormal()
super(LinearBlock, self).__init__()
with self.init_scope():
self.fc = L.Linear(None, 1024, initialW=w)
self.drop = drop

def __call__(self, x):
h = F.relu(self.fc(x))
if self.drop:
h = F.dropout(h)
return h

class DeepCNN(chainer.ChainList):

def __init__(self, n_output):
super(DeepCNN, self).__init__(
ConvBlock(64),
ConvBlock(64, True),
ConvBlock(128),
ConvBlock(128, True),
ConvBlock(256),
ConvBlock(256),
ConvBlock(256),
ConvBlock(256, True),
LinearBlock(),
LinearBlock(),
L.Linear(None, n_output)
)

def __call__(self, x):
for f in self:
x = f(x)
return x

# あとで使います
def get_hidden(self, layer=10, x=None ):
for i, f in enumerate(self):
x = f(x)
if i == layer:
return x

def train(batchsize=128, gpu_id=0, max_epoch=20, base_lr=0.01):
# 1. Dataset
train_val, test = chainer.datasets.get_cifar10()
train_size = int(len(train_val) * 0.9)
train, valid = split_dataset_random(train_val, train_size, seed=0)

# 2. Iterator
train_iter = iterators.MultiprocessIterator(train, batchsize)
valid_iter = iterators.MultiprocessIterator(valid, batchsize, False, False)

# 3. Model
model = DeepCNN(10)
net = L.Classifier(model)

# 4. Optimizer
optimizer = optimizers.MomentumSGD(lr=0.01).setup(net)
optimizer.add_hook(chainer.optimizer.WeightDecay(0.0005))

# 5. Updater
updater = training.StandardUpdater(train_iter, optimizer, device=gpu_id)

# 6. Trainer
trainer = training.Trainer(updater, (max_epoch, 'epoch'), out='result-cifar-10')

# 7. Trainer extensions
trainer.extend(extensions.LogReport())
trainer.extend(extensions.observe_lr())
trainer.extend(extensions.Evaluator(valid_iter, net, device=gpu_id), name='val')
trainer.extend(extensions.PrintReport(['epoch', 'main/loss', 'main/accuracy', 'val/main/loss', 'val/main/accuracy', 'elapsed_time', 'lr']))
trainer.extend(extensions.PlotReport(['main/loss', 'val/main/loss'], x_key='epoch', file_name='loss.png'))
trainer.extend(extensions.PlotReport(['main/accuracy', 'val/main/accuracy'], x_key='epoch', file_name='accuracy.png'))
trainer.extend(extensions.snapshot_object(model
, "model_epoch_{.updater.epoch}"))
trainer.extend(extensions.dump_graph("main/loss"))

trainer.run()
del trainer

# 8. Evaluation
test_iter = iterators.MultiprocessIterator(test, batchsize, False, False)
test_evaluator = extensions.Evaluator(test_iter, net, device=gpu_id)
results = test_evaluator()
print('Test accuracy:', results['main/accuracy'])
return model

少し長くなりましたが、ネットワークを定義し、ChainerのTrainerを使った学習用の関数を定義しました。


学習開始

学習を開始します。

model = train()

epoch main/loss main/accuracy val/main/loss val/main/accuracy elapsed_time lr
1 1.54913 0.433105 1.50426 0.483594 56.1975 0.01
2 1.13509 0.589067 1.40465 0.535937 113.01 0.01
3 0.930945 0.664797 1.03801 0.6375 169.588 0.01
4 0.791013 0.717174 0.891207 0.70332 226.381 0.01
5 0.687618 0.755698 0.751768 0.743555 282.989 0.01
6 0.614297 0.782071 0.707172 0.756445 339.734 0.01
7 0.551641 0.804732 0.67634 0.766406 396.23 0.01
8 0.497919 0.823619 0.598127 0.804492 452.982 0.01
9 0.452935 0.837846 0.576993 0.806641 509.731 0.01
10 0.410008 0.85499 0.58311 0.808398 566.321 0.01
11 0.375396 0.866189 0.579353 0.81543 622.945 0.01
12 0.338473 0.88052 0.608432 0.805469 679.473 0.01
13 0.308283 0.887917 0.581245 0.816797 736.22 0.01
14 0.280307 0.900396 0.61594 0.814844 792.716 0.01
15 0.259444 0.908048 0.571848 0.826758 849.41 0.01
16 0.234903 0.915932 0.625492 0.813086 905.911 0.01
17 0.20928 0.92425 0.581202 0.832422 962.66 0.01
18 0.189567 0.932373 0.559933 0.841211 1019.34 0.01
19 0.17559 0.936254 0.543546 0.847461 1075.94 0.01
20 0.15712 0.942938 0.570619 0.838672 1132.59 0.01
Test accuracy: 0.83425635

学習結果を確認してみます。

from IPython.display import Image

Image(filename='result-cifar-10/accuracy.png')

accuracy.png

Image(filename='result-cifar-10/loss.png')

loss.png

改善できそうな部分はありますが、今回の目的とはずれてしまうので画像分類の精度はおいといて先に進みます。


Annoyのインストール

近似最近傍探索ライブラリのAnnoyをインストールします。

!pip install annoy


特徴ベクトルをデータベースへの追加

Annoyで類似度計算を行うために、必要となる特徴ベクトルをAnnoyモデルに追加する必要があります。

今回は学習済みモデルを使って、画像分類の推論を行う際の10層目の出力を特徴ベクトルとして用います。

CuPyとNumPyの変換を忘れるとランタイムエラーで

Google Colabratoryが落ちるので注意してください。

from annoy import AnnoyIndex

train, test = chainer.datasets.get_cifar10()

# 画像分類モデルの10層目の出力要素数
dim = 1024

# Annoyモデルの宣言
annoy_model = AnnoyIndex(dim)

with chainer.using_config('train', False), chainer.using_config('enable_backprop', False):
for i in range(len(train)):
img, _ = train[i]
# numpy -> cupy
x = model.xp.asarray(img[None, ...])

# 学習データの推論結果のうち画像分類モデルの10層目の出力を得る
x = get_hidden(model ,9 ,x=x).data
x = x.reshape(-1)

#cupy -> numpy
x = chainer.cuda.to_cpu(x)

# Annoyモデルにインデックスと特徴ベクトルを追加
annoy_model.add_item(i, x)

# Annoyモデルのビルド(以後データの追加は行えない)
annoy_model.build(1000)
annoy_model.save("cifar-10-1000tree.ann")


類似度計算と検索結果の表示

上記と同じように学習済みモデルを使って得た画像分類の推論を行う際の10層目の出力を特徴ベクトルとして、Annoyモデルにわたすと類似度計算を行い、結果を返してくれます。


rows_count = 10
columns_count = 6
images_count = rows_count * columns_count
axes = []
fig = plt.figure(figsize=(15, 15))
with chainer.using_config('train', False), chainer.using_config('enable_backprop', False):
for i in range(10):
index =random.randint(0, len(test))
x, t = test[index]
answer_ax = fig.add_subplot(rows_count, columns_count, i*columns_count + 1)
answer_ax.imshow(x.transpose(1, 2, 0))
plt.axis("off")
answer_ax.set_title("Correct Answer")

x = model.xp.asarray(x[None, ...])
x = get_hidden(model ,9 ,x= x).data
x = x.reshape(-1)
x = chainer.cuda.to_cpu(x)

# 特徴ベクトルxをわたすと、類似度計算を行い、類似度の大きいものを返す
predict_indexes = annoy_model.get_nns_by_vector(x, 5, search_k=-1)
for j, predict_i in enumerate(predict_indexes):
predict_x, predict_t = train[predict_i]
ax = fig.add_subplot(rows_count, columns_count, i*columns_count + 2+ j)
ax.imshow(predict_x.transpose(1, 2, 0))
plt.axis("off")

fig.subplots_adjust(wspace=0.2, hspace=0.2)
plt.tight_layout()
plt.savefig("cifar10-1000tree-sample5.png")
plt.show()


結果  

左の列がクエリ画像で各行の右の5個が検索結果です。 

結構似ている画像が検索できていると思います。

cifar10-1000tree-sample3.png

cifar10-1000tree-sample4.png


まとめ


  • Chainerと近似最近傍探索ライブラリAnnoyを使うと類似画像検索できる

  • 大規模なデータセット & 深いネットワークのCNN出力を利用したversionも追記予定

  • 次はDeep Learningの類似度学習を使った類似画像検索を実装してみる


参考