torch.rbとannoy.rbを使ってRubyで類似画像検索

6 min read読了の目安(約5700字

はじめに

torch.rbはlibtorch(PyTorchのC++版)をbindingしたもので、Rubyで深層学習を行うことができます。別記事では、学習済みのVGG-16 Networkを使って画像認識を行う方法を紹介しました。今回は、VGG-16から画像の特徴ベクトルを抽出して、それをベクトル検索ライブラリであるannoy.rbによりインデックス化して、類似画像検索を実現する方法を紹介します。

https://zenn.dev/yoshoku/articles/91ac1d0671ed04

インストール

この記事での作業は、macOS Big Sur 11.2で行いました。必要なライブラリはhomebrewでインストールしていきます。まずは、前節で紹介している記事と同様に、torch.rbとtorchvision、画像処理を行うためのmagroをインストールします。

$ brew install libtorch automake libpng jpeg
$ gem install torch-rb torchvision magro

そして、類似検索のためのインデックスを作成するannoy.rbをインストールします。annoy.rbは外部ライブラリを必要としません。

$ gem install annoy-rb

したごしらえ

torch.rbはlibtorchをbindingしたものですが、libtorchにちょっとしたバグがあり(2021/2/11現在)、PyTorchの学習済みモデルをそのまま読み込むことができません。「はじめに」で紹介している記事にある方法で、一度読み込めるかたちに修正したものをご準備ください。

類似画像検索

画像からの特徴ベクトルの抽出にtorch.rbを、特徴ベクトルの検索インデックスにannoy.rbを使用して、類似画像検索を行うスクリプトを、パートにわけて解説します。わりと長いスクリプトになりましたので、全体をGistに置きました。

https://gist.github.com/yoshoku/846a980aafd8c7dcb3caf8eaf0c140c6

データセットの準備

今回は、類似検索ができることを確認できればよいので、小ぶりなデータセットを使うことにしました。ワシントン大学のObject and Concept Recognition for Content-Based Image Retrievalで提供されているデータセットを使います。

http://imagedatabase.cs.washington.edu/

Ground Truth Databaseというリンクからたどれるデータセットをダウンロードします。rarファイルを解凍するとできる、icpr2004.imgset/groundtruthディレクトリ以下にあるjpg画像を用います。

$ brew install unar
$ wget http://imagedatabase.cs.washington.edu/groundtruth/icpr2004.imgset.rar
$ unar icpr2004.imgset.rar

画像の読み込みと前処理

画像の読み込みにはmagroを用います。magroは、画像をNumo::NArray形式で扱います。Numo::NArrayは、Pythonでいうところのnumpyで、torch.rbでもtensorとの相互変換がサポートされています。画像をNumo::NArrayで扱ったほうが、取り回しが良いのでmagroを用います。torchvisionが提供する学習済みモデルの前処理は、以下に詳しく書かれています。

https://pytorch.org/vision/0.8/models.html

これに従って、前処理をメソッド化します。

retrieval.rb
# @param src [Numo::NArray] (shape: [height, width, channel]) 入力画像
def preprocessing(src)
  # 画像の中心を正方形に切り出す.
  height, width, = src.shape
  img_size = [height, width].min
  y_offset = (height - img_size) / 2
  x_offset = (width - img_size) / 2
  img = src[y_offset...(y_offset + img_size), x_offset...(x_offset + img_size), true].dup
  # 画像を224x224の大きさにする.
  img = Magro::Transform.resize(img, height: 224, width: 224)
  # 画素値を[0, 1]の範囲に正規化する.
  img = Numo::SFloat.cast(img) / 255.0
  img -= Numo::SFloat[0.485, 0.456, 0.406]
  img /= Numo::SFloat[0.229, 0.224, 0.225]
  # [チャンネル, 高さ, 幅]の順に入れ替える.
  img.transpose(2, 0, 1).dup
end

ダウンロードしたデータセットの画像を読み込み、前処理を施して、torch.rbのtensorに変換します。

retrieval.rb
require 'magro'

filelist = Dir.glob('icpr2004.imgset/groundtruth/*/*.jpg')
images = Numo::SFloat.cast filelist.map { |filename| preprocessing(Magro::IO.imread(filename)) }
images_tensor = Torch.from_numo(images)

VGG-16を使った特徴抽出

「したごしらえ」の節で得られた、学習済みモデルを読み込んで特徴抽出を行います。ネットワークの定義はtorchvisionのものを使います。

retrieval.rb
require 'torch'
require 'torchvision'

vgg = TorchVision::Models::VGG16.new
vgg.load_state_dict(Torch.load('vgg16_.pth'))

VGG-16は画像認識のためのニューラルネットワークであるため、最終層は分類器の働きをします。以下はそれをpメソッドで表示したものですが、6番のLinearが分類器の働きをしています。これを除いてやれば、手前の4096ユニットの出力が得られます。

> p vgg.classifier
Sequential(
  (0): Linear(in_features: 25088, out_features: 4096, bias: true)
  (1): ReLU(inplace: true)
  (2): Dropout(p: 0.5, inplace: false)
  (3): Linear(in_features: 4096, out_features: 4096, bias: true)
  (4): ReLU(inplace: true)
  (5): Dropout(p: 0.5, inplace: false)
  (6): Linear(in_features: 4096, out_features: 1000, bias: true)
)

いまのところ、ネットワーク構成を変更するメソッドはtorch.rbにはない様子です。このネットワーク構成は、modulesというインスタンス変数にあります。これを、instance_variable_getで取得して変更します。

retrieval.rb
vgg.classifier.instance_variable_get(:@modules).delete("6")

前節で得られた、画像データセットのtensorをVGG-16に与えて、各画像の特徴ベクトルを抽出します。annoy.rbでは、ベクトルをRuby Arrayで扱うので、to_aメソッドでtensorをRuby Arrayに変換しています。

retrieval.rb
vgg.eval
features = vgg.forward(images_tensor).to_a

annory.rbによる類似画像検索

annoy.rbは近似最近傍探索ライブラリのAnnoyのbindingライブラリです。AnnoyはPythonライブラリがよく使われていますが、本体はC++で書かれた割とコンパクトなものです。annoy.rbはこのC++のコードをラップしたものなので、annoy.rbの他に外部ライブラリをインストールする必要はありません。

さっそく、前節で抽出した画像の特徴ベクトルを登録していきましょう。annoy.rbでは、インデックスを生成する際に、登録するベクトルの次元数と類似度を評価する距離を指定します。最終層は4096ユニットですので、特徴ベクトルは4096次元となります。距離には、コサイン類似度による距離(L2正規化されたベクトル間のユークリッド距離)であるangularを選択しました。

retrieval.rb
require 'annoy'

# ベクトルの次元数と距離を指定する.
index = Annoy::AnnoyIndex.new(n_features: 4096, metric: 'angular')

# Ruby Arrayによるベクトルを添え字とともに登録する.
features.each_with_index { |vec, idx| index.add_item(idx, vec) }

Annoyは木構造による検索インデックスです。2-meansクラスタリングに似た方法でデータを二分割することで木構造を作りますが、このときにランダムにデータを選択します。木の作成にランダム性があるので、複数の木を組み合わせて、検索結果を安定させます。検索インデックスの作成のパラメータとして、木の数があります。

retrieval.rb
# 検索インデックスを作成する. 木の数は10本とした.
index.build(10)

これで検索インデックスができたので、実際に検索してみましょう。登録したデータの添字を指定して検索する get_nns_by_item メソッドがありますが、今回は、クエリが与えられたことを想定して、ベクトルで検索する方法を用いました。

retrieval.rb
# 添字55の特徴ベクトルをクエリに見立てて、5-近傍を検索する.
query = features[55]
retrieved_ids = index.get_nns_by_vector(query, 5)

# 検索結果を出力する.
pp retrieved_ids.map { |i| filelist[i] }

これを実行すると、次のようになりました。

$ ruby retrieval.rb
["icpr2004.imgset/groundtruth/cherries/cherries09.jpg",
 "icpr2004.imgset/groundtruth/cherries/cherries10.jpg",
 "icpr2004.imgset/groundtruth/cherries/cherries30.jpg",
 "icpr2004.imgset/groundtruth/cherries/cherries03.jpg",
 "icpr2004.imgset/groundtruth/cherries/cherries04.jpg"]

実際に見てみると...検索できてますね 👍

おわりに

画像類似検索システムを作る場合、ニューラルネットワークで画像の特徴ベクトルを抽出し、Annoyの様な近似最近傍探索ライブラリでインデックスを作成し、検索APIを用意するということが行われていると思います。この記事では紹介しませんでしたが、annoy.rbは save メソッドで、インデックスを保存することもできます。あとは、RailsでAPIを用意することで、画像類似検索システムがRubyで作れます 💪

https://github.com/yoshoku/annoy.rb

この記事に贈られたバッジ