Milvusによるテキスト画像検索

Text-to-image検索は、ユーザが自然言語のテキスト記述を使って画像を検索できる高度な技術です。これは、事前に訓練されたマルチモーダルモデルを活用し、テキストと画像の両方を共有された意味空間の埋め込みに変換し、類似性に基づいた比較を可能にします。

このチュートリアルでは、OpenAIのCLIP(Contrastive Language-Image Pretraining)モデルとmilvusを使って、テキストベースの画像検索を実装する方法を探ります。CLIPを用いて画像埋め込みを生成し、Milvusに保存し、効率的な類似検索を行います。

前提条件

始める前に、必要なパッケージとサンプルデータが揃っていることを確認してください。

依存関係のインストール

• pymilvus>=2.4.2:Milvusデータベースとやりとりするため。
• CLIPモデルを扱うためのclip
• 画像処理と可視化のためのpillow

$ pip install --upgrade pymilvus pillow
$ pip install git+https://github.com/openai/CLIP.git

サンプルデータのダウンロード

ImageNetデータセットのサブセット（100クラス、各クラス10画像）をサンプル画像として使います。以下のコマンドでサンプルデータをダウンロードし、ローカルフォルダ./images_folder に展開します：

$ wget https://github.com/towhee-io/examples/releases/download/data/reverse_image_search.zip
$ unzip -q reverse_image_search.zip -d images_folder

Milvusのセットアップ

先に進む前に、Milvusサーバのセットアップを行い、URI（オプションでトークン）を使って接続してください：

• Milvus Lite (便宜上推奨)：URIを./milvus.dbのようなローカルファイルに設定します。これは自動的にMilvus Liteを活用し、すべてのデータを単一のファイルに保存します。

• DockerまたはKubernetes（大規模データ用）：より大規模なデータセットを扱うには、DockerまたはKubernetesを使用して、よりパフォーマンスの高いMilvusサーバをデプロイします。この場合、http://localhost:19530 のようなサーバURIを使用して接続します。

• Zillizクラウド（マネージドサービス）：MilvusのフルマネージドクラウドサービスであるZilliz Cloudをご利用の場合は、URIにPublic Endpoint、トークンにAPI Keyを設定してください。

from pymilvus import MilvusClient

milvus_client = MilvusClient(uri="milvus.db")

開始

必要な依存関係やデータが揃ったところで、いよいよ機能抽出ツールをセットアップしてMilvusを使い始めましょう。このセクションでは、テキストから画像への検索システムを構築するための重要なステップを説明します。最後に、テキストクエリに基づいて画像を検索し、視覚化する方法を示します。

特徴抽出器の定義

画像とテキストの埋め込みを生成するために、事前に学習されたCLIPモデルを使用します。このセクションでは、事前に学習されたViT-B/32variant of CLIPをロードし、画像とテキストをエンコードするためのヘルパー関数を定義する：

• encode_image(image_path):画像を処理して特徴ベクトルにエンコードする。
• encode_text(text):テキストクエリを特徴ベクトルにエンコード

両関数とも、正確な余弦類似度計算に不可欠な単位長にベクトルを変換することにより、一貫した比較を保証するために出力特徴を正規化する。

import clip
from PIL import Image


# Load CLIP model
model_name = "ViT-B/32"
model, preprocess = clip.load(model_name)
model.eval()


# Define a function to encode images
def encode_image(image_path):
    image = preprocess(Image.open(image_path)).unsqueeze(0)
    image_features = model.encode_image(image)
    image_features /= image_features.norm(
        dim=-1, keepdim=True
    )  # Normalize the image features
    return image_features.squeeze().tolist()


# Define a function to encode text
def encode_text(text):
    text_tokens = clip.tokenize(text)
    text_features = model.encode_text(text_tokens)
    text_features /= text_features.norm(
        dim=-1, keepdim=True
    )  # Normalize the text features
    return text_features.squeeze().tolist()

データの取り込み

セマンティック画像検索を可能にするために、まずすべての画像の埋め込みを生成し、効率的なインデックス付けと検索のためにベクトルデータベースに格納する必要があります。このセクションでは、画像データをmilvusに取り込むためのステップバイステップのガイドを提供します。

1.Milvusコレクションの作成

画像の埋め込みを保存する前に、Milvusコレクションを作成する必要があります。以下のコードは、デフォルトのCOSINE メトリックタイプでクイックセットアップモードでコレクションを作成する方法を示しています。コレクションは以下のフィールドを含む：

• id:オートIDが有効なプライマリフィールド。

• vector:浮動小数点ベクトル埋め込みを格納するフィールド。

カスタムスキーマが必要な場合は、Milvusのドキュメントを参照してください。

collection_name = "image_collection"

# Drop the collection if it already exists
if milvus_client.has_collection(collection_name):
    milvus_client.drop_collection(collection_name)

# Create a new collection in quickstart mode
milvus_client.create_collection(
    collection_name=collection_name,
    dimension=512,  # this should match the dimension of the image embedding
    auto_id=True,  # auto generate id and store in the id field
    enable_dynamic_field=True,  # enable dynamic field for scalar fields
)

2.Milvusへのデータ挿入

このステップでは、あらかじめ定義された画像エンコーダを使用して、サンプル・データ・ディレクトリ内のすべてのJPEG画像の埋め込みデータを生成します。これらのエンベッディングは、対応するファイルパスとともにMilvusコレクションに挿入されます。コレクション内の各エントリは以下から構成される：

• 埋め込みベクトル：画像の数値表現。vector フィールドに格納される。
• ファイルパス：参照用の画像ファイルの場所。動的フィールドとしてフィールドfilepath に格納される。

import os
from glob import glob


image_dir = "./images_folder/train"
raw_data = []

for image_path in glob(os.path.join(image_dir, "**/*.JPEG")):
    image_embedding = encode_image(image_path)
    image_dict = {"vector": image_embedding, "filepath": image_path}
    raw_data.append(image_dict)
insert_result = milvus_client.insert(collection_name=collection_name, data=raw_data)

print("Inserted", insert_result["insert_count"], "images into Milvus.")

Inserted 1000 images into Milvus.

検索の実行

それでは、テキストクエリの例を使って検索を実行してみましょう。これは、指定されたテキスト記述との意味的類似性に基づいて、最も関連性の高い画像を検索します。

query_text = "a white dog"
query_embedding = encode_text(query_text)

search_results = milvus_client.search(
    collection_name=collection_name,
    data=[query_embedding],
    limit=10,  # return top 10 results
    output_fields=["filepath"],  # return the filepath field
)

結果を視覚化してみましょう：

from IPython.display import display


width = 150 * 5
height = 150 * 2
concatenated_image = Image.new("RGB", (width, height))

result_images = []
for result in search_results:
    for hit in result:
        filename = hit["entity"]["filepath"]
        img = Image.open(filename)
        img = img.resize((150, 150))
        result_images.append(img)

for idx, img in enumerate(result_images):
    x = idx % 5
    y = idx // 5
    concatenated_image.paste(img, (x * 150, y * 150))
print(f"Query text: {query_text}")
print("\nSearch results:")
display(concatenated_image)

Query text: a white dog

Search results:

png