Text-zu-Bildern-Suche mit Milvus

Die Text-zu-Bild-Suche ist eine fortschrittliche Technologie, die es Benutzern ermöglicht, anhand von Textbeschreibungen in natürlicher Sprache nach Bildern zu suchen. Sie nutzt ein vorab trainiertes multimodales Modell, um sowohl Text als auch Bilder in Einbettungen in einem gemeinsamen semantischen Raum zu konvertieren und ermöglicht so Ähnlichkeitsvergleiche.

In diesem Tutorial werden wir untersuchen, wie textbasierte Bildsuche mit Hilfe des CLIP-Modells (Contrastive Language-Image Pretraining) von OpenAI und Milvus implementiert werden kann. Wir werden Bildeinbettungen mit CLIP erzeugen, sie in Milvus speichern und effiziente Ähnlichkeitssuchen durchführen.

Voraussetzungen

Bevor Sie beginnen, stellen Sie sicher, dass Sie alle erforderlichen Pakete und Beispieldaten bereithalten.

Installieren Sie die Abhängigkeiten

• pymilvus>=2.4.2 für die Interaktion mit der Milvus-Datenbank
• clip für die Arbeit mit dem CLIP-Modell
• pillow für die Bildverarbeitung und Visualisierung

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

Herunterladen von Beispieldaten

Wir werden eine Teilmenge des ImageNet-Datensatzes (100 Klassen, 10 Bilder für jede Klasse) als Beispielbilder verwenden. Der folgende Befehl lädt die Beispieldaten herunter und extrahiert sie in den lokalen Ordner ./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 einrichten

Bevor Sie fortfahren, richten Sie Ihren Milvus-Server ein und verbinden Sie sich mit Ihrer URI (und optional mit einem Token):

• Milvus Lite (aus Gründen der Bequemlichkeit empfohlen): Setzen Sie den URI auf eine lokale Datei, z. B. ./milvus.db. Dadurch wird Milvus Lite automatisch genutzt, um alle Daten in einer einzigen Datei zu speichern.

• Docker oder Kubernetes (für große Datenmengen): Für die Verarbeitung größerer Datenmengen können Sie einen leistungsfähigeren Milvus-Server mit Docker oder Kubernetes bereitstellen. Verwenden Sie in diesem Fall die Server-URI, z. B. http://localhost:19530, um eine Verbindung herzustellen.

• Zilliz Cloud (Managed Service): Wenn Sie Zilliz Cloud verwenden, den vollständig verwalteten Cloud-Service von Milvus, legen Sie den öffentlichen Endpunkt als URI und den API-Schlüssel als Token fest.

from pymilvus import MilvusClient

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

Erste Schritte

Jetzt, da Sie die notwendigen Abhängigkeiten und Daten haben, ist es an der Zeit, Feature-Extraktoren einzurichten und mit Milvus zu arbeiten. Dieser Abschnitt führt Sie durch die wichtigsten Schritte beim Aufbau eines Text-zu-Bild-Suchsystems. Abschließend zeigen wir Ihnen, wie Sie Bilder auf der Grundlage von Textabfragen abrufen und visualisieren können.

Definieren von Merkmalsextraktoren

Wir werden ein vortrainiertes CLIP-Modell verwenden, um Bild- und Texteinbettungen zu erzeugen. In diesem Abschnitt laden wir die vortrainierte ViT-B/32-Variante von CLIP und definieren Hilfsfunktionen zur Kodierung von Bild und Text:

• encode_image(image_path): Verarbeitet und kodiert Bilder in Feature-Vektoren
• encode_text(text): Kodierung von Textanfragen in Merkmalsvektoren

Beide Funktionen normalisieren die ausgegebenen Merkmale, um konsistente Vergleiche zu gewährleisten, indem sie die Vektoren in Einheitslängen umwandeln, was für genaue Kosinusähnlichkeitsberechnungen unerlässlich ist.

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

Dateneingabe

Um eine semantische Bildsuche zu ermöglichen, müssen wir zunächst Einbettungen für alle Bilder erzeugen und sie in einer Vektordatenbank speichern, um sie effizient indizieren und abrufen zu können. Dieser Abschnitt enthält eine schrittweise Anleitung zur Aufnahme von Bilddaten in Milvus.

1. Milvus-Sammlung erstellen

Bevor Sie Bildeinbettungen speichern können, müssen Sie eine Milvus-Sammlung erstellen. Der folgende Code veranschaulicht, wie eine Sammlung im Schnelleinstellungsmodus mit dem Standard-Metrik-Typ COSINE erstellt wird. Die Sammlung enthält die folgenden Felder:

• id: Ein Primärfeld mit aktivierter automatischer ID.

• vector: Ein Feld zum Speichern von Fließkomma-Vektoreinbettungen.

Wenn Sie ein benutzerdefiniertes Schema benötigen, finden Sie detaillierte Anweisungen in der Milvus-Dokumentation.

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. Daten in Milvus einfügen

In diesem Schritt verwenden wir einen vordefinierten Bildkodierer, um Einbettungen für alle JPEG-Bilder im Beispieldatenverzeichnis zu erzeugen. Diese Einbettungen werden dann in die Milvus-Sammlung eingefügt, zusammen mit den entsprechenden Dateipfaden. Jeder Eintrag in der Sammlung besteht aus:

• Einbettungsvektor: Die numerische Darstellung des Bildes. Sie wird im Feld vector gespeichert.
• Dateipfad: Der Speicherort der Bilddatei als Referenz. Wird im Feld filepath als dynamisches Feld gespeichert.

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.

Durchführen einer Suche

Führen wir nun eine Suche mit einer Beispiel-Textabfrage durch. Dadurch werden die relevantesten Bilder auf der Grundlage ihrer semantischen Ähnlichkeit mit der angegebenen Textbeschreibung abgerufen.

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
)

Visualisieren Sie die Ergebnisse:

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