SentenceTransformersを用いた画像類似度検索

背景

画像の類似度を検索する方法を試してみようと思い、sentence_transformers を用いて6つの画像から、検索された文字列 "Alessandro Squarzi Jacket Style" に最も近い画像を検索することにしてみた。

やったこと

画像の準備

必要なライブラリインポート

import os
import pprint
from PIL import Image
from sentence_transformers import SentenceTransformer

画像をベクトル化

# モデルのロード
model = SentenceTransformer('clip-ViT-B-32')

# 画像ディレクトリの設定
image_folder = 'backend/images'

# 画像のリストを取得
image_files = [f for f in os.listdir(image_folder) if f.endswith(('.png', '.jpg', '.jpeg'))]

# 画像のベクトル化
image_vectors = {}
for image_name in image_files:
    image_path = os.path.join(image_folder, image_name)
    image = Image.open(image_path)
    image_vector = model.encode([image])
    image_vectors[image_name] = image_vector

入力された文字から画像を検索

def find_matching_image(prompt, model, image_vectors):
    text_emb = model.encode([prompt])
    similarities = {}
    for image_name, img_emb in image_vectors.items():
        similarity = model.similarity(img_emb, text_emb).item()
        similarities[image_name] = similarity
    pprint.pprint(similarities)
    return max(similarities, key=similarities.get)


# 例: ユーザー入力プロンプト
user_prompt = "Alessandro Squarzi Jacket Style"
selected_image = find_matching_image(user_prompt, model, image_vectors)
if selected_image:
    print(f"正解: alessandro_squarzi3.jpg, 選択された画像: {selected_image}")
else:
    print("一致する画像が見つかりませんでした。")

結果

最終的な結果は以下のようになった

正解: alessandro_squarzi3.jpg, 選択された画像: unknown2.jpg

類似度も以下の通り。

{'alessandro_squarzi1.jpg': 0.22763186693191528,
 'alessandro_squarzi2.jpg': 0.2647595703601837,
 'alessandro_squarzi3.jpg': 0.27001553773880005,
 'junk1.jpg': 0.15288576483726501,
 'unknown1.jpg': 0.2811787724494934,
 'unknown2.jpg': 0.2970343232154846}

まとめ

期待通りの結果が得られなかった。
原因調べていきます。