Búsqueda texto-imagen con Milvus

La búsqueda de texto a imagen es una tecnología avanzada que permite a los usuarios buscar imágenes utilizando descripciones de texto en lenguaje natural. Aprovecha un modelo multimodal previamente entrenado para convertir tanto el texto como las imágenes en incrustaciones en un espacio semántico compartido, lo que permite realizar comparaciones basadas en similitudes.

En este tutorial, exploraremos cómo implementar la recuperación de imágenes basada en texto utilizando el modelo CLIP (Contrastive Language-Image Pretraining) de OpenAI y Milvus. Generaremos incrustaciones de imágenes con CLIP, las almacenaremos en Milvus y realizaremos búsquedas de similitud eficientes.

Requisitos previos

Antes de empezar, asegúrese de tener listos todos los paquetes necesarios y los datos de ejemplo.

Instale las dependencias

• pymilvus>=2.4.2 para interactuar con la base de datos Milvus
• clip para trabajar con el modelo CLIP
• pillow para el procesamiento y visualización de imágenes

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

Descargar datos de ejemplo

Utilizaremos un subconjunto del conjunto de datos ImageNet (100 clases, 10 imágenes para cada clase) como imágenes de ejemplo. El siguiente comando descargará los datos de ejemplo y los extraerá a la carpeta 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

Configurar Milvus

Antes de continuar, configure su servidor Milvus y conéctese utilizando su URI (y opcionalmente, un token):

• Milvus Lite (Recomendado por conveniencia): Establezca el URI en un archivo local, como ./milvus.db. Esto aprovecha automáticamente Milvus Lite para almacenar todos los datos en un único archivo.

• Docker o Kubernetes (para datos a gran escala): Para manejar conjuntos de datos más grandes, despliegue un servidor Milvus de mayor rendimiento utilizando Docker o Kubernetes. En este caso, utilice el URI del servidor, como http://localhost:19530, para conectarse.

• Zilliz Cloud (servicio gestionado): Si está utilizando Zilliz Cloud, el servicio en la nube totalmente gestionado de Milvus, establezca el Public Endpoint como URI y API Key como token.

from pymilvus import MilvusClient

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

Primeros pasos

Ahora que tiene las dependencias y los datos necesarios, es hora de configurar los extractores de características y empezar a trabajar con Milvus. Esta sección le guiará a través de los pasos clave para construir un sistema de búsqueda de texto a imagen. Finalmente, demostraremos cómo recuperar y visualizar imágenes basadas en consultas de texto.

Definir extractores de características

Utilizaremos un modelo CLIP preentrenado para generar incrustaciones de imagen y texto. En esta sección, cargaremos la variante ViT-B/32 preentrenada de CLIP y definiremos funciones de ayuda para la codificación de imágenes y texto:

• encode_image(image_path): Procesa y codifica imágenes en vectores de características.
• encode_text(text): Codifica consultas de texto en vectores de características

Ambas funciones normalizan las características de salida para garantizar comparaciones coherentes convirtiendo los vectores a longitud unitaria, lo que resulta esencial para calcular con precisión la similitud 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()

Ingestión de datos

Para permitir la búsqueda semántica de imágenes, primero necesitamos generar incrustaciones para todas las imágenes y almacenarlas en una base de datos vectorial para una indexación y recuperación eficientes. Esta sección proporciona una guía paso a paso para introducir datos de imágenes en Milvus.

1. Crear una colección Milvus

Antes de almacenar las incrustaciones de imágenes, debe crear una colección Milvus. El siguiente código muestra cómo crear una colección en modo de configuración rápida con el tipo de métrica COSINE por defecto. La colección incluye los siguientes campos:

• id: Un campo primario con ID automático activado.

• vector: Un campo para almacenar incrustaciones de vectores de punto flotante.

Si necesita un esquema personalizado, consulte la documentación de Milvus para obtener instrucciones detalladas.

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. Insertar datos en Milvus

En este paso, utilizamos un codificador de imágenes predefinido para generar incrustaciones para todas las imágenes JPEG del directorio de datos de ejemplo. Estas incrustaciones se insertan en la colección Milvus, junto con sus correspondientes rutas de archivo. Cada entrada de la colección consta de

• Vector de incrustación: La representación numérica de la imagen. Se almacena en el campo vector.
• Ruta del archivo: La ubicación del archivo de imagen como referencia. Se almacena en el campo filepath como campo dinámico.

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.

Realizar una búsqueda

Ahora vamos a realizar una búsqueda utilizando una consulta de texto de ejemplo. Esto recuperará las imágenes más relevantes basándose en su similitud semántica con la descripción de texto dada.

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
)

Visualice los resultados:

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