Recherche texte-image avec Milvus

La recherche texte-image est une technologie avancée qui permet aux utilisateurs de rechercher des images à l'aide de descriptions textuelles en langage naturel. Elle s'appuie sur un modèle multimodal pré-entraîné pour convertir à la fois le texte et les images en encastrements dans un espace sémantique partagé, permettant des comparaisons basées sur la similarité.

Dans ce tutoriel, nous allons explorer comment mettre en œuvre la recherche d'images basée sur le texte en utilisant le modèle CLIP (Contrastive Language-Image Pretraining) de l'OpenAI et Milvus. Nous allons générer des encastrements d'images avec CLIP, les stocker dans Milvus et effectuer des recherches de similarité efficaces.

Conditions préalables

Avant de commencer, assurez-vous que vous disposez de tous les paquets nécessaires et des données d'exemple.

Installer les dépendances

• pymilvus>=2.4.2 pour interagir avec la base de données Milvus
• clip pour travailler avec le modèle CLIP
• pillow pour le traitement et la visualisation d'images

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

Télécharger des données d'exemple

Nous utiliserons un sous-ensemble du jeu de données ImageNet (100 classes, 10 images pour chaque classe) comme images d'exemple. La commande suivante télécharge les données d'exemple et les extrait dans le dossier local ./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

Configuration de Milvus

Avant de poursuivre, configurez votre serveur Milvus et connectez-vous à l'aide de votre URI (et éventuellement d'un jeton) :

• Milvus Lite (recommandé pour des raisons de commodité): Définissez l'URI sur un fichier local, tel que ./milvus.db. Cela permet de tirer automatiquement parti de Milvus Lite pour stocker toutes les données dans un seul fichier.

• Docker ou Kubernetes (pour les données à grande échelle): Pour traiter des ensembles de données plus importants, déployez un serveur Milvus plus performant à l'aide de Docker ou Kubernetes. Dans ce cas, utilisez l'URI du serveur, tel que http://localhost:19530, pour vous connecter.

• Zilliz Cloud (Managed Service): Si vous utilisez Zilliz Cloud, le service en nuage entièrement géré de Milvus, définissez le point de terminaison public comme URI et la clé API comme jeton.

from pymilvus import MilvusClient

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

Démarrage

Maintenant que vous disposez des dépendances et des données nécessaires, il est temps de configurer les extracteurs de fonctionnalités et de commencer à travailler avec Milvus. Cette section vous guidera à travers les étapes clés de la construction d'un système de recherche texte-image. Enfin, nous montrerons comment récupérer et visualiser des images sur la base de requêtes textuelles.

Définir les extracteurs de caractéristiques

Nous utiliserons un modèle CLIP pré-entraîné pour générer des encastrements d'images et de textes. Dans cette section, nous chargeons la variante ViT-B/32 pré-entraînée de CLIP et définissons les fonctions d'aide pour encoder les images et le texte :

• encode_image(image_path): Traite et encode les images en vecteurs de caractéristiques
• encode_text(text): Encode les requêtes textuelles en vecteurs de caractéristiques

Les deux fonctions normalisent les caractéristiques de sortie pour garantir des comparaisons cohérentes en convertissant les vecteurs en longueur unitaire, ce qui est essentiel pour des calculs de similarité en cosinus précis.

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

Ingestion des données

Pour permettre la recherche sémantique d'images, nous devons d'abord générer des embeddings pour toutes les images et les stocker dans une base de données vectorielle pour une indexation et une recherche efficaces. Cette section fournit un guide étape par étape pour l'ingestion de données d'images dans Milvus.

1. Créer une collection Milvus

Avant de stocker les incorporations d'images, vous devez créer une collection Milvus. Le code suivant montre comment créer une collection en mode installation rapide avec le type de métrique COSINE par défaut. La collection comprend les champs suivants :

• id: Un champ primaire avec ID automatique activé.

• vector: Un champ pour stocker les intégrations de vecteurs en virgule flottante.

Si vous avez besoin d'un schéma personnalisé, reportez-vous à la documentation Milvus pour obtenir des instructions détaillées.

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. Insérer des données dans Milvus

Dans cette étape, nous utilisons un encodeur d'images prédéfini pour générer des embeddings pour toutes les images JPEG dans le répertoire de données de l'exemple. Ces embeddings sont ensuite insérés dans la collection Milvus, avec les chemins d'accès aux fichiers correspondants. Chaque entrée de la collection se compose des éléments suivants

• Vecteur d'intégration: La représentation numérique de l'image. Stockée dans le champ vector.
• Chemin d'accès au fichier: L'emplacement du fichier d'image pour référence. Stocké dans le champ filepath en tant que champ dynamique.

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.

Effectuer une recherche

Lançons maintenant une recherche à l'aide d'un exemple de requête textuelle. Cette recherche permettra d'extraire les images les plus pertinentes en fonction de leur similarité sémantique avec la description textuelle donnée.

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
)

Visualiser les résultats :

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