類似画像検索
facebook が作ったベクトル検索エンジンの faiss を使って類似画像検索ができる、というなにかの記事を見てへーと思い、自分の日記画像の中から類似画像を見つけるというのをやってみたら、思っていた以上に簡単だったのでお気持ちを記事に。
類似画像検索は画像を embedding (ベクトル表現)に変えて類似度が高いものを見つければ良い。のだけど、そもども画像をどうやってベクトルに変えるのか、というのが、CNNの流れの画像分類用の学習モデルを使えば良い、というのを知ってなるほど！という感じであった。メルカリも MobileNet V2 を使っている(いた?)らしい。そしてこいつの良いところが、基本が学習済みモデルに転移学習すらさせることなく、学習済み重みを含むモデルに普通に推論させるだけで(大抵の場合)良い。
なおこの記事では faiss を使っての類似画像検索は行っていないので、faiss について知りたい人には有益でない内容。
具体的なやり方
やり方は簡単で、基本大抵のCNN系画像分類モデルは最後の層で softmax 関数を通して特定クラスの確率を出力するのだけど、その直前の層あたりではいい感じに特徴が含まれた層があるので、それをばばーんと取り出すだけ。画像分類モデルはお好きなものを使えば良く、自分は TensorFlow を使い慣れているので、そのへんの最近のモデルである EfficientNet を使う。
model = tf.keras.applications.EfficientNetB0(
include_top=False,
weights='imagenet',
pooling='avg',
)EfficientNetB0 の引数には、top を除き、学習済み重みは imagenet を使い、pooling では global avg pooling になるようにする。これで適当な画像を model に推論(というか変換だけど)させると、1280次元のベクトルへと変換された。
images_tensor = image_files_to_tensor([image_path])
print(images_tensor.shape)
image_tensor_embedding = model.predict(images_tensor)
print(image_tensor_embedding)
print(image_tensor_embedding.shape)
(1, 224, 224, 3)
[[-0.03652444 -0.05792176 -0.10945524 ... 0.11447129 -0.13936797
0.21209691]]
(1, 1280)これだけだとイマイチピンと来ないと思うので、私の写真から適当な画像10枚で見てみる。
この10枚の画像では、例えば0,1は似てるとか、8と9は似てそう、みたいに人間が見ると思える。先程と同じようにmodelに適用させる。
images_tensor = image_files_to_tensor(images_path)
images_tensor_embedding = model.predict(images_tensor)
print(images_tensor_embedding)
print(images_tensor_embedding.shape)
[[-0.10945629 -0.09343709 -0.08223398 ... 1.0024792 -0.1611993
0.01493335]
[-0.03652436 -0.05792165 -0.10945524 ... 0.11447226 -0.1393679
0.21209645]
[-0.10479932 -0.10463432 -0.08078936 ... -0.023954 -0.03211505
-0.13908449]
...
[-0.19969083 -0.1597161 -0.175438 ... -0.2206025 -0.04785771
0.03392322]
[ 0.1287042 -0.07741662 -0.10888506 ... 0.8499286 -0.20345287
0.06938681]
[-0.01477872 -0.09264231 -0.15021208 ... 0.30814508 -0.19284736
0.28705594]]
(10, 1280)10枚の画像から、10x1280次元のembeddings表現が取れた。この embeddings からコサイン類似度をとれば、似ているものが解る。
from sklearn.metrics.pairwise import cosine_similarity
sim_res = cosine_similarity(images_tensor_embedding, images_tensor_embedding)
sim_res
array([[ 1. , 0.44505394, -0.04161356, 0.01666282, 0.04543769,
0.24226916, 0.06981121, 0.00515541, 0.01611625, 0.06467595],
[ 0.44505394, 1.0000002 , 0.00313169, 0.06676795, 0.07920375,
0.25869107, 0.05456577, 0.07960878, -0.0146864 , 0.02522465],
[-0.04161356, 0.00313169, 0.9999999 , 0.5507376 , 0.4239326 ,
-0.02977732, 0.3225319 , 0.62029684, 0.07330316, 0.03695932],
[ 0.01666282, 0.06676795, 0.5507376 , 1.0000001 , 0.4721935 ,
0.05535758, 0.43453646, 0.6822141 , 0.17620735, 0.12906513],
[ 0.04543769, 0.07920375, 0.4239326 , 0.4721935 , 0.9999998 ,
0.08597542, 0.17120099, 0.40617043, 0.0032471 , 0.02258891],
[ 0.24226916, 0.25869107, -0.02977732, 0.05535758, 0.08597542,
0.99999976, 0.06380752, 0.07490596, 0.08929685, 0.13794066],
[ 0.06981121, 0.05456577, 0.3225319 , 0.43453646, 0.17120099,
0.06380752, 1. , 0.49228954, 0.22836211, 0.29995295],
[ 0.00515541, 0.07960878, 0.62029684, 0.6822141 , 0.40617043,
0.07490596, 0.49228954, 1.0000004 , 0.18286222, 0.11434785],
[ 0.01611625, -0.0146864 , 0.07330316, 0.17620735, 0.0032471 ,
0.08929685, 0.22836211, 0.18286222, 0.9999999 , 0.60561 ],
[ 0.06467595, 0.02522465, 0.03695932, 0.12906513, 0.02258891,
0.13794066, 0.29995295, 0.11434785, 0.60561 , 1.0000004 ]],
dtype=float32)自分自身の画像とは一番似ている(というか同じ)なので、対角成分は float のゆらぎを無視すればほぼ1となる。ので、対角成分を0としてしまい、argsort でスコア順のインデクスを取り出す。
# 対角成分を 0 に
sim_res_diag = np.diag(sim_res)
sim_res_diag.flags.writeable = True
sim_res_diag.fill(0)
sim_res_diag.flags.writeable = False
# スコア順に index を取得
max_score_indexes = (-sim_res).argsort(axis=1)
max_score_indexes
array([[1, 5, 6, 9, 4, 3, 8, 7, 0, 2],
[0, 5, 7, 4, 3, 6, 9, 2, 1, 8],
[7, 3, 4, 6, 8, 9, 1, 2, 5, 0],
[7, 2, 4, 6, 8, 9, 1, 5, 0, 3],
[3, 2, 7, 6, 5, 1, 0, 9, 8, 4],
[1, 0, 9, 8, 4, 7, 6, 3, 5, 2],
[7, 3, 2, 9, 8, 4, 0, 5, 1, 6],
[3, 2, 6, 4, 8, 9, 1, 5, 0, 7],
[9, 6, 7, 3, 5, 2, 0, 4, 8, 1],
[8, 6, 5, 3, 7, 0, 2, 1, 4, 9]])さて、実際の類似度スコアの上位3つがどれだけ似てるのか。0と1は似ている、0や1と5(ラーメン)は食べ物と皿という点では似ている。0のtop-3に6があるが、スコアを見るとそもそも低いため全然似ていない。2 と 7 は空だったり構図だったりが似ている、8と9も建物内で構図も似ている、といった感じで似てる感がわかる。
この日記の画像のみから類似画像を算出
この日記にはられている、5600枚ぐらいの画像でやってみる。image_embeddingsというライブラリを使うと、EfficientNetB0 使いつつ、機械学習パイプラインで良い感じに回すための tfrecords の作成やその他諸々をやってくれる。というかこの記事の大半は image_embeddings のレポジトリに含まれる notebook で書かれていることのシンプル版だ。
5600枚の画像をGPUのRTX3900にくべると、だいたい推論に11秒。1枚あたり2msぐらいで終わる。なおRyzen 3900XTでは推論に94秒で1枚あたり17msぐらい。推論タスクなのでCPUでもそこまで時間がかからない。推論自体の時間は線形増加だろうから、もし560万枚の画像にたいして推論を行ったとしても、推論自体の時間はGPU1つで11000秒ぐらいで終わりそうだ。
5600枚の embedding の中から20枚をランダムに抽出してスコア順に並べるとこんな感じ。おおー、だいぶ類似っぽい画像が出せた、簡単。どれも自分が撮った写真なので、自分自身が見てて面白い。
production での利用を考える
embedding 表現は、学習済み重みのモデルに変更がなければ、毎回同じ embedding が得られるので、データを追加していく分には問題ない。
ただ、増えたデータに対して類似する embedding を算出するのがどんどん高コストになっていく。この日記の類推画像は高々(5600, 1280)のデータの計算なので、PCであっという間にコサイン類似度を出せるが、数千万、数億…みたいに増えると大変になっていく。
ある程度の規模以上になると、前述の大規模運用事例もある faiss のようなベクトルに特化した検索エンジンをつかうことで、本番環境でも(運用の大変さはあるだろうが)スケールさせつつ使えそうだ。
また、最近出た GCP の Vertex AI にはVertex Matching Engineというフルマネージドなベクトル検索エンジンがあって、値段的なコストはあれど、Public Cloud におまかせもできそう。
最近は Elasticsearch にもベクトル検索が入ったらしいので、ESでもできそう。
汎用化していく機械学習
今回類似画像検索を行ってみて、転移学習すら必要なく学習済みモデルを使うだけで embedding が取得でき、かつそれを検索するための高性能のベクトル検索エンジンもOSSからフルマネージドなものまで今は揃っている。
類似画像検索のような embedding を使った汎用タスクも、API感覚で今以上にだれでも使えるようになっていくのだろうなーと改めて思ったのであった。今回の記事の内容も、機械学習的な難しさは一切なくて、ツールの使い方さえ知っていれば誰でもできる系の内容だしね。
というか、ちょっと調べたら Apple の TuriCreate にズバリの Image Similarity があったので、サクッと試したい人やiOSで類似画像検索を動かしたい人は、TuriCreate を使うのも良いかも知れない。
