使用 Milvus 進行文字轉圖像搜尋

文字到圖像搜尋是一項先進的技術，可讓使用者使用自然語言文字描述搜尋圖像。它利用預先訓練的多模態模型，將文字和圖像轉換成共享語意空間中的嵌入，從而實現基於相似性的比較。

在本教程中，我們將探討如何使用 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 伺服器。在這種情況下，請使用伺服器 URI 連線，例如 http://localhost:19530。

• Zilliz Cloud (管理服務)：如果您使用的是Zilliz Cloud，Milvus 的完全管理雲端服務，請將公共端點設定為 URI，並將 API Key 設定為 token。

from pymilvus import MilvusClient

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

開始使用

現在您已經擁有必要的相依性與資料，是時候設定功能擷取器並開始使用 Milvus。本節將介紹建立文字到圖片搜尋系統的關鍵步驟。最後，我們將示範如何根據文字查詢擷取圖片並將其視覺化。

定義特徵萃取器

我們將使用預先訓練好的 CLIP 模型來產生圖像和文字嵌入。在本節中，我們將載入 CLIP 的預訓ViT-B/32變體，並定義用於編碼圖像和文字的輔助函式：

• 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