Pencarian Teks-ke-Gambar dengan Milvus

Pencarian teks-ke-gambar adalah teknologi canggih yang memungkinkan pengguna mencari gambar menggunakan deskripsi teks bahasa alami. Teknologi ini memanfaatkan model multimodal yang telah terlatih untuk mengubah teks dan gambar menjadi sematan dalam ruang semantik bersama, sehingga memungkinkan perbandingan berbasis kemiripan.

Dalam tutorial ini, kita akan mengeksplorasi cara mengimplementasikan pengambilan gambar berbasis teks menggunakan model CLIP (Contrastive Language-Image Pretraining) dari OpenAI dan Milvus. Kita akan membuat penyematan gambar dengan CLIP, menyimpannya di Milvus, dan melakukan pencarian kemiripan yang efisien.

Prasyarat

Sebelum memulai, pastikan Anda sudah menyiapkan semua paket yang diperlukan dan data contoh.

Instal ketergantungan

• pymilvus>=2.4.2 untuk berinteraksi dengan basis data Milvus
• clip untuk bekerja dengan model CLIP
• pillow untuk pemrosesan gambar dan visualisasi

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

Mengunduh data contoh

Kami akan menggunakan subset dari dataset ImageNet (100 kelas, 10 gambar untuk setiap kelas) sebagai gambar contoh. Perintah berikut ini akan mengunduh data contoh dan mengekstraknya ke folder lokal ./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

Menyiapkan Milvus

Sebelum melanjutkan, siapkan server Milvus Anda dan sambungkan dengan menggunakan URI Anda (dan secara opsional, token):

• Milvus Lite (Disarankan untuk kenyamanan): Atur URI ke berkas lokal, seperti ./milvus.db. Ini secara otomatis memanfaatkan Milvus Lite untuk menyimpan semua data dalam satu berkas.

• Docker atau Kubernetes (Untuk Data Berskala Besar): Untuk menangani kumpulan data yang lebih besar, gunakan server Milvus yang lebih berkinerja dengan menggunakan Docker atau Kubernetes. Dalam hal ini, gunakan URI server, seperti http://localhost:19530, untuk terhubung.

• Zilliz Cloud (Layanan Terkelola): Jika Anda menggunakan Zilliz Cloud, layanan cloud terkelola Milvus, tetapkan Public Endpoint sebagai URI dan API Key sebagai token.

from pymilvus import MilvusClient

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

Memulai

Sekarang setelah Anda memiliki dependensi dan data yang diperlukan, saatnya untuk menyiapkan ekstraktor fitur dan mulai bekerja dengan Milvus. Bagian ini akan memandu Anda melalui langkah-langkah penting dalam membangun sistem pencarian teks-ke-gambar. Terakhir, kami akan mendemonstrasikan cara mengambil dan memvisualisasikan gambar berdasarkan kueri teks.

Menentukan ekstraktor fitur

Kita akan menggunakan model CLIP yang telah dilatih untuk menghasilkan penyematan gambar dan teks. Pada bagian ini, kita akan memuat varian ViT-B/32 yang telah dilatih dari CLIP dan mendefinisikan fungsi-fungsi pembantu untuk menyandikan gambar dan teks:

• encode_image(image_path): Memproses dan menyandikan gambar menjadi vektor fitur
• encode_text(text): Menyandikan kueri teks menjadi vektor fitur

Kedua fungsi tersebut menormalkan fitur keluaran untuk memastikan perbandingan yang konsisten dengan mengubah vektor menjadi satuan panjang, yang sangat penting untuk penghitungan kemiripan kosinus yang akurat.

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()

Konsumsi Data

Untuk mengaktifkan pencarian gambar semantik, pertama-tama kita perlu menghasilkan penyematan untuk semua gambar dan menyimpannya dalam basis data vektor untuk pengindeksan dan pengambilan yang efisien. Bagian ini memberikan panduan langkah demi langkah untuk memasukkan data gambar ke dalam Milvus.

1. Membuat Koleksi Milvus

Sebelum menyimpan penyematan gambar, Anda perlu membuat koleksi Milvus. Kode berikut ini menunjukkan cara membuat koleksi dalam mode penyiapan cepat dengan tipe metrik COSINE default. Koleksi ini mencakup bidang-bidang berikut ini:

• id: Bidang utama dengan ID otomatis diaktifkan.

• vector: Bidang untuk menyimpan penyematan vektor floating-point.

Jika Anda memerlukan skema khusus, lihat dokumentasi Milvus untuk instruksi terperinci.

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. Memasukkan Data ke dalam Milvus

Pada langkah ini, kita menggunakan penyandi gambar yang sudah ditentukan untuk menghasilkan penyematan untuk semua gambar JPEG dalam direktori data contoh. Embedding ini kemudian dimasukkan ke dalam koleksi Milvus, bersama dengan jalur file yang sesuai. Setiap entri dalam koleksi terdiri dari:

• Vektor penyematan: Representasi numerik dari gambar. Disimpan di bidang vector.
• Jalur file: Lokasi file gambar untuk referensi. Disimpan di bidang filepath sebagai bidang dinamis.

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.

Melakukan Pencarian

Sekarang, mari kita jalankan pencarian menggunakan contoh kueri teks. Ini akan mengambil gambar yang paling relevan berdasarkan kemiripan semantik dengan deskripsi teks yang diberikan.

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
)

Memvisualisasikan hasil:

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