torch.rbとannoy.rbを使ってRubyで類似画像検索
はじめに
torch.rbはlibtorch（PyTorchのC++版）をbindingしたもので、Rubyで深層学習を行うことができます。別記事では、学習済みのVGG-16 Networkを使って画像認識を行う方法を紹介しました。今回は、VGG-16から画像の特徴ベクトルを抽出して、それをベクトル検索ライブラリであるannoy.rbによりインデックス化して、類似画像検索を実現する方法を紹介します。
インストール
この記事での作業は、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に置きました。
データセットの準備
今回は、類似検索ができることを確認できればよいので、小ぶりなデータセットを使うことにしました。ワシントン大学のObject and Concept Recognition for Content-Based Image Retrievalで提供されているデータセットを使います。
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が提供する学習済みモデルの前処理は、以下に詳しく書かれています。
これに従って、前処理をメソッド化します。
# @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に変換します。
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のものを使います。
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で取得して変更します。
vgg.classifier.instance_variable_get(:@modules).delete("6")
前節で得られた、画像データセットのtensorをVGG-16に与えて、各画像の特徴ベクトルを抽出します。annoy.rbでは、ベクトルをRuby Arrayで扱うので、to_aメソッドでtensorをRuby Arrayに変換しています。
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を選択しました。
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クラスタリングに似た方法でデータを二分割することで木構造を作りますが、このときにランダムにデータを選択します。木の作成にランダム性があるので、複数の木を組み合わせて、検索結果を安定させます。検索インデックスの作成のパラメータとして、木の数があります。
# 検索インデックスを作成する. 木の数は10本とした.
index.build(10)
これで検索インデックスができたので、実際に検索してみましょう。登録したデータの添字を指定して検索する get_nns_by_item メソッドがありますが、今回は、クエリが与えられたことを想定して、ベクトルで検索する方法を用いました。
# 添字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で作れます 💪
Discussion