Ricerca da testo a immagine con Milvus

La ricerca da testo a immagine è una tecnologia avanzata che consente agli utenti di cercare immagini utilizzando descrizioni testuali in linguaggio naturale. Sfrutta un modello multimodale preaddestrato per convertire sia il testo che le immagini in incorporazioni in uno spazio semantico condiviso, consentendo confronti basati sulla somiglianza.

In questa esercitazione esploreremo come implementare il recupero di immagini basato sul testo utilizzando il modello CLIP (Contrastive Language-Image Pretraining) di OpenAI e Milvus. Genereremo embeddings di immagini con CLIP, li memorizzeremo in Milvus ed eseguiremo ricerche di similarità efficienti.

Prerequisiti

Prima di iniziare, assicuratevi di avere pronti tutti i pacchetti necessari e i dati di esempio.

Installare le dipendenze

• pymilvus>=2.4.2 per interagire con il database di Milvus
• clip per lavorare con il modello CLIP
• pillow per l'elaborazione e la visualizzazione delle immagini

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

Scaricare i dati di esempio

Utilizzeremo un sottoinsieme del dataset ImageNet (100 classi, 10 immagini per ogni classe) come immagini di esempio. Il comando seguente scarica i dati di esempio e li estrae nella cartella locale ./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

Configurare Milvus

Prima di procedere, configurare il server Milvus e connettersi utilizzando il proprio URI (e, facoltativamente, un token):

• Milvus Lite (consigliato per comodità): Impostare l'URI su un file locale, come ./milvus.db. In questo modo si sfrutta automaticamente Milvus Lite per memorizzare tutti i dati in un unico file.

• Docker o Kubernetes (per dati su larga scala): Per gestire insiemi di dati più grandi, è possibile distribuire un server Milvus più performante utilizzando Docker o Kubernetes. In questo caso, per connettersi utilizzare l'URI del server, ad esempio http://localhost:19530.

• Zilliz Cloud (servizio gestito): Se si utilizza Zilliz Cloud, il servizio cloud completamente gestito da Milvus, impostare l'Endpoint pubblico come URI e la Chiave API come token.

from pymilvus import MilvusClient

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

Iniziare

Ora che si dispone delle dipendenze e dei dati necessari, è il momento di impostare gli estrattori di funzioni e iniziare a lavorare con Milvus. Questa sezione illustra le fasi principali della costruzione di un sistema di ricerca da testo a immagine. Infine, dimostreremo come recuperare e visualizzare le immagini in base alle query di testo.

Definire gli estrattori di caratteristiche

Utilizzeremo un modello CLIP preaddestrato per generare embeddings di immagini e testo. In questa sezione, carichiamo la variante preaddestrata ViT-B/32 di CLIP e definiamo le funzioni di aiuto per la codifica di immagini e testo:

• encode_image(image_path): elabora e codifica le immagini in vettori di caratteristiche
• encode_text(text): Codifica le query di testo in vettori di caratteristiche

Entrambe le funzioni normalizzano le caratteristiche in uscita per garantire confronti coerenti, convertendo i vettori in vettori di lunghezza unitaria, essenziali per un calcolo accurato della similarità del coseno.

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

Ingestione dei dati

Per consentire la ricerca semantica delle immagini, dobbiamo innanzitutto generare le incorporazioni per tutte le immagini e memorizzarle in un database vettoriale per un'indicizzazione e un recupero efficienti. Questa sezione fornisce una guida passo passo per l'inserimento dei dati delle immagini in Milvus.

1. Creare la raccolta Milvus

Prima di memorizzare le incorporazioni di immagini, è necessario creare una raccolta Milvus. Il codice seguente mostra come creare una raccolta in modalità rapida con il tipo di metrica predefinito COSINE. La collezione comprende i seguenti campi:

• id: Un campo primario con ID automatico abilitato.

• vector: Un campo per memorizzare le incorporazioni vettoriali in virgola mobile.

Se si desidera uno schema personalizzato, consultare la documentazione di Milvus per istruzioni dettagliate.

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. Inserire i dati in Milvus

In questa fase si utilizza un codificatore di immagini predefinito per generare le incorporazioni di tutte le immagini JPEG presenti nella directory dei dati di esempio. Questi embeddings vengono poi inseriti nella raccolta di Milvus, insieme ai percorsi dei file corrispondenti. Ogni voce della raccolta è composta da:

• Vettore di incorporazione: La rappresentazione numerica dell'immagine. Memorizzato nel campo vector.
• Percorso del file: La posizione del file dell'immagine come riferimento. Memorizzato nel campo filepath come campo dinamico.

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.

Eseguire una ricerca

Eseguiamo ora una ricerca utilizzando una query di testo di esempio. In questo modo verranno recuperate le immagini più rilevanti in base alla loro somiglianza semantica con la descrizione testuale fornita.

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
)

Visualizzare i risultati:

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