ディープラーニングで誰でも類似画像検索

#背景
「ある画像と最も似ている画像を探索したい」場合に、どのようにディープラーニングを活用できるか知見がなかったので、以下の2つを参考に改変しています。

近似最近傍探索ライブラリ比較
機械学習モデルという『関数』を使ったはじめての類似画像検索

annoyはインストールが面倒だったので、nmslibを使っています。

#結果

内容解説の前に結果を示します。

ある画像（猫)に似ている画像をデータベース（２０枚の子猫、子犬の画像）内を検索した結果です。
１枚を除き、猫画像が似ていると判定されていました。

＊誤判定の１枚は被り物をして、とても猫には見えないので誤判定も納得の結果です。

#概要
探索対象の画像群にはラベリングはされていないため、教師あり学習をすることはできません。
一つのアプローチとしては、学習済みのモデルを活用して全結合層の出力等を特徴量として使うことができます。

今回はVGG19でImagenetを学習した重みづけを用いました。

また、tkinterでデータベースと画像の選択を出来るようにしてあるので、
プログラム未経験者でも動かすことはできるかと思います。

#コード

import glob
from pathlib import Path
import tkinter
import tkinter.filedialog

#License
#The MIT License
import keras
from keras.models import Model
from keras.layers import Input, Dense
from keras.preprocessing import image
from keras.applications.vgg19 import preprocess_input

#License
#These weights are ported from the ones released by VGG at Oxford under the Creative Commons Attribution License.
#https://keras.io/applications/
from keras.applications.vgg19 import VGG19, preprocess_input

#Apache License Version 2.0
#https://github.com/nmslib/nmslib/blob/master/README.md
import nmslib

#https://numpy.org/license.html
import numpy as np

current_path = Path.cwd()

# refer https://qiita.com/wasnot/items/20c4f30a529ae3ed5f52
# refer https://qiita.com/K-jun/items/cab923d49a939a8486fc

def main():
    print("データベースを選択してください")
    print("サブディレクトリ内の画像もすべて検索対象となります")

    data_folder_path = tkinter.filedialog.askdirectory(initialdir = current_path,
                        title = 'choose data folder')

    print("データベースと比較したい画像を選択してください")
    test_img_path = tkinter.filedialog.askopenfilename(initialdir = current_path,
                        title = 'choose test image', filetypes = [('image file', '*.jpeg;*jpg;*png')])

    base_model = VGG19(weights="imagenet")
    base_model.summary()
    #outputsを"fc2"と指定し、2番目の全結合層を出力します
    model = Model(inputs=base_model.input, outputs=base_model.get_layer("fc2").output)

    test_img = image.load_img(test_img_path, target_size=(224, 224))
    x = image.img_to_array(test_img)
    x = np.expand_dims(x, axis=0)
    x = preprocess_input(x)
    test_fc2_features = model.predict(x)
    
    #選択したフォルダに存在するpng,jpeg,jpgをサブディレクトリも含めて抽出
    png_list  = glob.glob(data_folder_path + "/**/*.png", recursive=True)
    jpeg_list = glob.glob(data_folder_path + "/**/*.jpeg", recursive=True)
    jpg_list  = glob.glob(data_folder_path + "/**/*.jpg", recursive=True)
    image_list = png_list + jpeg_list + jpg_list

    fc2_list = []
    for image_path in image_list:
        img = image.load_img(image_path, target_size=(224, 224))
        x = image.img_to_array(img)
        x = np.expand_dims(x, axis=0)
        x = preprocess_input(x)
        fc2_features = model.predict(x)
        fc2_list.append(fc2_features[0])

    index = nmslib.init(method='hnsw', space='cosinesimil')
    index.addDataPointBatch(fc2_list)
    index.createIndex({'post': 2}, print_progress=True)
    ids, distances = index.knnQuery(test_fc2_features, k=len(image_list))
    result = [image_list[i] for i in ids]

    print(ids)
    print(distances)
    print(result)

    print("選択した画像は " , test_img_path, " です")
    print("選択した画像に似ている順に表示します")
    for i, id in enumerate(ids):
        print(image_list[id], " : 距離： ", distances[i])

main()