Pesquisa de texto para imagem com Milvus

A pesquisa de texto para imagem é uma tecnologia avançada que permite aos utilizadores pesquisar imagens utilizando descrições de texto em linguagem natural. Esta tecnologia utiliza um modelo multimodal pré-treinado para converter texto e imagens em embeddings num espaço semântico partilhado, permitindo comparações baseadas em semelhanças.

Neste tutorial, vamos explorar como implementar a recuperação de imagens baseadas em texto usando o modelo CLIP (Contrastive Language-Image Pretraining) da OpenAI e o Milvus. Vamos gerar embeddings de imagens com o CLIP, armazená-los no Milvus e realizar pesquisas de similaridade eficientes.

Pré-requisitos

Antes de começar, certifique-se de que tem todos os pacotes necessários e dados de exemplo prontos.

Instale as dependências

• pymilvus>=2.4.2 para interagir com a base de dados Milvus
• clip para trabalhar com o modelo CLIP
• pillow para processamento e visualização de imagens

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

Descarregar dados de exemplo

Iremos utilizar um subconjunto do conjunto de dados ImageNet (100 classes, 10 imagens para cada classe) como imagens de exemplo. O comando seguinte descarrega os dados de exemplo e extrai-os para a pasta 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 o Milvus

Antes de continuar, configure o seu servidor Milvus e ligue-se utilizando o seu URI (e, opcionalmente, um token):

• Milvus Lite (recomendado por conveniência): Defina o URI para um ficheiro local, como ./milvus.db. Isso aproveita automaticamente o Milvus Lite para armazenar todos os dados em um único arquivo.

• Docker ou Kubernetes (para dados em grande escala): Para lidar com conjuntos de dados maiores, implante um servidor Milvus de melhor desempenho usando Docker ou Kubernetes. Neste caso, utilize o URI do servidor, como http://localhost:19530, para se ligar.

• Zilliz Cloud (Serviço Gerido): Se estiver a utilizar o Zilliz Cloud, o serviço de nuvem totalmente gerido do Milvus, defina o Public Endpoint como URI e a API Key como token.

from pymilvus import MilvusClient

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

Começar a utilizar

Agora que tem as dependências e os dados necessários, está na altura de configurar os extractores de caraterísticas e começar a trabalhar com o Milvus. Esta secção irá guiá-lo através dos principais passos da construção de um sistema de pesquisa de texto para imagem. Por fim, demonstraremos como recuperar e visualizar imagens com base em consultas de texto.

Definir extractores de caraterísticas

Utilizaremos um modelo CLIP pré-treinado para gerar embeddings de imagem e texto. Nesta secção, carregamos a variante ViT-B/32 pré-treinada do CLIP e definimos funções auxiliares para codificar imagem e texto:

• encode_image(image_path): Processa e codifica imagens em vectores de caraterísticas
• encode_text(text): Codifica consultas de texto em vectores de caraterísticas

Ambas as funções normalizam as caraterísticas de saída para garantir comparações consistentes, convertendo os vectores para comprimento unitário, o que é essencial para cálculos precisos de semelhança de cosseno.

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

Ingestão de dados

Para permitir a pesquisa semântica de imagens, precisamos primeiro de gerar embeddings para todas as imagens e armazená-las numa base de dados de vectores para uma indexação e recuperação eficientes. Esta secção fornece um guia passo-a-passo para ingerir dados de imagens no Milvus.

1. Criar uma coleção Milvus

Antes de armazenar os embeddings de imagens, é necessário criar uma coleção Milvus. O código seguinte demonstra como criar uma coleção num modo de configuração rápida com o tipo de métrica COSINE predefinido. A coleção inclui os seguintes campos:

• id: Um campo primário com a ID automática activada.

• vector: Um campo para armazenar as incorporações de vectores de vírgula flutuante.

Se precisar de um esquema personalizado, consulte a documentação do Milvus para obter instruções detalhadas.

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. Inserir dados no Milvus

Neste passo, usamos um codificador de imagem predefinido para gerar embeddings para todas as imagens JPEG no diretório de dados de exemplo. Estes embeddings são depois inseridos na coleção do Milvus, juntamente com os caminhos dos ficheiros correspondentes. Cada entrada na coleção é composta por:

• Vetor de incorporação: A representação numérica da imagem. Armazenado no campo vector.
• Caminho do ficheiro: A localização do ficheiro de imagem para referência. Armazenado no campo filepath como um 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.

Efetuar uma pesquisa

Agora, vamos executar uma pesquisa usando uma consulta de texto de exemplo. Isto irá recuperar as imagens mais relevantes com base na sua semelhança semântica com a descrição de texto fornecida.

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
)

Visualize os 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