Python
画像処理
機械学習
DeepLearning
Chainer

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の類似度学習を使った類似画像検索を実装してみる

参考